Skip to content

fix: emit [0] offsets for empty structural list streams (alternative to #10533) #10534

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danking wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

fix: emit [0] offsets for empty structural list streams (alternative to #10533) #10534

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 101 additions & 0 deletions vortex-file/src/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ use vortex_layout::LayoutStrategy;
use vortex_layout::layouts::buffered::BufferedStrategy;
use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy;
use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
use vortex_layout::layouts::list::ListLayoutEncoding;
use vortex_layout::layouts::struct_::StructStrategy;
use vortex_layout::layouts::table::TableStrategy;
use vortex_layout::layouts::zoned::LegacyStats;
Expand Down Expand Up @@ -2069,6 +2070,106 @@ async fn nested_list_of_list_roundtrip() -> VortexResult<()> {
Ok(())
}

#[rstest]
#[case::empty_lists(Validity::NonNullable)]
#[case::nullable_empty_lists(Validity::AllValid)]
#[case::null_lists(Validity::AllInvalid)]
#[case::mixed_empty_and_null(
Validity::Array(BoolArray::from_iter([true, false]).into_array())
)]
#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn nested_empty_list_roundtrip(#[case] validity: Validity) -> VortexResult<()> {
let inner = ListArray::try_new(
PrimitiveArray::new::<i32>(buffer![], Validity::NonNullable).into_array(),
buffer![0u32].into_array(),
Validity::NonNullable,
)?
.into_array();
let outer = ListArray::try_new(inner, buffer![0u32, 0, 0].into_array(), validity)?.into_array();
let array = StructArray::from_fields(&[("nested", outer)])?.into_array();

let result = write_read_roundtrip(array.clone()).await?;
assert_arrays_eq!(result, array, &mut SESSION.create_execution_ctx());
Ok(())
}

#[rstest]
#[case::one_level(1)]
#[case::two_levels(2)]
#[case::five_levels(5)]
#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn empty_nested_list_layouts_have_zero_offsets(
#[case] depth: usize,
#[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability,
) -> VortexResult<()> {
let mut elements =
PrimitiveArray::new::<i32>(buffer![], Validity::from(nullability)).into_array();
for _ in 1..depth {
elements = ListArray::try_new(
elements,
buffer![0u32].into_array(),
Validity::from(nullability),
)?
.into_array();
}
let lists = ListArray::try_new(
elements,
buffer![0u32, 0, 0].into_array(),
Validity::from(nullability),
)?
.into_array();
let array = StructArray::from_fields(&[("nested", lists)])?.into_array();
let strategy = crate::strategy::WriteStrategyBuilder::from_session(&SESSION)
.with_list_layout()
.build();
let mut buf = ByteBufferMut::empty();
SESSION
.write_options()
.with_strategy(strategy)
.write(&mut buf, array.to_array_stream())
.await?;
let file = SESSION.open_options().open_buffer(buf)?;
let result = file.scan()?.into_array_stream()?.read_all().await?;
assert_arrays_eq!(result, array, &mut SESSION.create_execution_ctx());

let mut list_levels = 0;
let mut empty_levels = 0;
for layout in file.footer().layout().depth_first_traversal() {
let layout = layout?;
let Some(list) = layout.as_opt::<ListLayoutEncoding>() else {
continue;
};
list_levels += 1;
if list.row_count() == 0 {
empty_levels += 1;
assert_eq!(list.elements()?.row_count(), 0);
if let Some(validity) = list.validity()? {
assert_eq!(validity.row_count(), 0);
}
let reader = list.offsets()?.new_reader(
"offsets".into(),
file.segment_source(),
&SESSION,
&Default::default(),
)?;
let offsets = ScanBuilder::new(SESSION.clone(), reader)
.into_array_stream()?
.read_all()
.await?;
assert_arrays_eq!(
offsets,
buffer![0u64].into_array(),
&mut SESSION.create_execution_ctx()
);
}
}
assert_eq!(list_levels, depth);
assert_eq!(empty_levels, depth - 1);
Ok(())
}

type MapEntryFixture<'a> = (i32, Option<&'a str>);
type MapRowFixture<'a> = Option<Vec<MapEntryFixture<'a>>>;

Expand Down
35 changes: 29 additions & 6 deletions vortex-layout/src/layouts/list/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use vortex_array::dtype::Nullability;
use vortex_array::dtype::PType;
use vortex_array::matcher::Matcher;
use vortex_array::scalar_fn::fns::operators::Operator;
use vortex_buffer::buffer;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
Expand Down Expand Up @@ -149,11 +150,16 @@ impl LayoutStrategy for ListLayoutStrategy {
(None, None)
};

// Reserve the sentinel's sequence before the child EOFs so the offsets
// writer can publish it even when the input contains no chunks.
let empty_offsets_sequence = eof.split_off().downgrade();

// Transpose the list column into its child sub-streams and rebase offsets to global
// positions. Kept joined with the child writers below so producer errors surface rather
// than being hidden as an early channel close.
let fanout_fut = transpose_list_column(
stream,
empty_offsets_sequence,
session.clone(),
elements_tx,
offsets_tx,
Expand Down Expand Up @@ -214,6 +220,7 @@ impl LayoutStrategy for ListLayoutStrategy {
/// joins this against the child writers, rather than being hidden as an early channel close.
async fn transpose_list_column(
mut stream: SendableSequentialStream,
empty_offsets_sequence: SequenceId,
session: VortexSession,
elements_tx: kanal::AsyncSender<ChildChunk>,
offsets_tx: kanal::AsyncSender<ChildChunk>,
Expand Down Expand Up @@ -261,7 +268,15 @@ async fn transpose_list_column(
}
}
if !saw_chunk {
vortex_bail!("ListLayoutStrategy needs at least one chunk");
// Zero lists still need one boundary offset. Closing the other channels
// leaves typed empty elements and validity streams for their writers.
if offsets_tx
.send(Ok((empty_offsets_sequence, buffer![0u64].into_array())))
.await
.is_err()
{
vortex_bail!("list offset writer finished before the empty offset was sent");
}
}
Ok(())
}
Expand Down Expand Up @@ -322,6 +337,7 @@ impl Matcher for AnyList {
#[cfg(test)]
mod tests {
use futures::stream;
use rstest::rstest;
use vortex_array::ArrayContext;
use vortex_array::arrays::BoolArray;
use vortex_array::arrays::ChunkedArray;
Expand Down Expand Up @@ -431,24 +447,31 @@ mod tests {
Ok(())
}

#[rstest]
#[case::non_nullable(false)]
#[case::nullable(true)]
#[tokio::test]
async fn empty_stream_errors() {
async fn empty_stream_preserves_dtype_and_row_count(
#[case] nullable: bool,
) -> VortexResult<()> {
let segments = Arc::new(TestSegments::default());
let (_, eof) = SequenceId::root().split();
let empty = stream::empty::<VortexResult<(SequenceId, ArrayRef)>>().boxed();
let stream = SequentialStreamAdapter::new(i32_list_dtype(false), empty).sendable();
let stream = SequentialStreamAdapter::new(i32_list_dtype(nullable), empty).sendable();
let session = layout_test_session();

let res = flat_list_strategy()
let layout = flat_list_strategy()
.write_stream(
ArrayContext::empty().into(),
segments,
stream,
eof,
&session,
)
.await;
assert!(res.is_err())
.await?;
assert_eq!(layout.dtype(), &i32_list_dtype(nullable));
assert_eq!(layout.row_count(), 0);
Ok(())
}

#[tokio::test]
Expand Down
Loading