diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 0098fbeae52..7ad0dc33a9b 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2069,6 +2069,30 @@ 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::(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(()) +} + type MapEntryFixture<'a> = (i32, Option<&'a str>); type MapRowFixture<'a> = Option>>; diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 4f41dc3b20d..e23af8af3c5 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -33,6 +33,8 @@ use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; use crate::LayoutWriterContext; +use crate::children::OwnedLayoutChildren; +use crate::layouts::chunked::ChunkedLayout; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::ListLayout; use crate::segments::SegmentSinkRef; @@ -117,7 +119,7 @@ impl LayoutStrategy for ListLayoutStrategy { &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, + mut stream: SendableSequentialStream, mut eof: SequencePointer, session: &VortexSession, ) -> VortexResult { @@ -129,6 +131,20 @@ impl LayoutStrategy for ListLayoutStrategy { .await; } + // Empty nested element streams are valid below empty or all-null lists. + // As in the flat writer, no input chunks require no child segments. + let first_chunk = match stream.next().await { + Some(chunk) => chunk?, + None => { + return Ok(ChunkedLayout::new( + 0, + dtype, + OwnedLayoutChildren::layout_children(vec![]), + ) + .into_layout()); + } + }; + let is_nullable = dtype.is_nullable(); let element_dtype = dtype .as_list_element_opt() @@ -153,6 +169,7 @@ impl LayoutStrategy for ListLayoutStrategy { // 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( + first_chunk, stream, session.clone(), elements_tx, @@ -210,10 +227,12 @@ impl LayoutStrategy for ListLayoutStrategy { /// sub-streams, rebasing each chunk's local `offsets` to global `u64` positions so the single /// `offsets` child indexes into the concatenated `elements` child. /// +/// `first_chunk` establishes nonempty input; `stream` contains the remaining chunks. /// `validity_tx` is `Some` exactly when the list is nullable. Errors surface to the caller, which /// joins this against the child writers, rather than being hidden as an early channel close. async fn transpose_list_column( - mut stream: SendableSequentialStream, + first_chunk: (SequenceId, ArrayRef), + stream: SendableSequentialStream, session: VortexSession, elements_tx: kanal::AsyncSender, offsets_tx: kanal::AsyncSender, @@ -222,10 +241,9 @@ async fn transpose_list_column( let mut exec_ctx = session.create_execution_ctx(); let mut element_base: u64 = 0; let mut first = true; - let mut saw_chunk = false; - while let Some(chunk) = stream.next().await { + let mut chunks = futures::stream::iter([Ok(first_chunk)]).chain(stream); + while let Some(chunk) = chunks.next().await { let (sequence_id, array) = chunk?; - saw_chunk = true; let mut sp = sequence_id.descend(); let ListDataParts { elements, @@ -260,9 +278,6 @@ async fn transpose_list_column( } } } - if !saw_chunk { - vortex_bail!("ListLayoutStrategy needs at least one chunk"); - } Ok(()) } @@ -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; @@ -431,15 +447,20 @@ 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::>().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, @@ -447,8 +468,10 @@ mod tests { 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]