From b4028cddadc4c36ef485a48b2f6336ccbc263bd4 Mon Sep 17 00:00:00 2001 From: Matthew Katz Date: Fri, 22 May 2026 14:58:03 -0400 Subject: [PATCH 01/37] first pass Signed-off-by: Matthew Katz Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 308 ++++++++++++++++ vortex-layout/src/layouts/list/reader.rs | 190 ++++++++++ vortex-layout/src/layouts/list/writer.rs | 440 +++++++++++++++++++++++ vortex-layout/src/layouts/mod.rs | 1 + vortex-layout/src/session.rs | 2 + 5 files changed, 941 insertions(+) create mode 100644 vortex-layout/src/layouts/list/mod.rs create mode 100644 vortex-layout/src/layouts/list/reader.rs create mode 100644 vortex-layout/src/layouts/list/writer.rs diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs new file mode 100644 index 00000000000..645c25ffdff --- /dev/null +++ b/vortex-layout/src/layouts/list/mod.rs @@ -0,0 +1,308 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod reader; +pub mod writer; + +use std::sync::Arc; + +use reader::ListReader; +use vortex_array::DeserializeMetadata; +use vortex_array::ProstMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::LayoutChildType; +use crate::LayoutEncodingRef; +use crate::LayoutId; +use crate::LayoutReaderRef; +use crate::LayoutRef; +use crate::VTable; +use crate::children::LayoutChildren; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; +use crate::vtable; + +/// Child index of the `elements` layout. +pub const ELEMENTS_CHILD_INDEX: usize = 0; +/// Child index of the `offsets` layout. +pub const OFFSETS_CHILD_INDEX: usize = 1; +/// Child index of the `validity` layout (only present when the list dtype is nullable). +pub const VALIDITY_CHILD_INDEX: usize = 2; + +/// Number of children when the list dtype is non-nullable. +pub const NUM_CHILDREN_NON_NULLABLE: usize = 2; +/// Number of children when the list dtype is nullable. +pub const NUM_CHILDREN_NULLABLE: usize = 3; + +vtable!(List); + +impl VTable for List { + type Layout = ListLayout; + type Encoding = ListLayoutEncoding; + type Metadata = ProstMetadata; + + fn id(_encoding: &Self::Encoding) -> LayoutId { + LayoutId::new("vortex.list") + } + + fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { + LayoutEncodingRef::new_ref(ListLayoutEncoding.as_ref()) + } + + fn row_count(layout: &Self::Layout) -> u64 { + layout.row_count() + } + + fn dtype(layout: &Self::Layout) -> &DType { + &layout.dtype + } + + fn metadata(layout: &Self::Layout) -> Self::Metadata { + ProstMetadata(ListLayoutMetadata::new(layout.offsets_ptype)) + } + + fn segment_ids(_layout: &Self::Layout) -> Vec { + vec![] + } + + fn nchildren(layout: &Self::Layout) -> usize { + if layout.dtype.is_nullable() { + NUM_CHILDREN_NULLABLE + } else { + NUM_CHILDREN_NON_NULLABLE + } + } + + fn child(layout: &Self::Layout, idx: usize) -> VortexResult { + match (idx, layout.validity.as_ref()) { + (ELEMENTS_CHILD_INDEX, _) => Ok(Arc::clone(&layout.elements)), + (OFFSETS_CHILD_INDEX, _) => Ok(Arc::clone(&layout.offsets)), + (VALIDITY_CHILD_INDEX, Some(validity)) => Ok(Arc::clone(validity)), + _ => vortex_bail!("Invalid child index {idx} for ListLayout"), + } + } + + fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType { + match (idx, layout.validity.is_some()) { + (ELEMENTS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("elements".into()), + (OFFSETS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("offsets".into()), + (VALIDITY_CHILD_INDEX, true) => LayoutChildType::Transparent("validity".into()), + _ => vortex_panic!("Invalid child index {idx} for ListLayout"), + } + } + + fn new_reader( + layout: &Self::Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ) -> VortexResult { + Ok(Arc::new(ListReader::try_new( + layout.clone(), + name, + segment_source, + session.clone(), + )?)) + } + + fn build( + _encoding: &Self::Encoding, + dtype: &DType, + _row_count: u64, + metadata: &::Output, + _segment_ids: Vec, + children: &dyn LayoutChildren, + _ctx: &ReadContext, + ) -> VortexResult { + let elements_dtype = dtype + .as_list_element_opt() + .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {dtype}"))?; + + let expected_children = if dtype.is_nullable() { + NUM_CHILDREN_NULLABLE + } else { + NUM_CHILDREN_NON_NULLABLE + }; + vortex_ensure!( + children.nchildren() == expected_children, + "ListLayout expects {expected_children} children, got {}", + children.nchildren(), + ); + + let elements = children.child(ELEMENTS_CHILD_INDEX, elements_dtype.as_ref())?; + let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); + let offsets = children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; + let validity = if dtype.is_nullable() { + Some(children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))?) + } else { + None + }; + + Ok(ListLayout { + dtype: dtype.clone(), + elements, + offsets, + validity, + offsets_ptype: metadata.offsets_ptype(), + }) + } + + fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { + let expected = if layout.dtype.is_nullable() { + NUM_CHILDREN_NULLABLE + } else { + NUM_CHILDREN_NON_NULLABLE + }; + vortex_ensure!( + children.len() == expected, + "ListLayout expects {expected} children, got {}", + children.len() + ); + let mut iter = children.into_iter(); + layout.elements = iter + .next() + .ok_or_else(|| vortex_err!("missing elements child"))?; + layout.offsets = iter + .next() + .ok_or_else(|| vortex_err!("missing offsets child"))?; + layout.validity = if layout.dtype.is_nullable() { + Some( + iter.next() + .ok_or_else(|| vortex_err!("missing validity child"))?, + ) + } else { + None + }; + Ok(()) + } +} + +#[derive(Debug)] +pub struct ListLayoutEncoding; + +/// Stores a list-typed column in Apache Arrow-compatible form: a flattened `elements` buffer +/// plus an `offsets` buffer of length `n+1` (and optionally a `validity` bitmap of length `n`). +/// +/// Mirrors the in-memory [`ListArray`](vortex_array::arrays::ListArray): offsets are monotonic +/// and `list[i] = elements[offsets[i]..offsets[i+1]]`. +/// +/// To write a [`ListViewArray`](vortex_array::arrays::ListViewArray) (the canonical encoding for +/// [`DType::List`]) into a `ListLayout`, the writer rebuilds it into zero-copy-to-list form via +/// [`list_from_list_view`](vortex_array::arrays::listview::list_from_list_view). +#[derive(Clone, Debug)] +pub struct ListLayout { + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, + offsets_ptype: PType, +} + +impl ListLayout { + /// Construct a new `ListLayout` from its components. + /// + /// `dtype` must be a [`DType::List`]. The `validity` child must be supplied iff the dtype is + /// nullable. The list's row count is derived from the `offsets` child as + /// `offsets.row_count() - 1`. + pub fn try_new( + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, + ) -> VortexResult { + if !dtype.is_list() { + vortex_bail!("ListLayout requires a List dtype, got {dtype}"); + } + vortex_ensure!( + dtype.is_nullable() == validity.is_some(), + "validity must be supplied iff dtype is nullable (dtype: {dtype})", + ); + if !offsets.dtype().is_int() || offsets.dtype().is_nullable() { + vortex_bail!( + "offsets must be a non-nullable integer layout, got {}", + offsets.dtype() + ); + } + vortex_ensure!( + offsets.row_count() >= 1, + "ListLayout offsets must have at least 1 row (n+1 entries for n lists), got 0", + ); + let row_count = offsets.row_count() - 1; + if let Some(validity) = validity.as_ref() { + vortex_ensure!( + validity.row_count() == row_count, + "validity row count ({}) must match list row count ({row_count})", + validity.row_count(), + ); + } + + let offsets_ptype = offsets.dtype().as_ptype(); + + Ok(Self { + dtype, + elements, + offsets, + validity, + offsets_ptype, + }) + } + + /// Number of lists in this layout. + /// + /// Equal to `offsets.row_count() - 1` by construction. + #[inline] + pub fn row_count(&self) -> u64 { + self.offsets.row_count().saturating_sub(1) + } + + #[inline] + pub fn elements(&self) -> &LayoutRef { + &self.elements + } + + #[inline] + pub fn offsets(&self) -> &LayoutRef { + &self.offsets + } + + #[inline] + pub fn validity(&self) -> Option<&LayoutRef> { + self.validity.as_ref() + } + + #[inline] + pub fn offsets_ptype(&self) -> PType { + self.offsets_ptype + } + + /// The dtype of the inner elements column. + pub fn elements_dtype(&self) -> &DType { + self.dtype + .as_list_element_opt() + .vortex_expect("ListLayout dtype must be a List") + } +} + +#[derive(prost::Message)] +pub struct ListLayoutMetadata { + #[prost(enumeration = "PType", tag = "1")] + offsets_ptype: i32, +} + +impl ListLayoutMetadata { + pub fn new(offsets_ptype: PType) -> Self { + let mut metadata = Self::default(); + metadata.set_offsets_ptype(offsets_ptype); + metadata + } +} diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs new file mode 100644 index 00000000000..72164eeb155 --- /dev/null +++ b/vortex-layout/src/layouts/list/reader.rs @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::BTreeSet; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::try_join; +use vortex_array::Array; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::arrays::List; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::dtype::Nullability; +use vortex_array::expr::Expression; +use vortex_array::expr::root; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::ArrayFuture; +use crate::LayoutReader; +use crate::LayoutReaderRef; +use crate::SplitRange; +use crate::layouts::list::ListLayout; +use crate::segments::SegmentSource; + +/// Reader for [`ListLayout`]. +/// +/// Reads the underlying `elements`, `offsets`, and (optional) `validity` children in parallel +/// and reassembles them into a [`ListArray`](vortex_array::arrays::ListArray) before applying +/// the requested projection. +pub struct ListReader { + layout: ListLayout, + name: Arc, + _session: VortexSession, + elements: LayoutReaderRef, + offsets: LayoutReaderRef, + validity: Option, +} + +impl ListReader { + pub(super) fn try_new( + layout: ListLayout, + name: Arc, + segment_source: Arc, + session: VortexSession, + ) -> VortexResult { + let elements = layout.elements().new_reader( + format!("{name}.elements").into(), + Arc::clone(&segment_source), + &session, + )?; + let offsets = layout.offsets().new_reader( + format!("{name}.offsets").into(), + Arc::clone(&segment_source), + &session, + )?; + let validity = layout + .validity() + .map(|v| { + v.new_reader( + format!("{name}.validity").into(), + Arc::clone(&segment_source), + &session, + ) + }) + .transpose()?; + + Ok(Self { + layout, + name, + _session: session, + elements, + offsets, + validity, + }) + } +} + +impl LayoutReader for ListReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn dtype(&self) -> &DType { + self.layout.dtype() + } + + fn row_count(&self) -> u64 { + self.layout.row_count() + } + + fn register_splits( + &self, + field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut BTreeSet, + ) -> VortexResult<()> { + // Offsets has one more row than the list itself but shares the list's chunking + // structure, so it's the appropriate child to drive scan splits. + self.offsets.register_splits(field_mask, split_range, splits) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let elements_len = self.layout.elements().row_count(); + let elements_range = 0..elements_len; + let elements_mask = MaskFuture::new_true(usize::try_from(elements_len)?); + + let row_count = usize::try_from(row_range.end - row_range.start)?; + // The offsets child has n+1 entries; reading row_range maps to offsets in + // [row_range.start..row_range.end + 1). + let offsets_range = row_range.start..row_range.end + 1; + let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; + + let elements_fut = self + .elements + .projection_evaluation(&elements_range, &root(), elements_mask)?; + let offsets_fut = self.offsets.projection_evaluation( + &offsets_range, + &root(), + MaskFuture::new_true(offsets_count), + )?; + let validity_fut = self + .validity + .as_ref() + .map(|v| v.projection_evaluation(row_range, &root(), MaskFuture::new_true(row_count))) + .transpose()?; + + let nullability = self.layout.dtype().nullability(); + let expr = expr.clone(); + + Ok(async move { + let (elements, offsets) = try_join!(elements_fut, offsets_fut)?; + let validity = match validity_fut { + Some(fut) => Validity::Array(fut.await?), + None => match nullability { + Nullability::Nullable => Validity::AllValid, + Nullability::NonNullable => Validity::NonNullable, + }, + }; + + // SAFETY: the layout was validated at write time, so offsets are monotonic, + // non-nullable, integer, of length n+1. + let array: ArrayRef = + unsafe { Array::::new_unchecked(elements, offsets, validity) }.into_array(); + + let mask = mask.await?; + let array = if !mask.all_true() { + array.filter(mask)? + } else { + array + }; + + array.apply(&expr) + } + .boxed()) + } +} diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs new file mode 100644 index 00000000000..d6bb727502b --- /dev/null +++ b/vortex-layout/src/layouts/list/writer.rs @@ -0,0 +1,440 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use async_trait::async_trait; +use futures::StreamExt; +use futures::future::try_join_all; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::session::RuntimeSessionExt; +use vortex_session::VortexSession; + +use crate::IntoLayout; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::children::OwnedLayoutChildren; +use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::list::ListLayout; +use crate::segments::SegmentSinkRef; +use crate::sequence::SendableSequentialStream; +use crate::sequence::SequenceId; +use crate::sequence::SequencePointer; +use crate::sequence::SequentialStream; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; + +/// Strategy for writing list-typed arrays as an Apache Arrow-style [`ListLayout`]. +/// +/// For each input chunk, the strategy: +/// 1. Canonicalizes the chunk into a [`ListViewArray`]. +/// 2. Calls [`list_from_list_view`] to rebuild it into zero-copy-to-list form +/// (sorted, gapless, non-overlapping offsets) and produce a [`ListArray`]. +/// 3. Writes the `elements`, `n+1` `offsets`, and (when nullable) `validity` columns into +/// separately configurable downstream strategies, producing a single [`ListLayout`] for +/// that chunk. +/// +/// When the input stream contains multiple chunks, each chunk becomes one `ListLayout` and the +/// strategy wraps them in a [`ChunkedLayout`]. +/// +/// [`ListArray`]: vortex_array::arrays::ListArray +pub struct ListLayoutStrategy { + elements: Arc, + offsets: Arc, + validity: Arc, +} + +impl ListLayoutStrategy { + pub fn new( + elements: Arc, + offsets: Arc, + validity: Arc, + ) -> Self { + Self { + elements, + offsets, + validity, + } + } +} + +#[async_trait] +impl LayoutStrategy for ListLayoutStrategy { + async fn write_stream( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + mut eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + if !dtype.is_list() { + return Err(vortex_err!( + "ListLayoutStrategy requires a List dtype, got {dtype}" + )); + } + let is_nullable = dtype.is_nullable(); + + let mut stream = stream; + let mut chunk_layouts: Vec = Vec::new(); + + while let Some(item) = stream.next().await { + let (sequence_id, array) = item?; + let chunk_eof = eof.split_off(); + + let chunk_layout = self + .write_chunk( + ctx.clone(), + Arc::clone(&segment_sink), + sequence_id, + chunk_eof, + session, + &dtype, + is_nullable, + array, + ) + .await?; + chunk_layouts.push(chunk_layout); + } + + if chunk_layouts.is_empty() { + // Empty stream: emit an empty ListLayout. We write empty children of the same dtype + // (with u32 offsets) so the children's encoding ID is still derivable. + let chunk_layout = self + .write_empty_chunk(ctx, segment_sink, eof, session, dtype.clone(), is_nullable) + .await?; + return Ok(chunk_layout); + } + + if chunk_layouts.len() == 1 { + return Ok(chunk_layouts.pop().expect("len == 1")); + } + + let row_count = chunk_layouts.iter().map(|l| l.row_count()).sum(); + Ok(ChunkedLayout::new( + row_count, + dtype, + OwnedLayoutChildren::layout_children(chunk_layouts), + ) + .into_layout()) + } + + fn buffered_bytes(&self) -> u64 { + self.elements.buffered_bytes() + + self.offsets.buffered_bytes() + + self.validity.buffered_bytes() + } +} + +impl ListLayoutStrategy { + #[allow(clippy::too_many_arguments)] + async fn write_chunk( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + sequence_id: SequenceId, + mut chunk_eof: SequencePointer, + session: &VortexSession, + dtype: &DType, + is_nullable: bool, + array: ArrayRef, + ) -> VortexResult { + let mut exec_ctx = session.create_execution_ctx(); + + // Canonicalize then rebuild into ZCTL form. + let listview = array.execute::(&mut exec_ctx)?; + let list = list_from_list_view(listview)?; + let parts = list.into_data_parts(); + let row_count = parts.offsets.len().saturating_sub(1); + + let elements_dtype = parts.elements.dtype().clone(); + let offsets_dtype = parts.offsets.dtype().clone(); + + // Build single-chunk streams for each child. Each child sees one array and EOF. + let mut sequence_pointer = sequence_id.descend(); + let elements_seq = sequence_pointer.advance(); + let offsets_seq = sequence_pointer.advance(); + let validity_seq = if is_nullable { + Some(sequence_pointer.advance()) + } else { + None + }; + drop(sequence_pointer); + + let elements_eof = chunk_eof.split_off(); + let offsets_eof = chunk_eof.split_off(); + let validity_eof = if is_nullable { + Some(chunk_eof.split_off()) + } else { + None + }; + drop(chunk_eof); + + let elements_stream = SequentialStreamAdapter::new( + elements_dtype, + futures::stream::once(async move { Ok((elements_seq, parts.elements)) }).boxed(), + ) + .sendable(); + let offsets_stream = SequentialStreamAdapter::new( + offsets_dtype, + futures::stream::once(async move { Ok((offsets_seq, parts.offsets)) }).boxed(), + ) + .sendable(); + + let validity_array = if is_nullable { + Some( + parts + .validity + .execute_mask(row_count, &mut exec_ctx)? + .into_array(), + ) + } else { + None + }; + let validity_stream = validity_array.map(|arr| { + let seq = validity_seq.expect("validity sequence id"); + SequentialStreamAdapter::new( + DType::Bool(Nullability::NonNullable), + futures::stream::once(async move { Ok((seq, arr)) }).boxed(), + ) + .sendable() + }); + + let handle = session.handle(); + + let elements_strategy = Arc::clone(&self.elements); + let elements_ctx = ctx.clone(); + let elements_sink = Arc::clone(&segment_sink); + let elements_session = session.clone(); + let elements_task = handle.spawn_nested(move |h| async move { + let session = elements_session.with_handle(h); + elements_strategy + .write_stream( + elements_ctx, + elements_sink, + elements_stream, + elements_eof, + &session, + ) + .await + }); + + let offsets_strategy = Arc::clone(&self.offsets); + let offsets_ctx = ctx.clone(); + let offsets_sink = Arc::clone(&segment_sink); + let offsets_session = session.clone(); + let offsets_task = handle.spawn_nested(move |h| async move { + let session = offsets_session.with_handle(h); + offsets_strategy + .write_stream( + offsets_ctx, + offsets_sink, + offsets_stream, + offsets_eof, + &session, + ) + .await + }); + + let mut tasks = vec![elements_task, offsets_task]; + if let (Some(validity_stream), Some(validity_eof)) = (validity_stream, validity_eof) { + let validity_strategy = Arc::clone(&self.validity); + let validity_ctx = ctx; + let validity_sink = segment_sink; + let validity_session = session.clone(); + tasks.push(handle.spawn_nested(move |h| async move { + let session = validity_session.with_handle(h); + validity_strategy + .write_stream( + validity_ctx, + validity_sink, + validity_stream, + validity_eof, + &session, + ) + .await + })); + } + + let mut child_layouts = try_join_all(tasks).await?; + let elements_layout = child_layouts.remove(0); + let offsets_layout = child_layouts.remove(0); + let validity_layout = if is_nullable { + Some(child_layouts.remove(0)) + } else { + None + }; + + Ok(ListLayout::try_new( + dtype.clone(), + elements_layout, + offsets_layout, + validity_layout, + )? + .into_layout()) + } + + async fn write_empty_chunk( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + mut eof: SequencePointer, + session: &VortexSession, + dtype: DType, + is_nullable: bool, + ) -> VortexResult { + let elements_dtype = dtype + .as_list_element_opt() + .ok_or_else(|| vortex_err!("ListLayoutStrategy requires a List dtype, got {dtype}"))? + .as_ref() + .clone(); + + let elements_eof = eof.split_off(); + let offsets_eof = eof.split_off(); + let validity_eof = if is_nullable { Some(eof.split_off()) } else { None }; + + let elements_layout = self + .elements + .write_stream( + ctx.clone(), + Arc::clone(&segment_sink), + empty_stream(elements_dtype), + elements_eof, + session, + ) + .await?; + // Offsets buffer of length 1 (a single 0 boundary) representing 0 lists. + // We can't easily synthesize a single-element primitive array here without coupling to + // a builder, so we just emit an empty offsets buffer and rely on `row_count()` returning + // 0 via `saturating_sub`. The reader treats `offsets.row_count() == 0` as 0 rows. + let offsets_layout = self + .offsets + .write_stream( + ctx.clone(), + Arc::clone(&segment_sink), + empty_stream(DType::Primitive(PType::U32, Nullability::NonNullable)), + offsets_eof, + session, + ) + .await?; + let validity_layout = if let Some(validity_eof) = validity_eof { + Some( + self.validity + .write_stream( + ctx, + segment_sink, + empty_stream(DType::Bool(Nullability::NonNullable)), + validity_eof, + session, + ) + .await?, + ) + } else { + None + }; + + Ok(ListLayout::try_new(dtype, elements_layout, offsets_layout, validity_layout)? + .into_layout()) + } +} + +fn empty_stream(dtype: DType) -> SendableSequentialStream { + SequentialStreamAdapter::new(dtype, futures::stream::empty().boxed()).sendable() +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vortex_array::ArrayContext; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ListArray; + use vortex_array::assert_arrays_eq; + use vortex_array::expr::root; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + + use crate::LayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::list::writer::ListLayoutStrategy; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialArrayStreamExt; + use crate::test::SESSION; + + async fn round_trip(list: ArrayRef) { + let segments = Arc::new(TestSegments::default()); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let writer = ListLayoutStrategy::new( + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + ); + + let (ptr, eof) = SequenceId::root().split(); + let stream = list.clone().to_array_stream().sequenced(ptr); + + let layout = writer + .write_stream( + ArrayContext::empty(), + Arc::::clone(&segments), + stream, + eof, + &SESSION, + ) + .await + .unwrap(); + + let reader = layout + .new_reader(Arc::from("test"), segments, &SESSION) + .unwrap(); + + let row_count = usize::try_from(layout.row_count()).unwrap(); + let result = reader + .projection_evaluation( + &(0..layout.row_count()), + &root(), + MaskFuture::new_true(row_count), + ) + .unwrap() + .await + .unwrap(); + + assert_arrays_eq!(result, list); + } + + #[tokio::test] + async fn round_trip_non_nullable() { + let elements = buffer![1i32, 2, 3, 4, 5].into_array(); + let offsets = buffer![0u32, 2, 5, 5].into_array(); // 3 lists: [1,2], [3,4,5], [] + let list = ListArray::try_new(elements, offsets, Validity::NonNullable) + .unwrap() + .into_array(); + round_trip(list).await; + } + + #[tokio::test] + async fn round_trip_nullable() { + let elements = buffer![10i32, 20, 30, 40, 50].into_array(); + let offsets = buffer![0u32, 2, 3, 5].into_array(); // 3 lists + let validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); + let list = ListArray::try_new(elements, offsets, validity) + .unwrap() + .into_array(); + round_trip(list).await; + } +} diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 18df5b8f347..47fa31aa3d9 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -16,6 +16,7 @@ pub mod dict; pub mod file_stats; pub mod flat; pub(crate) mod foreign; +pub mod list; pub(crate) mod partitioned; pub mod repartition; pub mod row_idx; diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index 42c8906d107..a6831abcdcd 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -11,6 +11,7 @@ use crate::LayoutEncodingRef; use crate::layouts::chunked::ChunkedLayoutEncoding; use crate::layouts::dict::DictLayoutEncoding; use crate::layouts::flat::FlatLayoutEncoding; +use crate::layouts::list::ListLayoutEncoding; use crate::layouts::struct_::StructLayoutEncoding; use crate::layouts::zoned::LegacyStatsLayoutEncoding; use crate::layouts::zoned::ZonedLayoutEncoding; @@ -56,6 +57,7 @@ impl Default for LayoutSession { LegacyStatsLayoutEncoding.as_ref(), ); layouts.register(DictLayoutEncoding.id(), DictLayoutEncoding.as_ref()); + layouts.register(ListLayoutEncoding.id(), ListLayoutEncoding.as_ref()); Self { registry: layouts } } From 4655c388858883a04484af4cd2c8324c0ae1c3d7 Mon Sep 17 00:00:00 2001 From: Matthew Katz Date: Fri, 22 May 2026 15:10:37 -0400 Subject: [PATCH 02/37] fix Signed-off-by: Matthew Katz Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 645c25ffdff..794f2a2a776 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -96,7 +96,7 @@ impl VTable for List { match (idx, layout.validity.is_some()) { (ELEMENTS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("elements".into()), (OFFSETS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("offsets".into()), - (VALIDITY_CHILD_INDEX, true) => LayoutChildType::Transparent("validity".into()), + (VALIDITY_CHILD_INDEX, true) => LayoutChildType::Auxiliary("validity".into()), _ => vortex_panic!("Invalid child index {idx} for ListLayout"), } } From 801bc958b54c973b5a9c2d0ca343d2222f026f1f Mon Sep 17 00:00:00 2001 From: Matthew Katz Date: Fri, 22 May 2026 17:43:12 -0400 Subject: [PATCH 03/37] second claude pass Signed-off-by: Matthew Katz Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 29 +- vortex-layout/src/layouts/list/reader.rs | 17 +- vortex-layout/src/layouts/list/writer.rs | 365 ++++++++++++----------- 3 files changed, 221 insertions(+), 190 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 794f2a2a776..519a2e8cd42 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -220,29 +220,28 @@ impl ListLayout { offsets: LayoutRef, validity: Option, ) -> VortexResult { - if !dtype.is_list() { - vortex_bail!("ListLayout requires a List dtype, got {dtype}"); - } + vortex_ensure!( + dtype.is_list(), + "ListLayout requires a List dtype, got {dtype}" + ); + vortex_ensure!( dtype.is_nullable() == validity.is_some(), "validity must be supplied iff dtype is nullable (dtype: {dtype})", ); - if !offsets.dtype().is_int() || offsets.dtype().is_nullable() { - vortex_bail!( - "offsets must be a non-nullable integer layout, got {}", - offsets.dtype() - ); - } + + let offsets_dtype = offsets.dtype(); vortex_ensure!( - offsets.row_count() >= 1, - "ListLayout offsets must have at least 1 row (n+1 entries for n lists), got 0", + offsets_dtype.is_int() && !offsets_dtype.is_nullable(), + "offsets must be a non-nullable integer layout, got {offsets_dtype}", ); - let row_count = offsets.row_count() - 1; + + let offsets_row_count = offsets.row_count().saturating_sub(1); if let Some(validity) = validity.as_ref() { + let validity_row_count = validity.row_count(); vortex_ensure!( - validity.row_count() == row_count, - "validity row count ({}) must match list row count ({row_count})", - validity.row_count(), + validity_row_count == offsets_row_count, + "validity row count ({validity_row_count}) must match list row count ({offsets_row_count})", ); } diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 72164eeb155..c5b8add236d 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -107,25 +107,26 @@ impl LayoutReader for ListReader { ) -> VortexResult<()> { // Offsets has one more row than the list itself but shares the list's chunking // structure, so it's the appropriate child to drive scan splits. - self.offsets.register_splits(field_mask, split_range, splits) + self.offsets + .register_splits(field_mask, split_range, splits) } fn pruning_evaluation( &self, _row_range: &Range, _expr: &Expression, - mask: Mask, + _mask: Mask, ) -> VortexResult { - Ok(MaskFuture::ready(mask)) + todo!() } fn filter_evaluation( &self, _row_range: &Range, _expr: &Expression, - mask: MaskFuture, + _mask: MaskFuture, ) -> VortexResult { - Ok(mask) + todo!() } fn projection_evaluation( @@ -144,9 +145,9 @@ impl LayoutReader for ListReader { let offsets_range = row_range.start..row_range.end + 1; let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; - let elements_fut = self - .elements - .projection_evaluation(&elements_range, &root(), elements_mask)?; + let elements_fut = + self.elements + .projection_evaluation(&elements_range, &root(), elements_mask)?; let offsets_fut = self.offsets.projection_evaluation( &offsets_range, &root(), diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index d6bb727502b..62d1e5917eb 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -5,12 +5,12 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt; -use futures::future::try_join_all; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::list::ListDataParts; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -34,7 +34,7 @@ use crate::sequence::SequentialStream; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; -/// Strategy for writing list-typed arrays as an Apache Arrow-style [`ListLayout`]. +/// Strategy for writing list-typed arrays. /// /// For each input chunk, the strategy: /// 1. Canonicalizes the chunk into a [`ListViewArray`]. @@ -44,8 +44,18 @@ use crate::sequence::SequentialStreamExt; /// separately configurable downstream strategies, producing a single [`ListLayout`] for /// that chunk. /// -/// When the input stream contains multiple chunks, each chunk becomes one `ListLayout` and the -/// strategy wraps them in a [`ChunkedLayout`]. +/// # Multi-chunk handling +/// +/// Each input chunk is a self-contained list array — its own elements buffer, its own +/// `n_i + 1` offsets starting from 0, its own validity. When the stream contains multiple +/// chunks, the strategy produces one `ListLayout` per chunk and wraps them in a +/// [`ChunkedLayout`], rather than merging them into a single `ListLayout` by rebasing offsets +/// across chunks. +/// +/// This mirrors how `ChunkedReader` reads back the column: it returns a `ChunkedArray` of +/// per-chunk `ListArray`s rather than concatenating into one big `ListArray`. The chunk +/// boundary is preserved end-to-end, consistent with how every other dtype's chunked column +/// is handled in Vortex. /// /// [`ListArray`]: vortex_array::arrays::ListArray pub struct ListLayoutStrategy { @@ -87,7 +97,7 @@ impl LayoutStrategy for ListLayoutStrategy { let is_nullable = dtype.is_nullable(); let mut stream = stream; - let mut chunk_layouts: Vec = Vec::new(); + let mut chunk_layouts = Vec::::new(); while let Some(item) = stream.next().await { let (sequence_id, array) = item?; @@ -151,130 +161,55 @@ impl ListLayoutStrategy { array: ArrayRef, ) -> VortexResult { let mut exec_ctx = session.create_execution_ctx(); - - // Canonicalize then rebuild into ZCTL form. - let listview = array.execute::(&mut exec_ctx)?; - let list = list_from_list_view(listview)?; - let parts = list.into_data_parts(); - let row_count = parts.offsets.len().saturating_sub(1); - - let elements_dtype = parts.elements.dtype().clone(); - let offsets_dtype = parts.offsets.dtype().clone(); - - // Build single-chunk streams for each child. Each child sees one array and EOF. - let mut sequence_pointer = sequence_id.descend(); - let elements_seq = sequence_pointer.advance(); - let offsets_seq = sequence_pointer.advance(); - let validity_seq = if is_nullable { - Some(sequence_pointer.advance()) - } else { - None - }; - drop(sequence_pointer); - - let elements_eof = chunk_eof.split_off(); - let offsets_eof = chunk_eof.split_off(); - let validity_eof = if is_nullable { - Some(chunk_eof.split_off()) - } else { - None - }; - drop(chunk_eof); - - let elements_stream = SequentialStreamAdapter::new( - elements_dtype, - futures::stream::once(async move { Ok((elements_seq, parts.elements)) }).boxed(), - ) - .sendable(); - let offsets_stream = SequentialStreamAdapter::new( - offsets_dtype, - futures::stream::once(async move { Ok((offsets_seq, parts.offsets)) }).boxed(), - ) - .sendable(); - - let validity_array = if is_nullable { - Some( - parts - .validity - .execute_mask(row_count, &mut exec_ctx)? - .into_array(), + // Canonicalize to ListView, then rebuild into zero-copy-to-list form with `n+1` + // monotonic offsets. + let ListDataParts { + elements, + offsets, + validity, + .. + } = list_from_list_view(array.execute::(&mut exec_ctx)?)? + .into_data_parts(); + // `offsets` is the Arrow-canonical `n+1` entries, so the list count is one less. + let validity_array = is_nullable + .then(|| { + validity + .execute_mask(offsets.len().saturating_sub(1), &mut exec_ctx) + .map(|m| m.into_array()) + }) + .transpose()?; + + // Closure to spawn one child writer with its own sequence id and EOF pointer, advanced + // in invocation order: elements, offsets, validity (when nullable). + let mut sp = sequence_id.descend(); + let handle = session.handle(); + let mut spawn = |strategy: &Arc, array: ArrayRef| { + spawn_layout_write( + &handle, + Arc::clone(strategy), + ctx.clone(), + Arc::clone(&segment_sink), + session, + single_chunk_stream(array.dtype().clone(), sp.advance(), array), + chunk_eof.split_off(), ) - } else { - None }; - let validity_stream = validity_array.map(|arr| { - let seq = validity_seq.expect("validity sequence id"); - SequentialStreamAdapter::new( - DType::Bool(Nullability::NonNullable), - futures::stream::once(async move { Ok((seq, arr)) }).boxed(), - ) - .sendable() - }); - - let handle = session.handle(); - - let elements_strategy = Arc::clone(&self.elements); - let elements_ctx = ctx.clone(); - let elements_sink = Arc::clone(&segment_sink); - let elements_session = session.clone(); - let elements_task = handle.spawn_nested(move |h| async move { - let session = elements_session.with_handle(h); - elements_strategy - .write_stream( - elements_ctx, - elements_sink, - elements_stream, - elements_eof, - &session, - ) - .await - }); - - let offsets_strategy = Arc::clone(&self.offsets); - let offsets_ctx = ctx.clone(); - let offsets_sink = Arc::clone(&segment_sink); - let offsets_session = session.clone(); - let offsets_task = handle.spawn_nested(move |h| async move { - let session = offsets_session.with_handle(h); - offsets_strategy - .write_stream( - offsets_ctx, - offsets_sink, - offsets_stream, - offsets_eof, - &session, - ) - .await - }); - - let mut tasks = vec![elements_task, offsets_task]; - if let (Some(validity_stream), Some(validity_eof)) = (validity_stream, validity_eof) { - let validity_strategy = Arc::clone(&self.validity); - let validity_ctx = ctx; - let validity_sink = segment_sink; - let validity_session = session.clone(); - tasks.push(handle.spawn_nested(move |h| async move { - let session = validity_session.with_handle(h); - validity_strategy - .write_stream( - validity_ctx, - validity_sink, - validity_stream, - validity_eof, - &session, - ) - .await - })); - } - let mut child_layouts = try_join_all(tasks).await?; - let elements_layout = child_layouts.remove(0); - let offsets_layout = child_layouts.remove(0); - let validity_layout = if is_nullable { - Some(child_layouts.remove(0)) - } else { - None - }; + let elements_task = spawn(&self.elements, elements); + let offsets_task = spawn(&self.offsets, offsets); + let validity_task = validity_array.map(|arr| spawn(&self.validity, arr)); + drop(spawn); + + let (elements_layout, offsets_layout, validity_layout) = futures::try_join!( + elements_task, + offsets_task, + async { + match validity_task { + Some(task) => task.await.map(Some), + None => Ok(None), + } + } + )?; Ok(ListLayout::try_new( dtype.clone(), @@ -285,6 +220,8 @@ impl ListLayoutStrategy { .into_layout()) } + /// Empty-stream variant: produces a `ListLayout` whose children all encode 0 rows. + /// `offsets.row_count() == 0` is treated as 0 lists by [`ListLayout::row_count`]. async fn write_empty_chunk( &self, ctx: ArrayContext, @@ -300,52 +237,35 @@ impl ListLayoutStrategy { .as_ref() .clone(); - let elements_eof = eof.split_off(); - let offsets_eof = eof.split_off(); - let validity_eof = if is_nullable { Some(eof.split_off()) } else { None }; - - let elements_layout = self - .elements - .write_stream( - ctx.clone(), - Arc::clone(&segment_sink), - empty_stream(elements_dtype), - elements_eof, - session, - ) - .await?; - // Offsets buffer of length 1 (a single 0 boundary) representing 0 lists. - // We can't easily synthesize a single-element primitive array here without coupling to - // a builder, so we just emit an empty offsets buffer and rely on `row_count()` returning - // 0 via `saturating_sub`. The reader treats `offsets.row_count() == 0` as 0 rows. - let offsets_layout = self - .offsets - .write_stream( + let mut write_empty = |strategy: &Arc, child_dtype: DType| { + write_empty_child( + Arc::clone(strategy), ctx.clone(), Arc::clone(&segment_sink), - empty_stream(DType::Primitive(PType::U32, Nullability::NonNullable)), - offsets_eof, session, + child_dtype, + eof.split_off(), ) - .await?; - let validity_layout = if let Some(validity_eof) = validity_eof { + }; + + let elements_layout = write_empty(&self.elements, elements_dtype).await?; + let offsets_layout = write_empty( + &self.offsets, + DType::Primitive(PType::U32, Nullability::NonNullable), + ) + .await?; + let validity_layout = if is_nullable { Some( - self.validity - .write_stream( - ctx, - segment_sink, - empty_stream(DType::Bool(Nullability::NonNullable)), - validity_eof, - session, - ) - .await?, + write_empty(&self.validity, DType::Bool(Nullability::NonNullable)).await?, ) } else { None }; - Ok(ListLayout::try_new(dtype, elements_layout, offsets_layout, validity_layout)? - .into_layout()) + Ok( + ListLayout::try_new(dtype, elements_layout, offsets_layout, validity_layout)? + .into_layout(), + ) } } @@ -353,6 +273,52 @@ fn empty_stream(dtype: DType) -> SendableSequentialStream { SequentialStreamAdapter::new(dtype, futures::stream::empty().boxed()).sendable() } +/// Drive a single child writer with an empty stream of the given `dtype`. +async fn write_empty_child( + strategy: Arc, + ctx: ArrayContext, + sink: SegmentSinkRef, + session: &VortexSession, + dtype: DType, + eof: SequencePointer, +) -> VortexResult { + strategy + .write_stream(ctx, sink, empty_stream(dtype), eof, session) + .await +} + +/// Wrap a single array as a one-shot [`SendableSequentialStream`] for handoff to a child writer. +fn single_chunk_stream( + dtype: DType, + sequence_id: SequenceId, + array: ArrayRef, +) -> SendableSequentialStream { + SequentialStreamAdapter::new( + dtype, + futures::stream::once(async move { Ok((sequence_id, array)) }).boxed(), + ) + .sendable() +} + +/// Spawn a child layout writer task onto the session handle. +/// +/// Captures the strategy, ctx, sink, and a cloned session so the spawned future is `'static`. +fn spawn_layout_write( + handle: &vortex_io::runtime::Handle, + strategy: Arc, + ctx: ArrayContext, + sink: SegmentSinkRef, + session: &VortexSession, + stream: SendableSequentialStream, + eof: SequencePointer, +) -> vortex_io::runtime::Task> { + let session = session.clone(); + handle.spawn_nested(move |h| async move { + let session = session.with_handle(h); + strategy.write_stream(ctx, sink, stream, eof, &session).await + }) +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -379,11 +345,8 @@ mod tests { async fn round_trip(list: ArrayRef) { let segments = Arc::new(TestSegments::default()); let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let writer = ListLayoutStrategy::new( - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - ); + let writer = + ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); let (ptr, eof) = SequenceId::root().split(); let stream = list.clone().to_array_stream().sequenced(ptr); @@ -437,4 +400,72 @@ mod tests { .into_array(); round_trip(list).await; } + + /// Writes a list, then reads back only a sub-range to exercise projection over a slice. + async fn round_trip_subset(list: ArrayRef, row_range: std::ops::Range) { + let segments = Arc::new(TestSegments::default()); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let writer = + ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); + + let (ptr, eof) = SequenceId::root().split(); + let stream = list.clone().to_array_stream().sequenced(ptr); + + let layout = writer + .write_stream( + ArrayContext::empty(), + Arc::::clone(&segments), + stream, + eof, + &SESSION, + ) + .await + .unwrap(); + + let reader = layout + .new_reader(Arc::from("test"), segments, &SESSION) + .unwrap(); + + let mask_len = usize::try_from(row_range.end - row_range.start).unwrap(); + let result = reader + .projection_evaluation(&row_range, &root(), MaskFuture::new_true(mask_len)) + .unwrap() + .await + .unwrap(); + + let expected = list + .slice( + usize::try_from(row_range.start).unwrap() + ..usize::try_from(row_range.end).unwrap(), + ) + .unwrap(); + assert_arrays_eq!(result, expected); + } + + #[tokio::test] + async fn round_trip_subset_non_nullable() { + // 5 lists: [1,2], [3], [], [4,5,6], [7] + let elements = buffer![1i32, 2, 3, 4, 5, 6, 7].into_array(); + let offsets = buffer![0u32, 2, 3, 3, 6, 7].into_array(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable) + .unwrap() + .into_array(); + // Read the middle three lists: [3], [], [4,5,6] + round_trip_subset(list, 1..4).await; + } + + #[tokio::test] + async fn round_trip_subset_nullable() { + // 4 lists with validity [true, false, true, true]: + // [10,20], null, [30], [40,50,60] + let elements = buffer![10i32, 20, 30, 40, 50, 60].into_array(); + let offsets = buffer![0u32, 2, 2, 3, 6].into_array(); + let validity = + Validity::Array(BoolArray::from_iter([true, false, true, true]).into_array()); + let list = ListArray::try_new(elements, offsets, validity) + .unwrap() + .into_array(); + // Read lists 1..3: null, [30] + round_trip_subset(list, 1..3).await; + } } From d72e53d9b471093da8f7c501c7c32a885885d9db Mon Sep 17 00:00:00 2001 From: Matthew Katz Date: Tue, 26 May 2026 11:44:25 -0400 Subject: [PATCH 04/37] small fixes Signed-off-by: Matthew Katz Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 67 ++++++++---------------- vortex-layout/src/layouts/list/reader.rs | 2 +- vortex-layout/src/layouts/list/writer.rs | 41 ++++++--------- 3 files changed, 38 insertions(+), 72 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 519a2e8cd42..7628ffaca81 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -68,7 +68,7 @@ impl VTable for List { } fn metadata(layout: &Self::Layout) -> Self::Metadata { - ProstMetadata(ListLayoutMetadata::new(layout.offsets_ptype)) + ProstMetadata(ListLayoutMetadata::new(layout.offsets_ptype())) } fn segment_ids(_layout: &Self::Layout) -> Vec { @@ -124,10 +124,6 @@ impl VTable for List { children: &dyn LayoutChildren, _ctx: &ReadContext, ) -> VortexResult { - let elements_dtype = dtype - .as_list_element_opt() - .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {dtype}"))?; - let expected_children = if dtype.is_nullable() { NUM_CHILDREN_NULLABLE } else { @@ -139,9 +135,14 @@ impl VTable for List { children.nchildren(), ); + let elements_dtype = dtype + .as_list_element_opt() + .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {dtype}"))?; let elements = children.child(ELEMENTS_CHILD_INDEX, elements_dtype.as_ref())?; + let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); let offsets = children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; + let validity = if dtype.is_nullable() { Some(children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))?) } else { @@ -153,7 +154,6 @@ impl VTable for List { elements, offsets, validity, - offsets_ptype: metadata.offsets_ptype(), }) } @@ -205,60 +205,36 @@ pub struct ListLayout { elements: LayoutRef, offsets: LayoutRef, validity: Option, - offsets_ptype: PType, } impl ListLayout { /// Construct a new `ListLayout` from its components. /// - /// `dtype` must be a [`DType::List`]. The `validity` child must be supplied iff the dtype is - /// nullable. The list's row count is derived from the `offsets` child as - /// `offsets.row_count() - 1`. - pub fn try_new( + /// # Invariants (caller's responsibility) + /// + /// - `dtype` must be a [`DType::List`]. + /// - `validity` must be `Some` iff `dtype.is_nullable()`. + /// - `offsets.dtype()` must be a non-nullable integer. + /// - `offsets.row_count()` is the Arrow-canonical `n+1` for `n` lists (or `0` for empty). + /// - When present, `validity.row_count() == offsets.row_count().saturating_sub(1)`. + /// + /// The [`ListLayoutStrategy`](writer::ListLayoutStrategy) writer upholds these invariants by + /// construction. + pub fn new( dtype: DType, elements: LayoutRef, offsets: LayoutRef, validity: Option, - ) -> VortexResult { - vortex_ensure!( - dtype.is_list(), - "ListLayout requires a List dtype, got {dtype}" - ); - - vortex_ensure!( - dtype.is_nullable() == validity.is_some(), - "validity must be supplied iff dtype is nullable (dtype: {dtype})", - ); - - let offsets_dtype = offsets.dtype(); - vortex_ensure!( - offsets_dtype.is_int() && !offsets_dtype.is_nullable(), - "offsets must be a non-nullable integer layout, got {offsets_dtype}", - ); - - let offsets_row_count = offsets.row_count().saturating_sub(1); - if let Some(validity) = validity.as_ref() { - let validity_row_count = validity.row_count(); - vortex_ensure!( - validity_row_count == offsets_row_count, - "validity row count ({validity_row_count}) must match list row count ({offsets_row_count})", - ); - } - - let offsets_ptype = offsets.dtype().as_ptype(); - - Ok(Self { + ) -> Self { + Self { dtype, elements, offsets, validity, - offsets_ptype, - }) + } } /// Number of lists in this layout. - /// - /// Equal to `offsets.row_count() - 1` by construction. #[inline] pub fn row_count(&self) -> u64 { self.offsets.row_count().saturating_sub(1) @@ -279,9 +255,10 @@ impl ListLayout { self.validity.as_ref() } + /// The integer type used for the `offsets` child layout. #[inline] pub fn offsets_ptype(&self) -> PType { - self.offsets_ptype + self.offsets.dtype().as_ptype() } /// The dtype of the inner elements column. diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index c5b8add236d..ff8fddddd0a 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -139,7 +139,7 @@ impl LayoutReader for ListReader { let elements_range = 0..elements_len; let elements_mask = MaskFuture::new_true(usize::try_from(elements_len)?); - let row_count = usize::try_from(row_range.end - row_range.start)?; + let row_count: usize = usize::try_from(row_range.end - row_range.start)?; // The offsets child has n+1 entries; reading row_range maps to offsets in // [row_range.start..row_range.end + 1). let offsets_range = row_range.start..row_range.end + 1; diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 62d1e5917eb..52eb2e28d3d 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -168,8 +168,7 @@ impl ListLayoutStrategy { offsets, validity, .. - } = list_from_list_view(array.execute::(&mut exec_ctx)?)? - .into_data_parts(); + } = list_from_list_view(array.execute::(&mut exec_ctx)?)?.into_data_parts(); // `offsets` is the Arrow-canonical `n+1` entries, so the list count is one less. let validity_array = is_nullable .then(|| { @@ -200,24 +199,18 @@ impl ListLayoutStrategy { let validity_task = validity_array.map(|arr| spawn(&self.validity, arr)); drop(spawn); - let (elements_layout, offsets_layout, validity_layout) = futures::try_join!( - elements_task, - offsets_task, - async { + let (elements_layout, offsets_layout, validity_layout) = + futures::try_join!(elements_task, offsets_task, async { match validity_task { Some(task) => task.await.map(Some), None => Ok(None), } - } - )?; - - Ok(ListLayout::try_new( - dtype.clone(), - elements_layout, - offsets_layout, - validity_layout, - )? - .into_layout()) + })?; + + Ok( + ListLayout::new(dtype.clone(), elements_layout, offsets_layout, validity_layout) + .into_layout(), + ) } /// Empty-stream variant: produces a `ListLayout` whose children all encode 0 rows. @@ -255,17 +248,12 @@ impl ListLayoutStrategy { ) .await?; let validity_layout = if is_nullable { - Some( - write_empty(&self.validity, DType::Bool(Nullability::NonNullable)).await?, - ) + Some(write_empty(&self.validity, DType::Bool(Nullability::NonNullable)).await?) } else { None }; - Ok( - ListLayout::try_new(dtype, elements_layout, offsets_layout, validity_layout)? - .into_layout(), - ) + Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) } } @@ -315,7 +303,9 @@ fn spawn_layout_write( let session = session.clone(); handle.spawn_nested(move |h| async move { let session = session.with_handle(h); - strategy.write_stream(ctx, sink, stream, eof, &session).await + strategy + .write_stream(ctx, sink, stream, eof, &session) + .await }) } @@ -435,8 +425,7 @@ mod tests { let expected = list .slice( - usize::try_from(row_range.start).unwrap() - ..usize::try_from(row_range.end).unwrap(), + usize::try_from(row_range.start).unwrap()..usize::try_from(row_range.end).unwrap(), ) .unwrap(); assert_arrays_eq!(result, expected); From 0337c2a4dad9aa075f8abb53ba7871ce069a8865 Mon Sep 17 00:00:00 2001 From: Matthew Katz Date: Tue, 26 May 2026 15:14:44 -0400 Subject: [PATCH 05/37] fix Signed-off-by: Matthew Katz Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 314 ++++++++++++----------- 1 file changed, 162 insertions(+), 152 deletions(-) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 52eb2e28d3d..0f30a1a7da9 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -13,18 +13,14 @@ use vortex_array::arrays::ListViewArray; use vortex_array::arrays::list::ListDataParts; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; use vortex_error::VortexResult; -use vortex_error::vortex_err; +use vortex_error::vortex_bail; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; -use crate::children::OwnedLayoutChildren; -use crate::layouts::chunked::ChunkedLayout; use crate::layouts::list::ListLayout; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; @@ -36,26 +32,21 @@ use crate::sequence::SequentialStreamExt; /// Strategy for writing list-typed arrays. /// -/// For each input chunk, the strategy: -/// 1. Canonicalizes the chunk into a [`ListViewArray`]. +/// Single-chunk only. The strategy: +/// 1. Canonicalizes the input chunk into a [`ListViewArray`]. /// 2. Calls [`list_from_list_view`] to rebuild it into zero-copy-to-list form /// (sorted, gapless, non-overlapping offsets) and produce a [`ListArray`]. /// 3. Writes the `elements`, `n+1` `offsets`, and (when nullable) `validity` columns into -/// separately configurable downstream strategies, producing a single [`ListLayout`] for -/// that chunk. +/// separately configurable downstream strategies, producing a single [`ListLayout`]. /// -/// # Multi-chunk handling +/// # Chunking /// -/// Each input chunk is a self-contained list array — its own elements buffer, its own -/// `n_i + 1` offsets starting from 0, its own validity. When the stream contains multiple -/// chunks, the strategy produces one `ListLayout` per chunk and wraps them in a -/// [`ChunkedLayout`], rather than merging them into a single `ListLayout` by rebasing offsets -/// across chunks. -/// -/// This mirrors how `ChunkedReader` reads back the column: it returns a `ChunkedArray` of -/// per-chunk `ListArray`s rather than concatenating into one big `ListArray`. The chunk -/// boundary is preserved end-to-end, consistent with how every other dtype's chunked column -/// is handled in Vortex. +/// `ListLayoutStrategy` bails on empty or multi-chunk input, matching the convention used by +/// [`FlatLayoutStrategy`](crate::layouts::flat::writer::FlatLayoutStrategy). To handle +/// multi-chunk input, wrap with +/// [`ChunkedLayoutStrategy`](crate::layouts::chunked::writer::ChunkedLayoutStrategy), which +/// slices the input into one-chunk substreams and produces +/// `ChunkedLayout(ListLayout × K)`. /// /// [`ListArray`]: vortex_array::arrays::ListArray pub struct ListLayoutStrategy { @@ -84,82 +75,20 @@ impl LayoutStrategy for ListLayoutStrategy { &self, ctx: ArrayContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, + mut stream: SendableSequentialStream, mut eof: SequencePointer, session: &VortexSession, ) -> VortexResult { let dtype = stream.dtype().clone(); if !dtype.is_list() { - return Err(vortex_err!( - "ListLayoutStrategy requires a List dtype, got {dtype}" - )); - } - let is_nullable = dtype.is_nullable(); - - let mut stream = stream; - let mut chunk_layouts = Vec::::new(); - - while let Some(item) = stream.next().await { - let (sequence_id, array) = item?; - let chunk_eof = eof.split_off(); - - let chunk_layout = self - .write_chunk( - ctx.clone(), - Arc::clone(&segment_sink), - sequence_id, - chunk_eof, - session, - &dtype, - is_nullable, - array, - ) - .await?; - chunk_layouts.push(chunk_layout); - } - - if chunk_layouts.is_empty() { - // Empty stream: emit an empty ListLayout. We write empty children of the same dtype - // (with u32 offsets) so the children's encoding ID is still derivable. - let chunk_layout = self - .write_empty_chunk(ctx, segment_sink, eof, session, dtype.clone(), is_nullable) - .await?; - return Ok(chunk_layout); + vortex_bail!("ListLayoutStrategy requires a List dtype, got {dtype}"); } - if chunk_layouts.len() == 1 { - return Ok(chunk_layouts.pop().expect("len == 1")); - } - - let row_count = chunk_layouts.iter().map(|l| l.row_count()).sum(); - Ok(ChunkedLayout::new( - row_count, - dtype, - OwnedLayoutChildren::layout_children(chunk_layouts), - ) - .into_layout()) - } - - fn buffered_bytes(&self) -> u64 { - self.elements.buffered_bytes() - + self.offsets.buffered_bytes() - + self.validity.buffered_bytes() - } -} + let Some(item) = stream.next().await else { + vortex_bail!("ListLayoutStrategy needs a single chunk"); + }; + let (sequence_id, array) = item?; -impl ListLayoutStrategy { - #[allow(clippy::too_many_arguments)] - async fn write_chunk( - &self, - ctx: ArrayContext, - segment_sink: SegmentSinkRef, - sequence_id: SequenceId, - mut chunk_eof: SequencePointer, - session: &VortexSession, - dtype: &DType, - is_nullable: bool, - array: ArrayRef, - ) -> VortexResult { let mut exec_ctx = session.create_execution_ctx(); // Canonicalize to ListView, then rebuild into zero-copy-to-list form with `n+1` // monotonic offsets. @@ -170,7 +99,8 @@ impl ListLayoutStrategy { .. } = list_from_list_view(array.execute::(&mut exec_ctx)?)?.into_data_parts(); // `offsets` is the Arrow-canonical `n+1` entries, so the list count is one less. - let validity_array = is_nullable + let validity_array = dtype + .is_nullable() .then(|| { validity .execute_mask(offsets.len().saturating_sub(1), &mut exec_ctx) @@ -178,8 +108,8 @@ impl ListLayoutStrategy { }) .transpose()?; - // Closure to spawn one child writer with its own sequence id and EOF pointer, advanced - // in invocation order: elements, offsets, validity (when nullable). + // Closure to spawn one child writer with its own sequence id and EOF pointer, + // advanced in invocation order: elements, offsets, validity (when nullable). let mut sp = sequence_id.descend(); let handle = session.handle(); let mut spawn = |strategy: &Arc, array: ArrayRef| { @@ -190,7 +120,7 @@ impl ListLayoutStrategy { Arc::clone(&segment_sink), session, single_chunk_stream(array.dtype().clone(), sp.advance(), array), - chunk_eof.split_off(), + eof.split_off(), ) }; @@ -207,74 +137,23 @@ impl ListLayoutStrategy { } })?; + if stream.next().await.is_some() { + vortex_bail!("ListLayoutStrategy received stream with more than a single chunk"); + } + Ok( - ListLayout::new(dtype.clone(), elements_layout, offsets_layout, validity_layout) + ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout) .into_layout(), ) } - /// Empty-stream variant: produces a `ListLayout` whose children all encode 0 rows. - /// `offsets.row_count() == 0` is treated as 0 lists by [`ListLayout::row_count`]. - async fn write_empty_chunk( - &self, - ctx: ArrayContext, - segment_sink: SegmentSinkRef, - mut eof: SequencePointer, - session: &VortexSession, - dtype: DType, - is_nullable: bool, - ) -> VortexResult { - let elements_dtype = dtype - .as_list_element_opt() - .ok_or_else(|| vortex_err!("ListLayoutStrategy requires a List dtype, got {dtype}"))? - .as_ref() - .clone(); - - let mut write_empty = |strategy: &Arc, child_dtype: DType| { - write_empty_child( - Arc::clone(strategy), - ctx.clone(), - Arc::clone(&segment_sink), - session, - child_dtype, - eof.split_off(), - ) - }; - - let elements_layout = write_empty(&self.elements, elements_dtype).await?; - let offsets_layout = write_empty( - &self.offsets, - DType::Primitive(PType::U32, Nullability::NonNullable), - ) - .await?; - let validity_layout = if is_nullable { - Some(write_empty(&self.validity, DType::Bool(Nullability::NonNullable)).await?) - } else { - None - }; - - Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) + fn buffered_bytes(&self) -> u64 { + self.elements.buffered_bytes() + + self.offsets.buffered_bytes() + + self.validity.buffered_bytes() } } -fn empty_stream(dtype: DType) -> SendableSequentialStream { - SequentialStreamAdapter::new(dtype, futures::stream::empty().boxed()).sendable() -} - -/// Drive a single child writer with an empty stream of the given `dtype`. -async fn write_empty_child( - strategy: Arc, - ctx: ArrayContext, - sink: SegmentSinkRef, - session: &VortexSession, - dtype: DType, - eof: SequencePointer, -) -> VortexResult { - strategy - .write_stream(ctx, sink, empty_stream(dtype), eof, session) - .await -} - /// Wrap a single array as a one-shot [`SendableSequentialStream`] for handoff to a child writer. fn single_chunk_stream( dtype: DType, @@ -457,4 +336,135 @@ mod tests { // Read lists 1..3: null, [30] round_trip_subset(list, 1..3).await; } + + // -- tree shape visualization --------------------------------------------------------- + // + // These tests are mostly for development/inspection — they show what the resulting + // layout tree looks like for various input shapes. Run with `--nocapture` to see the + // pretty-printed trees: + // + // cargo test -p vortex-layout layouts::list::writer::tests::tree -- --nocapture + + use vortex_array::ArrayContext as _ArrayContextAlias; + + /// Write `array` directly through `ListLayoutStrategy` (no ChunkedLayoutStrategy wrap) + /// and return the resulting top-level layout. + async fn write_through_list_strategy(array: ArrayRef) -> crate::LayoutRef { + let segments = Arc::new(TestSegments::default()); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let writer = + ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + writer + .write_stream( + _ArrayContextAlias::empty(), + segments, + stream, + eof, + &SESSION, + ) + .await + .unwrap() + } + + /// Wrap `ListLayoutStrategy` in `ChunkedLayoutStrategy` and write `array`. For a chunked + /// input, each chunk becomes one `ListLayout` under the outer `ChunkedLayout`. + async fn write_through_chunked_list_strategy(array: ArrayRef) -> crate::LayoutRef { + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + let segments = Arc::new(TestSegments::default()); + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let list_strategy = + ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); + let writer = ChunkedLayoutStrategy::new(list_strategy); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + writer + .write_stream( + _ArrayContextAlias::empty(), + segments, + stream, + eof, + &SESSION, + ) + .await + .unwrap() + } + + #[tokio::test] + async fn tree_shape_single_chunk_non_nullable() { + let elements = buffer![1i32, 2, 3, 4, 5].into_array(); + let offsets = buffer![0u32, 2, 5, 5].into_array(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable) + .unwrap() + .into_array(); + + let layout = write_through_list_strategy(list).await; + let tree = layout.display_tree().to_string(); + eprintln!("--- single-chunk non-nullable ---\n{tree}"); + // Top level is a single ListLayout with 2 children (elements, offsets). + assert!(tree.starts_with("vortex.list")); + assert!(tree.contains("elements")); + assert!(tree.contains("offsets")); + assert!(!tree.contains("validity")); + } + + #[tokio::test] + async fn tree_shape_single_chunk_nullable() { + let elements = buffer![10i32, 20, 30, 40, 50].into_array(); + let offsets = buffer![0u32, 2, 3, 5].into_array(); + let validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); + let list = ListArray::try_new(elements, offsets, validity) + .unwrap() + .into_array(); + + let layout = write_through_list_strategy(list).await; + let tree = layout.display_tree().to_string(); + eprintln!("--- single-chunk nullable ---\n{tree}"); + // Top level is a single ListLayout with 3 children (elements, offsets, validity). + assert!(tree.starts_with("vortex.list")); + assert!(tree.contains("elements")); + assert!(tree.contains("offsets")); + assert!(tree.contains("validity")); + } + + #[tokio::test] + async fn tree_shape_multi_chunk_via_chunked_strategy() { + use vortex_array::arrays::ChunkedArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use std::sync::Arc as StdArc; + + // Two list-array chunks fed through ChunkedArray -> ChunkedLayoutStrategy. + let chunk0 = ListArray::try_new( + buffer![1i32, 2, 3].into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let chunk1 = ListArray::try_new( + buffer![4i32, 5, 6, 7].into_array(), + buffer![0u32, 1, 4].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + + let list_dtype = DType::List( + StdArc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ); + let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], list_dtype) + .unwrap() + .into_array(); + + let layout = write_through_chunked_list_strategy(chunked).await; + let tree = layout.display_tree().to_string(); + eprintln!("--- multi-chunk via ChunkedLayoutStrategy ---\n{tree}"); + // Top level is a ChunkedLayout containing two ListLayouts. + assert!(tree.starts_with("vortex.chunked")); + assert_eq!(tree.matches("vortex.list").count(), 2); + } } From 90fd66ce9a283366681a6700ca3b6e59de9ac84a Mon Sep 17 00:00:00 2001 From: Matthew Katz Date: Tue, 26 May 2026 16:05:17 -0400 Subject: [PATCH 06/37] fix Signed-off-by: Matthew Katz Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 90 +++++++++++------------- 1 file changed, 42 insertions(+), 48 deletions(-) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 0f30a1a7da9..57a5773b4cd 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -83,68 +83,60 @@ impl LayoutStrategy for ListLayoutStrategy { if !dtype.is_list() { vortex_bail!("ListLayoutStrategy requires a List dtype, got {dtype}"); } + let (sequence_id, array) = take_single_chunk(&mut stream).await?; - let Some(item) = stream.next().await else { - vortex_bail!("ListLayoutStrategy needs a single chunk"); - }; - let (sequence_id, array) = item?; - - let mut exec_ctx = session.create_execution_ctx(); // Canonicalize to ListView, then rebuild into zero-copy-to-list form with `n+1` // monotonic offsets. + let mut exec_ctx = session.create_execution_ctx(); let ListDataParts { elements, offsets, validity, .. } = list_from_list_view(array.execute::(&mut exec_ctx)?)?.into_data_parts(); + // `offsets` is the Arrow-canonical `n+1` entries, so the list count is one less. + let row_count = offsets.len().saturating_sub(1); let validity_array = dtype .is_nullable() .then(|| { validity - .execute_mask(offsets.len().saturating_sub(1), &mut exec_ctx) + .execute_mask(row_count, &mut exec_ctx) .map(|m| m.into_array()) }) .transpose()?; - // Closure to spawn one child writer with its own sequence id and EOF pointer, - // advanced in invocation order: elements, offsets, validity (when nullable). + // Collect (strategy, array) pairs for each child to write. Validity is optional. + let mut child_writes: Vec<(Arc, ArrayRef)> = vec![ + (self.elements.clone(), elements), + (self.offsets.clone(), offsets), + ]; + if let Some(arr) = validity_array { + child_writes.push((self.validity.clone(), arr)); + } + + // Spawn one task per child writer. let mut sp = sequence_id.descend(); let handle = session.handle(); - let mut spawn = |strategy: &Arc, array: ArrayRef| { + let tasks = child_writes.into_iter().map(|(strategy, array)| { spawn_layout_write( &handle, - Arc::clone(strategy), + strategy, ctx.clone(), Arc::clone(&segment_sink), session, single_chunk_stream(array.dtype().clone(), sp.advance(), array), eof.split_off(), ) - }; - - let elements_task = spawn(&self.elements, elements); - let offsets_task = spawn(&self.offsets, offsets); - let validity_task = validity_array.map(|arr| spawn(&self.validity, arr)); - drop(spawn); - - let (elements_layout, offsets_layout, validity_layout) = - futures::try_join!(elements_task, offsets_task, async { - match validity_task { - Some(task) => task.await.map(Some), - None => Ok(None), - } - })?; - - if stream.next().await.is_some() { - vortex_bail!("ListLayoutStrategy received stream with more than a single chunk"); - } + }); + let mut layouts = futures::future::try_join_all(tasks).await?.into_iter(); - Ok( - ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout) - .into_layout(), - ) + // Same construction order as `child_writes`: elements, offsets, optional validity. + let elements_layout = layouts.next().expect("elements task"); + let offsets_layout = layouts.next().expect("offsets task"); + let validity_layout = layouts.next(); + + Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) } fn buffered_bytes(&self) -> u64 { @@ -154,6 +146,20 @@ impl LayoutStrategy for ListLayoutStrategy { } } +/// Pull the lone chunk out of a single-chunk stream, erroring on empty or multi-chunk input. +async fn take_single_chunk( + stream: &mut SendableSequentialStream, +) -> VortexResult<(SequenceId, ArrayRef)> { + let Some(chunk) = stream.next().await else { + vortex_bail!("ListLayoutStrategy needs a single chunk"); + }; + let chunk = chunk?; + if stream.next().await.is_some() { + vortex_bail!("ListLayoutStrategy received more than a single chunk"); + } + Ok(chunk) +} + /// Wrap a single array as a one-shot [`SendableSequentialStream`] for handoff to a child writer. fn single_chunk_stream( dtype: DType, @@ -357,13 +363,7 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let stream = array.to_array_stream().sequenced(ptr); writer - .write_stream( - _ArrayContextAlias::empty(), - segments, - stream, - eof, - &SESSION, - ) + .write_stream(_ArrayContextAlias::empty(), segments, stream, eof, &SESSION) .await .unwrap() } @@ -380,13 +380,7 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let stream = array.to_array_stream().sequenced(ptr); writer - .write_stream( - _ArrayContextAlias::empty(), - segments, - stream, - eof, - &SESSION, - ) + .write_stream(_ArrayContextAlias::empty(), segments, stream, eof, &SESSION) .await .unwrap() } @@ -430,11 +424,11 @@ mod tests { #[tokio::test] async fn tree_shape_multi_chunk_via_chunked_strategy() { + use std::sync::Arc as StdArc; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use std::sync::Arc as StdArc; // Two list-array chunks fed through ChunkedArray -> ChunkedLayoutStrategy. let chunk0 = ListArray::try_new( From 810eba35681488423bf918a16895ae1b659ec376 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 12:22:36 -0400 Subject: [PATCH 07/37] clean mod Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 55 +++++++++++---------------- 1 file changed, 22 insertions(+), 33 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 7628ffaca81..c8671bbb588 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -124,16 +124,7 @@ impl VTable for List { children: &dyn LayoutChildren, _ctx: &ReadContext, ) -> VortexResult { - let expected_children = if dtype.is_nullable() { - NUM_CHILDREN_NULLABLE - } else { - NUM_CHILDREN_NON_NULLABLE - }; - vortex_ensure!( - children.nchildren() == expected_children, - "ListLayout expects {expected_children} children, got {}", - children.nchildren(), - ); + validate_children(dtype, children.nchildren())?; let elements_dtype = dtype .as_list_element_opt() @@ -158,16 +149,8 @@ impl VTable for List { } fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { - let expected = if layout.dtype.is_nullable() { - NUM_CHILDREN_NULLABLE - } else { - NUM_CHILDREN_NON_NULLABLE - }; - vortex_ensure!( - children.len() == expected, - "ListLayout expects {expected} children, got {}", - children.len() - ); + validate_children(layout.dtype(), children.len())?; + let mut iter = children.into_iter(); layout.elements = iter .next() @@ -187,18 +170,27 @@ impl VTable for List { } } +/// Validates expected number of children based on `dtype` +#[inline] +fn validate_children(dtype: &DType, n_children: usize) -> VortexResult<()> { + let mut expected = NUM_CHILDREN_NON_NULLABLE; + + if dtype.is_nullable() { + expected += 1; + }; + + vortex_ensure!( + n_children == expected, + "ListLayout expects {expected} children, got {n_children}", + ); + + Ok(()) +} + #[derive(Debug)] pub struct ListLayoutEncoding; -/// Stores a list-typed column in Apache Arrow-compatible form: a flattened `elements` buffer -/// plus an `offsets` buffer of length `n+1` (and optionally a `validity` bitmap of length `n`). -/// -/// Mirrors the in-memory [`ListArray`](vortex_array::arrays::ListArray): offsets are monotonic -/// and `list[i] = elements[offsets[i]..offsets[i+1]]`. -/// -/// To write a [`ListViewArray`](vortex_array::arrays::ListViewArray) (the canonical encoding for -/// [`DType::List`]) into a `ListLayout`, the writer rebuilds it into zero-copy-to-list form via -/// [`list_from_list_view`](vortex_array::arrays::listview::list_from_list_view). +/// Stores a list-typed array by shredding `elements`, `offsets`, and optional `validity` children. #[derive(Clone, Debug)] pub struct ListLayout { dtype: DType, @@ -210,16 +202,13 @@ pub struct ListLayout { impl ListLayout { /// Construct a new `ListLayout` from its components. /// - /// # Invariants (caller's responsibility) + /// # Invariants /// /// - `dtype` must be a [`DType::List`]. /// - `validity` must be `Some` iff `dtype.is_nullable()`. /// - `offsets.dtype()` must be a non-nullable integer. /// - `offsets.row_count()` is the Arrow-canonical `n+1` for `n` lists (or `0` for empty). /// - When present, `validity.row_count() == offsets.row_count().saturating_sub(1)`. - /// - /// The [`ListLayoutStrategy`](writer::ListLayoutStrategy) writer upholds these invariants by - /// construction. pub fn new( dtype: DType, elements: LayoutRef, From b3a3c64b4d7fbc24cdc469cd0d44186d386d712c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 13:09:05 -0400 Subject: [PATCH 08/37] fix writer Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 132 +++++++++-------------- 1 file changed, 53 insertions(+), 79 deletions(-) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 57a5773b4cd..0657f94d624 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt; +use futures::stream; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -36,17 +37,13 @@ use crate::sequence::SequentialStreamExt; /// 1. Canonicalizes the input chunk into a [`ListViewArray`]. /// 2. Calls [`list_from_list_view`] to rebuild it into zero-copy-to-list form /// (sorted, gapless, non-overlapping offsets) and produce a [`ListArray`]. -/// 3. Writes the `elements`, `n+1` `offsets`, and (when nullable) `validity` columns into +/// 3. Writes the `elements`, `offsets`, and (when nullable) `validity` columns into /// separately configurable downstream strategies, producing a single [`ListLayout`]. /// /// # Chunking /// /// `ListLayoutStrategy` bails on empty or multi-chunk input, matching the convention used by -/// [`FlatLayoutStrategy`](crate::layouts::flat::writer::FlatLayoutStrategy). To handle -/// multi-chunk input, wrap with -/// [`ChunkedLayoutStrategy`](crate::layouts::chunked::writer::ChunkedLayoutStrategy), which -/// slices the input into one-chunk substreams and produces -/// `ChunkedLayout(ListLayout × K)`. +/// [`FlatLayoutStrategy`](crate::layouts::flat::writer::FlatLayoutStrategy). /// /// [`ListArray`]: vortex_array::arrays::ListArray pub struct ListLayoutStrategy { @@ -83,10 +80,14 @@ impl LayoutStrategy for ListLayoutStrategy { if !dtype.is_list() { vortex_bail!("ListLayoutStrategy requires a List dtype, got {dtype}"); } - let (sequence_id, array) = take_single_chunk(&mut stream).await?; - // Canonicalize to ListView, then rebuild into zero-copy-to-list form with `n+1` - // monotonic offsets. + // Writer wants exactly one chunk + let Some(chunk) = stream.next().await else { + vortex_bail!("ListLayoutStrategy needs a single chunk"); + }; + let (sequence_id, array) = chunk?; + + // Canonicalize to ListView, then rebuild into zctl let mut exec_ctx = session.create_execution_ctx(); let ListDataParts { elements, @@ -95,46 +96,54 @@ impl LayoutStrategy for ListLayoutStrategy { .. } = list_from_list_view(array.execute::(&mut exec_ctx)?)?.into_data_parts(); - // `offsets` is the Arrow-canonical `n+1` entries, so the list count is one less. + // There is one extra element in `offsets` let row_count = offsets.len().saturating_sub(1); - let validity_array = dtype - .is_nullable() - .then(|| { + let validity_array = if dtype.is_nullable() { + Some( validity - .execute_mask(row_count, &mut exec_ctx) - .map(|m| m.into_array()) - }) - .transpose()?; - - // Collect (strategy, array) pairs for each child to write. Validity is optional. - let mut child_writes: Vec<(Arc, ArrayRef)> = vec![ - (self.elements.clone(), elements), - (self.offsets.clone(), offsets), - ]; - if let Some(arr) = validity_array { - child_writes.push((self.validity.clone(), arr)); - } + .execute_mask(row_count, &mut exec_ctx)? + .into_array(), + ) + } else { + None + }; - // Spawn one task per child writer. - let mut sp = sequence_id.descend(); + // Spawn each child write onto the runtime so they run concurrently let handle = session.handle(); - let tasks = child_writes.into_iter().map(|(strategy, array)| { - spawn_layout_write( - &handle, - strategy, - ctx.clone(), - Arc::clone(&segment_sink), - session, - single_chunk_stream(array.dtype().clone(), sp.advance(), array), - eof.split_off(), + let (elements_task, offsets_task, validity_task) = { + let mut sp = sequence_id.descend(); + let mut spawn_layout_writer = |strategy: Arc, array: ArrayRef| { + let stream = single_chunk_stream(array.dtype().clone(), sp.advance(), array); + let child_eof = eof.split_off(); + let ctx = ctx.clone(); + let segment_sink = segment_sink.clone(); + let session = session.clone(); + handle.spawn_nested(move |h| async move { + let session = session.with_handle(h); + strategy + .write_stream(ctx, segment_sink, stream, child_eof, &session) + .await + }) + }; + ( + spawn_layout_writer(self.elements.clone(), elements), + spawn_layout_writer(self.offsets.clone(), offsets), + validity_array.map(|arr| spawn_layout_writer(self.validity.clone(), arr)), ) - }); - let mut layouts = futures::future::try_join_all(tasks).await?.into_iter(); + }; - // Same construction order as `child_writes`: elements, offsets, optional validity. - let elements_layout = layouts.next().expect("elements task"); - let offsets_layout = layouts.next().expect("offsets task"); - let validity_layout = layouts.next(); + // Should not have more than one chunk + if stream.next().await.is_some() { + vortex_bail!("ListLayoutStrategy received more than a single chunk"); + } + + let (elements_layout, offsets_layout, validity_layout) = + futures::try_join!(elements_task, offsets_task, async move { + match validity_task { + Some(t) => t.await.map(Some), + None => Ok(None), + } + },)?; Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) } @@ -146,20 +155,6 @@ impl LayoutStrategy for ListLayoutStrategy { } } -/// Pull the lone chunk out of a single-chunk stream, erroring on empty or multi-chunk input. -async fn take_single_chunk( - stream: &mut SendableSequentialStream, -) -> VortexResult<(SequenceId, ArrayRef)> { - let Some(chunk) = stream.next().await else { - vortex_bail!("ListLayoutStrategy needs a single chunk"); - }; - let chunk = chunk?; - if stream.next().await.is_some() { - vortex_bail!("ListLayoutStrategy received more than a single chunk"); - } - Ok(chunk) -} - /// Wrap a single array as a one-shot [`SendableSequentialStream`] for handoff to a child writer. fn single_chunk_stream( dtype: DType, @@ -168,32 +163,11 @@ fn single_chunk_stream( ) -> SendableSequentialStream { SequentialStreamAdapter::new( dtype, - futures::stream::once(async move { Ok((sequence_id, array)) }).boxed(), + stream::once(async move { Ok((sequence_id, array)) }).boxed(), ) .sendable() } -/// Spawn a child layout writer task onto the session handle. -/// -/// Captures the strategy, ctx, sink, and a cloned session so the spawned future is `'static`. -fn spawn_layout_write( - handle: &vortex_io::runtime::Handle, - strategy: Arc, - ctx: ArrayContext, - sink: SegmentSinkRef, - session: &VortexSession, - stream: SendableSequentialStream, - eof: SequencePointer, -) -> vortex_io::runtime::Task> { - let session = session.clone(); - handle.spawn_nested(move |h| async move { - let session = session.with_handle(h); - strategy - .write_stream(ctx, sink, stream, eof, &session) - .await - }) -} - #[cfg(test)] mod tests { use std::sync::Arc; From c46f05b6df6b314616bf09c9eab0b2808fbcf853 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 14:04:25 -0400 Subject: [PATCH 09/37] fix writer Signed-off-by: Matt Katz --- vortex-layout/Cargo.toml | 2 + vortex-layout/src/layouts/list/writer.rs | 315 ++++++++--------------- 2 files changed, 104 insertions(+), 213 deletions(-) diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 61b1253ef43..313d4580911 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -25,6 +25,7 @@ async-trait = { workspace = true } bit-vec = { workspace = true } flatbuffers = { workspace = true } futures = { workspace = true, features = ["alloc", "async-await", "executor"] } +insta.workspace = true itertools = { workspace = true } kanal = { workspace = true } moka = { workspace = true, features = ["future"] } @@ -55,6 +56,7 @@ vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] futures = { workspace = true, features = ["executor"] } +insta = { workspace = true } rstest = { workspace = true } temp-env = { workspace = true } tokio = { workspace = true, features = ["rt", "macros"] } diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 0657f94d624..608beebbc2f 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -170,241 +170,129 @@ fn single_chunk_stream( #[cfg(test)] mod tests { - use std::sync::Arc; - - use vortex_array::ArrayContext; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::MaskFuture; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; - use vortex_array::assert_arrays_eq; - use vortex_array::expr::root; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::buffer; - use crate::LayoutStrategy; + use super::*; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; - use crate::layouts::list::writer::ListLayoutStrategy; use crate::segments::TestSegments; - use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; - async fn round_trip(list: ArrayRef) { - let segments = Arc::new(TestSegments::default()); + fn flat_list_strategy() -> ListLayoutStrategy { let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let writer = - ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); - - let (ptr, eof) = SequenceId::root().split(); - let stream = list.clone().to_array_stream().sequenced(ptr); - - let layout = writer - .write_stream( - ArrayContext::empty(), - Arc::::clone(&segments), - stream, - eof, - &SESSION, - ) - .await - .unwrap(); - - let reader = layout - .new_reader(Arc::from("test"), segments, &SESSION) - .unwrap(); - - let row_count = usize::try_from(layout.row_count()).unwrap(); - let result = reader - .projection_evaluation( - &(0..layout.row_count()), - &root(), - MaskFuture::new_true(row_count), - ) - .unwrap() - .await - .unwrap(); - - assert_arrays_eq!(result, list); - } - - #[tokio::test] - async fn round_trip_non_nullable() { - let elements = buffer![1i32, 2, 3, 4, 5].into_array(); - let offsets = buffer![0u32, 2, 5, 5].into_array(); // 3 lists: [1,2], [3,4,5], [] - let list = ListArray::try_new(elements, offsets, Validity::NonNullable) - .unwrap() - .into_array(); - round_trip(list).await; - } - - #[tokio::test] - async fn round_trip_nullable() { - let elements = buffer![10i32, 20, 30, 40, 50].into_array(); - let offsets = buffer![0u32, 2, 3, 5].into_array(); // 3 lists - let validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); - let list = ListArray::try_new(elements, offsets, validity) - .unwrap() - .into_array(); - round_trip(list).await; + ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)) } - /// Writes a list, then reads back only a sub-range to exercise projection over a slice. - async fn round_trip_subset(list: ArrayRef, row_range: std::ops::Range) { + async fn write(strategy: &S, array: ArrayRef) -> VortexResult { let segments = Arc::new(TestSegments::default()); - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let writer = - ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); - let (ptr, eof) = SequenceId::root().split(); - let stream = list.clone().to_array_stream().sequenced(ptr); - - let layout = writer - .write_stream( - ArrayContext::empty(), - Arc::::clone(&segments), - stream, - eof, - &SESSION, - ) + let stream = array.to_array_stream().sequenced(ptr); + strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) .await - .unwrap(); - - let reader = layout - .new_reader(Arc::from("test"), segments, &SESSION) - .unwrap(); + } - let mask_len = usize::try_from(row_range.end - row_range.start).unwrap(); - let result = reader - .projection_evaluation(&row_range, &root(), MaskFuture::new_true(mask_len)) - .unwrap() - .await - .unwrap(); + fn i32_list_dtype(nullable: bool) -> DType { + DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }, + ) + } - let expected = list - .slice( - usize::try_from(row_range.start).unwrap()..usize::try_from(row_range.end).unwrap(), - ) - .unwrap(); - assert_arrays_eq!(result, expected); + fn create_basic_list(validity: Validity) -> ArrayRef { + ListArray::try_new( + buffer![1i32, 2, 3, 4, 5].into_array(), + buffer![0u32, 2, 5, 5].into_array(), + validity, + ) + .unwrap() + .into_array() } #[tokio::test] - async fn round_trip_subset_non_nullable() { - // 5 lists: [1,2], [3], [], [4,5,6], [7] - let elements = buffer![1i32, 2, 3, 4, 5, 6, 7].into_array(); - let offsets = buffer![0u32, 2, 3, 3, 6, 7].into_array(); - let list = ListArray::try_new(elements, offsets, Validity::NonNullable) - .unwrap() - .into_array(); - // Read the middle three lists: [3], [], [4,5,6] - round_trip_subset(list, 1..4).await; + async fn basic_non_nullable_input() -> VortexResult<()> { + let list = create_basic_list(Validity::NonNullable); + + let layout = write(&flat_list_strategy(), list).await?; + assert_eq!(layout.row_count(), 3); + + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.flat, dtype: i32, segment: 0 + └── offsets: vortex.flat, dtype: u32, segment: 1 + "); + Ok(()) } #[tokio::test] - async fn round_trip_subset_nullable() { - // 4 lists with validity [true, false, true, true]: - // [10,20], null, [30], [40,50,60] - let elements = buffer![10i32, 20, 30, 40, 50, 60].into_array(); - let offsets = buffer![0u32, 2, 2, 3, 6].into_array(); - let validity = - Validity::Array(BoolArray::from_iter([true, false, true, true]).into_array()); - let list = ListArray::try_new(elements, offsets, validity) - .unwrap() - .into_array(); - // Read lists 1..3: null, [30] - round_trip_subset(list, 1..3).await; + async fn basic_nullable_input() -> VortexResult<()> { + let list = create_basic_list(Validity::Array( + BoolArray::from_iter([true, false, true]).into_array(), + )); + + let layout = write(&flat_list_strategy(), list).await?; + assert_eq!(layout.row_count(), 3); + + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(i32)?, children: 3 + ├── elements: vortex.flat, dtype: i32, segment: 0 + ├── offsets: vortex.flat, dtype: u32, segment: 1 + └── validity: vortex.flat, dtype: bool, segment: 2 + "); + Ok(()) } - // -- tree shape visualization --------------------------------------------------------- - // - // These tests are mostly for development/inspection — they show what the resulting - // layout tree looks like for various input shapes. Run with `--nocapture` to see the - // pretty-printed trees: - // - // cargo test -p vortex-layout layouts::list::writer::tests::tree -- --nocapture - - use vortex_array::ArrayContext as _ArrayContextAlias; - - /// Write `array` directly through `ListLayoutStrategy` (no ChunkedLayoutStrategy wrap) - /// and return the resulting top-level layout. - async fn write_through_list_strategy(array: ArrayRef) -> crate::LayoutRef { + #[tokio::test] + async fn empty_stream_errors() { let segments = Arc::new(TestSegments::default()); - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let writer = - ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); - let (ptr, eof) = SequenceId::root().split(); - let stream = array.to_array_stream().sequenced(ptr); - writer - .write_stream(_ArrayContextAlias::empty(), segments, stream, eof, &SESSION) - .await - .unwrap() - } + let (_, eof) = SequenceId::root().split(); + let empty = stream::empty::>().boxed(); + let stream = SequentialStreamAdapter::new(i32_list_dtype(false), empty).sendable(); - /// Wrap `ListLayoutStrategy` in `ChunkedLayoutStrategy` and write `array`. For a chunked - /// input, each chunk becomes one `ListLayout` under the outer `ChunkedLayout`. - async fn write_through_chunked_list_strategy(array: ArrayRef) -> crate::LayoutRef { - use crate::layouts::chunked::writer::ChunkedLayoutStrategy; - let segments = Arc::new(TestSegments::default()); - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let list_strategy = - ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)); - let writer = ChunkedLayoutStrategy::new(list_strategy); - let (ptr, eof) = SequenceId::root().split(); - let stream = array.to_array_stream().sequenced(ptr); - writer - .write_stream(_ArrayContextAlias::empty(), segments, stream, eof, &SESSION) + let err = flat_list_strategy() + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) .await - .unwrap() + .unwrap_err(); + insta::assert_snapshot!(err.to_string(), @"Other error: ListLayoutStrategy needs a single chunk"); } #[tokio::test] - async fn tree_shape_single_chunk_non_nullable() { - let elements = buffer![1i32, 2, 3, 4, 5].into_array(); - let offsets = buffer![0u32, 2, 5, 5].into_array(); - let list = ListArray::try_new(elements, offsets, Validity::NonNullable) - .unwrap() - .into_array(); - - let layout = write_through_list_strategy(list).await; - let tree = layout.display_tree().to_string(); - eprintln!("--- single-chunk non-nullable ---\n{tree}"); - // Top level is a single ListLayout with 2 children (elements, offsets). - assert!(tree.starts_with("vortex.list")); - assert!(tree.contains("elements")); - assert!(tree.contains("offsets")); - assert!(!tree.contains("validity")); - } + async fn chunked_list_input_without_chunked_strategy_fails() -> VortexResult<()> { + let chunk0 = ListArray::try_new( + buffer![1i32, 2].into_array(), + buffer![0u32, 2].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let chunk1 = ListArray::try_new( + buffer![3i32, 4, 5].into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let chunked = + ChunkedArray::try_new(vec![chunk0, chunk1], i32_list_dtype(false))?.into_array(); - #[tokio::test] - async fn tree_shape_single_chunk_nullable() { - let elements = buffer![10i32, 20, 30, 40, 50].into_array(); - let offsets = buffer![0u32, 2, 3, 5].into_array(); - let validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); - let list = ListArray::try_new(elements, offsets, validity) - .unwrap() - .into_array(); - - let layout = write_through_list_strategy(list).await; - let tree = layout.display_tree().to_string(); - eprintln!("--- single-chunk nullable ---\n{tree}"); - // Top level is a single ListLayout with 3 children (elements, offsets, validity). - assert!(tree.starts_with("vortex.list")); - assert!(tree.contains("elements")); - assert!(tree.contains("offsets")); - assert!(tree.contains("validity")); + let err = write(&flat_list_strategy(), chunked).await.unwrap_err(); + insta::assert_snapshot!(err.to_string(), @"Other error: ListLayoutStrategy received more than a single chunk"); + Ok(()) } #[tokio::test] - async fn tree_shape_multi_chunk_via_chunked_strategy() { - use std::sync::Arc as StdArc; - use vortex_array::arrays::ChunkedArray; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - - // Two list-array chunks fed through ChunkedArray -> ChunkedLayoutStrategy. + async fn chunked_list_input_with_chunked_strategy_succeeds() -> VortexResult<()> { let chunk0 = ListArray::try_new( buffer![1i32, 2, 3].into_array(), buffer![0u32, 2, 3].into_array(), @@ -420,19 +308,20 @@ mod tests { .unwrap() .into_array(); - let list_dtype = DType::List( - StdArc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), - Nullability::NonNullable, - ); - let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], list_dtype) - .unwrap() - .into_array(); - - let layout = write_through_chunked_list_strategy(chunked).await; - let tree = layout.display_tree().to_string(); - eprintln!("--- multi-chunk via ChunkedLayoutStrategy ---\n{tree}"); - // Top level is a ChunkedLayout containing two ListLayouts. - assert!(tree.starts_with("vortex.chunked")); - assert_eq!(tree.matches("vortex.list").count(), 2); + let chunked = + ChunkedArray::try_new(vec![chunk0, chunk1], i32_list_dtype(false))?.into_array(); + + let layout = write(&ChunkedLayoutStrategy::new(flat_list_strategy()), chunked).await?; + + insta::assert_snapshot!(layout.display_tree(), @" + vortex.chunked, dtype: list(i32), children: 2 + ├── [0]: vortex.list, dtype: list(i32), children: 2 + │ ├── elements: vortex.flat, dtype: i32, segment: 0 + │ └── offsets: vortex.flat, dtype: u32, segment: 1 + └── [1]: vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.flat, dtype: i32, segment: 2 + └── offsets: vortex.flat, dtype: u32, segment: 3 + "); + Ok(()) } } From 7682427aabd06d56ebc0cd17e5876ef38913adb2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 14:05:37 -0400 Subject: [PATCH 10/37] fix Signed-off-by: Matt Katz --- vortex-layout/Cargo.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 313d4580911..daa96d4e138 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -25,7 +25,6 @@ async-trait = { workspace = true } bit-vec = { workspace = true } flatbuffers = { workspace = true } futures = { workspace = true, features = ["alloc", "async-await", "executor"] } -insta.workspace = true itertools = { workspace = true } kanal = { workspace = true } moka = { workspace = true, features = ["future"] } From d1b46d9c7c724b4294361f2321d6d9bb38e05660 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 15:39:36 -0400 Subject: [PATCH 11/37] improve projection eval Signed-off-by: Matt Katz --- Cargo.lock | 1 + vortex-layout/src/layouts/list/reader.rs | 254 +++++++++++++++++++---- 2 files changed, 218 insertions(+), 37 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c7c50f5047e..ec2fa196e98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10185,6 +10185,7 @@ dependencies = [ "bit-vec", "flatbuffers", "futures", + "insta", "itertools 0.14.0", "kanal", "moka", diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index ff8fddddd0a..d478d020b7d 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -7,17 +7,22 @@ use std::sync::Arc; use futures::FutureExt; use futures::try_join; -use vortex_array::Array; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::MaskFuture; -use vortex_array::arrays::List; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListArray; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; +use vortex_array::expr::is_root; use vortex_array::expr::root; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -30,14 +35,10 @@ use crate::layouts::list::ListLayout; use crate::segments::SegmentSource; /// Reader for [`ListLayout`]. -/// -/// Reads the underlying `elements`, `offsets`, and (optional) `validity` children in parallel -/// and reassembles them into a [`ListArray`](vortex_array::arrays::ListArray) before applying -/// the requested projection. pub struct ListReader { layout: ListLayout, name: Arc, - _session: VortexSession, + session: VortexSession, elements: LayoutReaderRef, offsets: LayoutReaderRef, validity: Option, @@ -74,12 +75,54 @@ impl ListReader { Ok(Self { layout, name, - _session: session, + session, elements, offsets, validity, }) } + + fn fetch_offsets(&self, row_range: &Range) -> VortexResult { + // The offsets child has an extra entry, so reading `row_range` maps to offsets in + // `[row_range.start..row_range.end + 1)`. + let offsets_range = row_range.start..(row_range.end + 1); + let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; + + self.offsets.projection_evaluation( + &offsets_range, + &root(), + MaskFuture::new_true(offsets_count), + ) + } +} + +/// Read `offsets[0]` and `offsets[-1]` and return the elements-buffer range they describe. +fn calculate_elements_range( + offsets: &ArrayRef, + session: &VortexSession, +) -> VortexResult> { + if offsets.is_empty() { + return Ok(0..0); + } + let mut exec_ctx = session.create_execution_ctx(); + let start = offsets + .execute_scalar(0, &mut exec_ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value fits in u64"); + let end = offsets + .execute_scalar(offsets.len() - 1, &mut exec_ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value fits in u64"); + Ok(start..end) +} + +/// Subtract `first` from every offset so the resulting offsets index into a sliced +/// `elements[first..]` buffer starting at zero. +fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { + let constant = ConstantArray::new(first, offsets.len()).into_array(); + offsets.binary(constant, Operator::Sub) } impl LayoutReader for ListReader { @@ -105,10 +148,12 @@ impl LayoutReader for ListReader { split_range: &SplitRange, splits: &mut BTreeSet, ) -> VortexResult<()> { - // Offsets has one more row than the list itself but shares the list's chunking - // structure, so it's the appropriate child to drive scan splits. self.offsets - .register_splits(field_mask, split_range, splits) + .register_splits(field_mask, split_range, splits)?; + if let Some(validity) = &self.validity { + validity.register_splits(field_mask, split_range, splits)?; + } + Ok(()) } fn pruning_evaluation( @@ -135,47 +180,55 @@ impl LayoutReader for ListReader { expr: &Expression, mask: MaskFuture, ) -> VortexResult { - let elements_len = self.layout.elements().row_count(); - let elements_range = 0..elements_len; - let elements_mask = MaskFuture::new_true(usize::try_from(elements_len)?); - - let row_count: usize = usize::try_from(row_range.end - row_range.start)?; - // The offsets child has n+1 entries; reading row_range maps to offsets in - // [row_range.start..row_range.end + 1). - let offsets_range = row_range.start..row_range.end + 1; - let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; - - let elements_fut = - self.elements - .projection_evaluation(&elements_range, &root(), elements_mask)?; - let offsets_fut = self.offsets.projection_evaluation( - &offsets_range, - &root(), - MaskFuture::new_true(offsets_count), - )?; + let offsets_fut = self.fetch_offsets(row_range)?; + // Validity shares the list's row space, so the caller's mask is the right shape to + // push down. Children get to fold it into their reads however they like. let validity_fut = self .validity .as_ref() - .map(|v| v.projection_evaluation(row_range, &root(), MaskFuture::new_true(row_count))) + .map(|v| v.projection_evaluation(row_range, &root(), mask.clone())) .transpose()?; + let elements_reader = self.elements.clone(); + let session = self.session.clone(); let nullability = self.layout.dtype().nullability(); let expr = expr.clone(); Ok(async move { - let (elements, offsets) = try_join!(elements_fut, offsets_fut)?; - let validity = match validity_fut { - Some(fut) => Validity::Array(fut.await?), + // Fetch offsets and validity in parallel. Elements waits until we know + // exactly which slice of the elements buffer it actually needs. + let (offsets, validity_array) = try_join!(offsets_fut, async move { + match validity_fut { + Some(fut) => fut.await.map(Some), + None => Ok(None), + } + },)?; + + // Bound the elements read using offsets[0] and offsets[-1] + let elements_range = calculate_elements_range(&offsets, &session)?; + + // Rebase the offsets so they start at zero + let rebased_offsets = rebase_offsets(offsets, elements_range.start)?; + + // Fetch only the elements we actually need. + let elements_len = elements_range.end - elements_range.start; + let elements = elements_reader + .projection_evaluation( + &elements_range, + &root(), + MaskFuture::new_true(usize::try_from(elements_len)?), + )? + .await?; + + let validity = match validity_array { + Some(arr) => Validity::Array(arr), None => match nullability { Nullability::Nullable => Validity::AllValid, Nullability::NonNullable => Validity::NonNullable, }, }; - // SAFETY: the layout was validated at write time, so offsets are monotonic, - // non-nullable, integer, of length n+1. - let array: ArrayRef = - unsafe { Array::::new_unchecked(elements, offsets, validity) }.into_array(); + let array = ListArray::try_new(elements, rebased_offsets, validity)?.into_array(); let mask = mask.await?; let array = if !mask.all_true() { @@ -189,3 +242,130 @@ impl LayoutReader for ListReader { .boxed()) } } + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + use std::ops::Range; + + use vortex_array::ArrayContext; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::ListArray; + use vortex_array::dtype::FieldMask; + use vortex_buffer::buffer; + + use super::*; + use crate::LayoutRef; + use crate::LayoutStrategy; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::list::writer::ListLayoutStrategy; + use crate::scan::split_by::SplitBy; + use crate::segments::SegmentSource; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialArrayStreamExt; + use crate::test::SESSION; + + fn flat_list_strategy() -> ListLayoutStrategy { + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)) + } + + async fn write_layout( + strategy: &S, + array: ArrayRef, + ) -> VortexResult<(Arc, LayoutRef)> { + let segments = Arc::new(TestSegments::default()); + let segments_ref: Arc = Arc::::clone(&segments); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + let layout = strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .await?; + Ok((segments_ref, layout)) + } + + fn collect_splits(reader: &dyn LayoutReader, row_range: Range) -> BTreeSet { + SplitBy::Layout + .splits(reader, &row_range, &[FieldMask::All]) + .unwrap() + } + + #[tokio::test] + async fn splits_non_nullable_flat() -> VortexResult<()> { + // 5 lists in one flat ListLayout. Offsets is flat (single segment), no validity. + // Only the row_range endpoints should appear in the split set. + let list = ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6, 7].into_array(), + buffer![0u32, 2, 3, 3, 6, 7].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &SESSION)?; + + assert_eq!( + collect_splits(reader.as_ref(), 0..5), + BTreeSet::from([0, 5]) + ); + assert_eq!( + collect_splits(reader.as_ref(), 1..4), + BTreeSet::from([1, 4]) + ); + Ok(()) + } + + #[tokio::test] + async fn splits_nullable_flat_dedups_validity_and_offsets() -> VortexResult<()> { + // Nullable list with flat validity. Offsets and validity both end at list-row 3; + // BTreeSet deduplicates, so the split set still has only the endpoints. + let list = ListArray::try_new( + buffer![10i32, 20, 30, 40, 50].into_array(), + buffer![0u32, 2, 3, 5].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + )? + .into_array(); + + let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &SESSION)?; + + assert_eq!( + collect_splits(reader.as_ref(), 0..3), + BTreeSet::from([0, 3]) + ); + Ok(()) + } + + #[tokio::test] + async fn splits_chunked_wrap_exposes_chunk_boundaries() -> VortexResult<()> { + // Two list chunks under a ChunkedLayoutStrategy. The chunk boundary at list-row 2 + // should appear as a split. + let chunk0 = ListArray::try_new( + buffer![1i32, 2, 3].into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let chunk1 = ListArray::try_new( + buffer![4i32, 5, 6, 7].into_array(), + buffer![0u32, 1, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let dtype = chunk0.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array(); + + let strategy = ChunkedLayoutStrategy::new(flat_list_strategy()); + let (segments, layout) = write_layout(&strategy, chunked).await?; + let reader = layout.new_reader("".into(), segments, &SESSION)?; + + assert_eq!( + collect_splits(reader.as_ref(), 0..4), + BTreeSet::from([0, 2, 4]) + ); + Ok(()) + } +} From b58bc03b068636e32477707df8aaf3ebe7ceaa06 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 15:46:32 -0400 Subject: [PATCH 12/37] projection evaluation Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index d478d020b7d..984582816c3 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -18,7 +18,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; -use vortex_array::expr::is_root; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; @@ -125,6 +124,16 @@ fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { offsets.binary(constant, Operator::Sub) } +fn create_validity(validity_array: Option, nullability: Nullability) -> Validity { + match validity_array { + Some(arr) => Validity::Array(arr), + None => match nullability { + Nullability::Nullable => Validity::AllValid, + Nullability::NonNullable => Validity::NonNullable, + }, + } +} + impl LayoutReader for ListReader { fn name(&self) -> &Arc { &self.name @@ -220,16 +229,11 @@ impl LayoutReader for ListReader { )? .await?; - let validity = match validity_array { - Some(arr) => Validity::Array(arr), - None => match nullability { - Nullability::Nullable => Validity::AllValid, - Nullability::NonNullable => Validity::NonNullable, - }, - }; - + // Create ListArray + let validity = create_validity(validity_array, nullability); let array = ListArray::try_new(elements, rebased_offsets, validity)?.into_array(); + // Apply mask and expression let mask = mask.await?; let array = if !mask.all_true() { array.filter(mask)? From c2496a41ad430ed41d5e282833c0fb61fb621b87 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 16:21:07 -0400 Subject: [PATCH 13/37] tests Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 163 ++++++++++++++--------- 1 file changed, 99 insertions(+), 64 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 984582816c3..16ea17900dc 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -118,9 +118,11 @@ fn calculate_elements_range( } /// Subtract `first` from every offset so the resulting offsets index into a sliced -/// `elements[first..]` buffer starting at zero. +/// `elements[first..]` buffer starting at zero. The constant array is cast to the offsets' dtype. fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { - let constant = ConstantArray::new(first, offsets.len()).into_array(); + let constant = ConstantArray::new(first, offsets.len()) + .into_array() + .cast(offsets.dtype().clone())?; offsets.binary(constant, Operator::Sub) } @@ -249,23 +251,22 @@ impl LayoutReader for ListReader { #[cfg(test)] mod tests { - use std::collections::BTreeSet; use std::ops::Range; + use rstest::rstest; use vortex_array::ArrayContext; use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; - use vortex_array::dtype::FieldMask; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::Nullability::NonNullable; use vortex_buffer::buffer; use super::*; use crate::LayoutRef; use crate::LayoutStrategy; - use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::writer::ListLayoutStrategy; - use crate::scan::split_by::SplitBy; use crate::segments::SegmentSource; use crate::segments::TestSegments; use crate::sequence::SequenceId; @@ -291,85 +292,119 @@ mod tests { Ok((segments_ref, layout)) } - fn collect_splits(reader: &dyn LayoutReader, row_range: Range) -> BTreeSet { - SplitBy::Layout - .splits(reader, &row_range, &[FieldMask::All]) + fn materialize_u32_array(array: ArrayRef) -> Vec { + let mut ctx = SESSION.create_execution_ctx(); + array + .execute::(&mut ctx) .unwrap() + .as_slice::() + .to_vec() } - #[tokio::test] - async fn splits_non_nullable_flat() -> VortexResult<()> { - // 5 lists in one flat ListLayout. Offsets is flat (single segment), no validity. - // Only the row_range endpoints should appear in the split set. - let list = ListArray::try_new( - buffer![1i32, 2, 3, 4, 5, 6, 7].into_array(), - buffer![0u32, 2, 3, 3, 6, 7].into_array(), - Validity::NonNullable, - )? - .into_array(); + #[rstest] + #[case::full(buffer![0u32, 2, 5, 5].into_array(), 0..5)] + #[case::partial_slice(buffer![2u32, 5, 5, 8].into_array(), 2..8)] + #[case::single_offset_is_empty(buffer![7u32].into_array(), 7..7)] + #[case::u64_offsets(buffer![10u64, 12, 15, 15].into_array(), 10..15)] + fn test_calculate_elements_range( + #[case] offsets: ArrayRef, + #[case] expected: Range, + ) -> VortexResult<()> { + assert_eq!(calculate_elements_range(&offsets, &SESSION)?, expected); + Ok(()) + } - let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; - let reader = layout.new_reader("".into(), segments, &SESSION)?; + #[test] + fn calculate_elements_range_empty_offsets() -> VortexResult<()> { + let offsets = PrimitiveArray::empty::(NonNullable).into_array(); + assert_eq!(calculate_elements_range(&offsets, &SESSION)?, 0..0); + Ok(()) + } - assert_eq!( - collect_splits(reader.as_ref(), 0..5), - BTreeSet::from([0, 5]) - ); - assert_eq!( - collect_splits(reader.as_ref(), 1..4), - BTreeSet::from([1, 4]) - ); + #[rstest] + #[case::first_zero_is_identity(buffer![0u32, 2, 5, 5].into_array(), 0, vec![0, 2, 5, 5])] + #[case::subtracts_first(buffer![3u32, 5, 8].into_array(), 3, vec![0, 2, 5])] + fn test_rebase_offsets( + #[case] offsets: ArrayRef, + #[case] first: u64, + #[case] expected: Vec, + ) -> VortexResult<()> { + let rebased = rebase_offsets(offsets, first)?; + assert_eq!(materialize_u32_array(rebased), expected); Ok(()) } #[tokio::test] - async fn splits_nullable_flat_dedups_validity_and_offsets() -> VortexResult<()> { - // Nullable list with flat validity. Offsets and validity both end at list-row 3; - // BTreeSet deduplicates, so the split set still has only the endpoints. + async fn fetch_offsets_includes_extra_endpoint() -> VortexResult<()> { let list = ListArray::try_new( - buffer![10i32, 20, 30, 40, 50].into_array(), - buffer![0u32, 2, 3, 5].into_array(), - Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + buffer![1i32, 2, 3, 4, 5].into_array(), + buffer![0u32, 2, 4, 5].into_array(), + Validity::NonNullable, )? .into_array(); let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; let reader = layout.new_reader("".into(), segments, &SESSION)?; - - assert_eq!( - collect_splits(reader.as_ref(), 0..3), - BTreeSet::from([0, 3]) - ); + let reader = reader + .as_any() + .downcast_ref::() + .expect("ListReader"); + + // row_range 1..3 should pull 3 offsets (indices 1, 2, 3) — the +1 endpoint matters. + let offsets = reader.fetch_offsets(&(1..3))?.await?; + assert_eq!(materialize_u32_array(offsets), vec![2, 4, 5]); + + // row_range 0..3 pulls all 4 offsets. + let offsets = reader.fetch_offsets(&(0..3))?.await?; + assert_eq!(materialize_u32_array(offsets), vec![0, 2, 4, 5]); Ok(()) } + fn create_basic_list_array(nullable: bool) -> ArrayRef { + let validity = if nullable { + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()) + } else { + Validity::NonNullable + }; + + ListArray::try_new( + buffer![1i32, 2, 3, 4, 5].into_array(), + buffer![0u32, 2, 4, 5].into_array(), + validity, + ) + .expect("array is valid") + .into_array() + } + + #[rstest] + #[case::full_range(0..3, false)] + #[case::partial_start(0..2, false)] + #[case::partial_end(1..3, false)] + #[case::middle_single(1..2, false)] + #[case::empty_range(1..1, false)] + #[case::full_range_null(0..3, true)] + #[case::partial_start_null(0..2, true)] + #[case::partial_end_null(1..3, true)] + #[case::middle_single_null(1..2, true)] + #[case::empty_range_null(1..1, true)] #[tokio::test] - async fn splits_chunked_wrap_exposes_chunk_boundaries() -> VortexResult<()> { - // Two list chunks under a ChunkedLayoutStrategy. The chunk boundary at list-row 2 - // should appear as a split. - let chunk0 = ListArray::try_new( - buffer![1i32, 2, 3].into_array(), - buffer![0u32, 2, 3].into_array(), - Validity::NonNullable, - )? - .into_array(); - let chunk1 = ListArray::try_new( - buffer![4i32, 5, 6, 7].into_array(), - buffer![0u32, 1, 4].into_array(), - Validity::NonNullable, - )? - .into_array(); - let dtype = chunk0.dtype().clone(); - let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array(); + async fn projection_evaluation_round_trips( + #[case] row_range: Range, + #[case] nullable: bool, + ) -> VortexResult<()> { + let list = create_basic_list_array(nullable); - let strategy = ChunkedLayoutStrategy::new(flat_list_strategy()); - let (segments, layout) = write_layout(&strategy, chunked).await?; + let len = usize::try_from(row_range.end - row_range.start)?; + let (segments, layout) = write_layout(&flat_list_strategy(), list.clone()).await?; let reader = layout.new_reader("".into(), segments, &SESSION)?; - assert_eq!( - collect_splits(reader.as_ref(), 0..4), - BTreeSet::from([0, 2, 4]) - ); + let result = reader + .projection_evaluation(&row_range, &root(), MaskFuture::new_true(len))? + .await?; + + let expected = + list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)?; + assert_arrays_eq!(result, expected); Ok(()) } } From 055301458d3e7d130d2a960e44f06b18641a237f Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 16:22:04 -0400 Subject: [PATCH 14/37] tests Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 16ea17900dc..23fd7bfe2fb 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -383,10 +383,6 @@ mod tests { #[case::middle_single(1..2, false)] #[case::empty_range(1..1, false)] #[case::full_range_null(0..3, true)] - #[case::partial_start_null(0..2, true)] - #[case::partial_end_null(1..3, true)] - #[case::middle_single_null(1..2, true)] - #[case::empty_range_null(1..1, true)] #[tokio::test] async fn projection_evaluation_round_trips( #[case] row_range: Range, From e324910a81f273bd05e0d0a6c46388393472eb93 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 16:22:32 -0400 Subject: [PATCH 15/37] fix test Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 59 +++++++++++++----------- 1 file changed, 33 insertions(+), 26 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 23fd7bfe2fb..87a78796333 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -334,32 +334,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn fetch_offsets_includes_extra_endpoint() -> VortexResult<()> { - let list = ListArray::try_new( - buffer![1i32, 2, 3, 4, 5].into_array(), - buffer![0u32, 2, 4, 5].into_array(), - Validity::NonNullable, - )? - .into_array(); - - let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; - let reader = layout.new_reader("".into(), segments, &SESSION)?; - let reader = reader - .as_any() - .downcast_ref::() - .expect("ListReader"); - - // row_range 1..3 should pull 3 offsets (indices 1, 2, 3) — the +1 endpoint matters. - let offsets = reader.fetch_offsets(&(1..3))?.await?; - assert_eq!(materialize_u32_array(offsets), vec![2, 4, 5]); - - // row_range 0..3 pulls all 4 offsets. - let offsets = reader.fetch_offsets(&(0..3))?.await?; - assert_eq!(materialize_u32_array(offsets), vec![0, 2, 4, 5]); - Ok(()) - } - fn create_basic_list_array(nullable: bool) -> ArrayRef { let validity = if nullable { Validity::Array(BoolArray::from_iter([true, false, true]).into_array()) @@ -376,6 +350,23 @@ mod tests { .into_array() } + #[tokio::test] + async fn fetch_offsets_includes_extra_endpoint() -> VortexResult<()> { + let list = create_basic_list_array(false); + + let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &SESSION)?; + let reader = reader + .as_any() + .downcast_ref::() + .expect("ListReader"); + + let offsets = reader.fetch_offsets(&(1..3))?.await?; + assert_eq!(materialize_u32_array(offsets), vec![2, 4, 5]); + + Ok(()) + } + #[rstest] #[case::full_range(0..3, false)] #[case::partial_start(0..2, false)] @@ -403,4 +394,20 @@ mod tests { assert_arrays_eq!(result, expected); Ok(()) } + + #[tokio::test] + async fn projection_evaluation_applies_mask() -> VortexResult<()> { + let list = create_basic_list_array(false); + let (segments, layout) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &SESSION)?; + + let mask = Mask::from_iter([true, false, true]); + let result = reader + .projection_evaluation(&(0..3), &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let expected = list.filter(mask)?; + assert_arrays_eq!(result, expected); + Ok(()) + } } From f84cce953720f92382cc33b26fdf209854748f84 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 16:40:57 -0400 Subject: [PATCH 16/37] skip pruning eval Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 87a78796333..dcedafb5282 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -171,9 +171,10 @@ impl LayoutReader for ListReader { &self, _row_range: &Range, _expr: &Expression, - _mask: Mask, + mask: Mask, ) -> VortexResult { - todo!() + // All stats based pruning should already be done upstream + Ok(MaskFuture::ready(mask)) } fn filter_evaluation( From ae89cacfe01872b2b5a3d272a6f260214d9ab901 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 27 May 2026 17:40:58 -0400 Subject: [PATCH 17/37] few more tests Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 69 ++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 608beebbc2f..47dd9b46ccd 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -173,6 +173,7 @@ mod tests { use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; + use vortex_array::arrays::StructArray; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::validity::Validity; @@ -181,6 +182,7 @@ mod tests { use super::*; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::table::TableStrategy; use crate::segments::TestSegments; use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; @@ -291,6 +293,73 @@ mod tests { Ok(()) } + #[tokio::test] + async fn list_of_struct_tree() -> VortexResult<()> { + let struct_array = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3, 4, 5].into_array()), + ("b", buffer![10i32, 20, 30, 40, 50].into_array()), + ] + .as_slice(), + )? + .into_array(); + let list = ListArray::try_new( + struct_array, + buffer![0u32, 2, 5, 5].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let table_strategy: Arc = + Arc::new(TableStrategy::new(Arc::clone(&flat), Arc::clone(&flat))); + let writer = ListLayoutStrategy::new(table_strategy, Arc::clone(&flat), Arc::clone(&flat)); + + let layout = write(&writer, list).await?; + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list({a=i32, b=i32}), children: 2 + ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 + │ ├── a: vortex.flat, dtype: i32, segment: 1 + │ └── b: vortex.flat, dtype: i32, segment: 2 + └── offsets: vortex.flat, dtype: u32, segment: 0 + "); + Ok(()) + } + + #[tokio::test] + async fn list_of_list_tree() -> VortexResult<()> { + let inner_list = ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 6].into_array(), + Validity::NonNullable, + )? + .into_array(); + let list = ListArray::try_new( + inner_list, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let inner_strategy: Arc = Arc::new(ListLayoutStrategy::new( + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + )); + let writer = ListLayoutStrategy::new(inner_strategy, Arc::clone(&flat), Arc::clone(&flat)); + + let layout = write(&writer, list).await?; + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(list(i32)), children: 2 + ├── elements: vortex.list, dtype: list(i32), children: 2 + │ ├── elements: vortex.flat, dtype: i32, segment: 1 + │ └── offsets: vortex.flat, dtype: u32, segment: 2 + └── offsets: vortex.flat, dtype: u32, segment: 0 + "); + Ok(()) + } + #[tokio::test] async fn chunked_list_input_with_chunked_strategy_succeeds() -> VortexResult<()> { let chunk0 = ListArray::try_new( From 5dfd384b30190307da44de3d91218cc6710cf793 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 28 May 2026 09:52:16 -0400 Subject: [PATCH 18/37] lint fix Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 23 +++++++++++------------ vortex-layout/src/layouts/list/reader.rs | 2 +- vortex-layout/src/layouts/list/writer.rs | 23 +++++++++++------------ 3 files changed, 23 insertions(+), 25 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index c8671bbb588..72d96361376 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -134,11 +134,10 @@ impl VTable for List { let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); let offsets = children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; - let validity = if dtype.is_nullable() { - Some(children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))?) - } else { - None - }; + let validity = dtype + .is_nullable() + .then(|| children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))) + .transpose()?; Ok(ListLayout { dtype: dtype.clone(), @@ -158,14 +157,14 @@ impl VTable for List { layout.offsets = iter .next() .ok_or_else(|| vortex_err!("missing offsets child"))?; - layout.validity = if layout.dtype.is_nullable() { - Some( + layout.validity = layout + .dtype + .is_nullable() + .then(|| { iter.next() - .ok_or_else(|| vortex_err!("missing validity child"))?, - ) - } else { - None - }; + .ok_or_else(|| vortex_err!("missing validity child")) + }) + .transpose()?; Ok(()) } } diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index dcedafb5282..bd3c8d8088d 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -201,7 +201,7 @@ impl LayoutReader for ListReader { .map(|v| v.projection_evaluation(row_range, &root(), mask.clone())) .transpose()?; - let elements_reader = self.elements.clone(); + let elements_reader = Arc::clone(&self.elements); let session = self.session.clone(); let nullability = self.layout.dtype().nullability(); let expr = expr.clone(); diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 47dd9b46ccd..c9c7ff75f2e 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -98,15 +98,14 @@ impl LayoutStrategy for ListLayoutStrategy { // There is one extra element in `offsets` let row_count = offsets.len().saturating_sub(1); - let validity_array = if dtype.is_nullable() { - Some( + let validity_array = dtype + .is_nullable() + .then(|| { validity - .execute_mask(row_count, &mut exec_ctx)? - .into_array(), - ) - } else { - None - }; + .execute_mask(row_count, &mut exec_ctx) + .map(|m| m.into_array()) + }) + .transpose()?; // Spawn each child write onto the runtime so they run concurrently let handle = session.handle(); @@ -116,7 +115,7 @@ impl LayoutStrategy for ListLayoutStrategy { let stream = single_chunk_stream(array.dtype().clone(), sp.advance(), array); let child_eof = eof.split_off(); let ctx = ctx.clone(); - let segment_sink = segment_sink.clone(); + let segment_sink = Arc::clone(&segment_sink); let session = session.clone(); handle.spawn_nested(move |h| async move { let session = session.with_handle(h); @@ -126,9 +125,9 @@ impl LayoutStrategy for ListLayoutStrategy { }) }; ( - spawn_layout_writer(self.elements.clone(), elements), - spawn_layout_writer(self.offsets.clone(), offsets), - validity_array.map(|arr| spawn_layout_writer(self.validity.clone(), arr)), + spawn_layout_writer(Arc::clone(&self.elements), elements), + spawn_layout_writer(Arc::clone(&self.offsets), offsets), + validity_array.map(|arr| spawn_layout_writer(Arc::clone(&self.validity), arr)), ) }; From f486908dd6e29ebdac8ded4038b97b7d1a3602ac Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 28 May 2026 10:07:14 -0400 Subject: [PATCH 19/37] fix test Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index c9c7ff75f2e..662bf3fb738 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -261,11 +261,10 @@ mod tests { let empty = stream::empty::>().boxed(); let stream = SequentialStreamAdapter::new(i32_list_dtype(false), empty).sendable(); - let err = flat_list_strategy() + let res = flat_list_strategy() .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) - .await - .unwrap_err(); - insta::assert_snapshot!(err.to_string(), @"Other error: ListLayoutStrategy needs a single chunk"); + .await; + assert!(res.is_err()) } #[tokio::test] @@ -287,8 +286,8 @@ mod tests { let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], i32_list_dtype(false))?.into_array(); - let err = write(&flat_list_strategy(), chunked).await.unwrap_err(); - insta::assert_snapshot!(err.to_string(), @"Other error: ListLayoutStrategy received more than a single chunk"); + let res = write(&flat_list_strategy(), chunked).await; + assert!(res.is_err()); Ok(()) } From 0b6a063fd22ea2ae4dcdd8051ba9fb87aeebb9e5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 28 May 2026 10:13:56 -0400 Subject: [PATCH 20/37] quick fix Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 72d96361376..b7627fff230 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -41,8 +41,6 @@ pub const VALIDITY_CHILD_INDEX: usize = 2; /// Number of children when the list dtype is non-nullable. pub const NUM_CHILDREN_NON_NULLABLE: usize = 2; -/// Number of children when the list dtype is nullable. -pub const NUM_CHILDREN_NULLABLE: usize = 3; vtable!(List); @@ -76,11 +74,12 @@ impl VTable for List { } fn nchildren(layout: &Self::Layout) -> usize { + let mut n = NUM_CHILDREN_NON_NULLABLE; if layout.dtype.is_nullable() { - NUM_CHILDREN_NULLABLE - } else { - NUM_CHILDREN_NON_NULLABLE + n += 1; } + + n } fn child(layout: &Self::Layout, idx: usize) -> VortexResult { From cf857b3ab0dc28993ccb2ab938887a4cd8a61877 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 28 May 2026 14:11:57 -0400 Subject: [PATCH 21/37] add anylist matcher Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 31 +++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 662bf3fb738..52e8fc09ae7 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -10,12 +10,15 @@ use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListView; use vortex_array::arrays::list::ListDataParts; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::dtype::DType; +use vortex_array::matcher::Matcher; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_panic; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; @@ -87,14 +90,23 @@ impl LayoutStrategy for ListLayoutStrategy { }; let (sequence_id, array) = chunk?; - // Canonicalize to ListView, then rebuild into zctl + // Run the execution loop until we have a List or ListView; skip work when the input is + // already in one of those forms. If the input is already a ListArray we extract its + // parts directly; if it's a ListViewArray we rebuild via `list_from_list_view`. let mut exec_ctx = session.create_execution_ctx(); + let canonical = array.execute_until::(&mut exec_ctx)?; let ListDataParts { elements, offsets, validity, .. - } = list_from_list_view(array.execute::(&mut exec_ctx)?)?.into_data_parts(); + } = if let Some(list) = canonical.as_opt::() { + list.into_owned().into_data_parts() + } else if let Some(view) = canonical.as_opt::() { + list_from_list_view(view.into_owned())?.into_data_parts() + } else { + unreachable!("AnyList matcher guarantees List or ListView"); + }; // There is one extra element in `offsets` let row_count = offsets.len().saturating_sub(1); @@ -167,6 +179,19 @@ fn single_chunk_stream( .sendable() } +/// Matcher for `Array` or `Array`. Used to short-circuit the execution loop +/// when the input is already in (or directly produces) a list form, avoiding a redundant +/// `ListView` round-trip when the writer already has the parts it needs. +struct AnyList; + +impl Matcher for AnyList { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option> { + (array.as_opt::().is_some() || array.as_opt::().is_some()).then_some(()) + } +} + #[cfg(test)] mod tests { use vortex_array::arrays::BoolArray; From da7794e3e56ef81c05ec2be54323ef270b99b883 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 28 May 2026 15:45:23 -0400 Subject: [PATCH 22/37] read validity with all-true mask, not caller's mask Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 14 +++++++++++--- vortex-layout/src/layouts/list/writer.rs | 1 - 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index bd3c8d8088d..9bd66d4aa16 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -193,12 +193,20 @@ impl LayoutReader for ListReader { mask: MaskFuture, ) -> VortexResult { let offsets_fut = self.fetch_offsets(row_range)?; - // Validity shares the list's row space, so the caller's mask is the right shape to - // push down. Children get to fold it into their reads however they like. + // Read validity at full `row_range` length (all-true mask), not the caller's mask. + // The three children (validity, offsets, elements) must produce compositionally + // consistent outputs that `ListArray::try_new` can assemble. A list-row mask cannot + // be pushed down independently to each child — filtering validity to `popcount(mask)` + // rows would mismatch full-length offsets/elements, breaking the resulting array. The + // final `array.filter(mask)` below applies the mask via list-aware filter semantics + // (re-derives offsets, concatenates kept elements). + let validity_row_count = usize::try_from(row_range.end - row_range.start)?; let validity_fut = self .validity .as_ref() - .map(|v| v.projection_evaluation(row_range, &root(), mask.clone())) + .map(|v| { + v.projection_evaluation(row_range, &root(), MaskFuture::new_true(validity_row_count)) + }) .transpose()?; let elements_reader = Arc::clone(&self.elements); diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 52e8fc09ae7..a2dbee5c546 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -18,7 +18,6 @@ use vortex_array::dtype::DType; use vortex_array::matcher::Matcher; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_panic; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; From e155fb197657f37a8f2b75cea3c4447e9e61a3ce Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 28 May 2026 16:19:09 -0400 Subject: [PATCH 23/37] narrow elements io for sparse mask instead of defering filtering Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 407 ++++++++++++++++++++--- 1 file changed, 358 insertions(+), 49 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 9bd66d4aa16..946386211c9 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -7,12 +7,15 @@ use std::sync::Arc; use futures::FutureExt; use futures::try_join; +use vortex_array::Array; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; @@ -21,6 +24,7 @@ use vortex_array::expr::Expression; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -136,6 +140,119 @@ fn create_validity(validity_array: Option, nullability: Nullability) - } } +/// Plan for fetching only the elements needed to materialize the kept list rows under a sparse +/// row mask, plus the offsets array we'll hand to `ListArray::try_new` for those kept rows. +/// +/// When the row mask is sparse, the alternative (read full row_range, build full list, then +/// `array.filter(mask)`) wastes IO on elements that get thrown away. This plan tells the reader: +/// +/// - which contiguous span of the elements buffer to fetch (`elements_range`), +/// - which positions inside that span belong to a kept row (`element_mask`), +/// - the offsets for the kept-row output, rebased to start at zero (`new_offsets`). +struct ScatterGather { + /// Tightest absolute elements range covering all kept rows. Empty range when no rows kept. + elements_range: Range, + /// `element_mask.len() == elements_range.end - elements_range.start`. A bit is set iff its + /// position in the elements buffer belongs to a kept list row. + element_mask: Mask, + /// Cumulative kept-list lengths starting at zero. `new_offsets.len() == kept_count + 1`. + new_offsets: ArrayRef, + /// Number of true bits in the input row mask. Read by unit tests only. + #[cfg_attr(not(test), allow(dead_code))] + kept_count: usize, +} + +/// Walk the row mask and the (canonicalized) offsets to plan the elements fetch + output offsets +/// for the sparse-mask path of `projection_evaluation`. Single linear pass; no IO. +/// +/// `offsets` is the offsets array we fetched for the full `row_range` (length n+1). `mask` is +/// the row-space mask (length n). Returns a plan suitable for handing the elements child a +/// bounded range + element-level mask, then constructing a kept-only `ListArray`. +// `usize::try_from` / `u64::try_from` are required by the macro arms whose `O` may be `u64` / +// `i64` (potentially fallible on 32-bit targets) but also expand to arms where `O` is `u8`, +// `u16`, etc. (where the conversion is trivially infallible). Suppress the resulting +// `unnecessary_fallible_conversions` lint from the latter arms — the uniform fallible form +// keeps the inner body identical across all expansions. +#[allow(clippy::unnecessary_fallible_conversions)] +fn compute_scatter_gather( + offsets: &ArrayRef, + mask: &Mask, + session: &VortexSession, +) -> VortexResult { + let kept_count = mask.true_count(); + let mut exec_ctx = session.create_execution_ctx(); + let prim_offsets = offsets.clone().execute::(&mut exec_ctx)?; + let ptype = prim_offsets.ptype(); + + if kept_count == 0 { + // Empty result: no elements to fetch, new_offsets is a single zero. + let new_offsets = vortex_array::match_each_integer_ptype!(ptype, |O| { + Array::::new::( + Buffer::::from(vec![O::default()]), + Validity::NonNullable, + ) + .into_array() + }); + return Ok(ScatterGather { + elements_range: 0..0, + element_mask: Mask::new_false(0), + new_offsets, + kept_count: 0, + }); + } + + // Within each macro arm, `O` is a concrete primitive integer type. Offsets are non-negative + // by construction; we materialize spans in `usize` so the element-mask construction is + // straightforward. + vortex_array::match_each_integer_ptype!(ptype, |O| { + let off = prim_offsets.as_slice::(); + + let mut spans: Vec<(usize, usize)> = Vec::with_capacity(kept_count); + let mut new_off: Vec = Vec::with_capacity(kept_count + 1); + new_off.push(O::default()); + let mut cumulative: O = O::default(); + + // `mask.indices()` returns the set bit positions for `Values` masks; `AllTrue` is rare + // here (caller checks density) but we handle it via fallback iteration. + let indices_owned: Vec = match mask.indices() { + vortex_mask::AllOr::All => (0..mask.len()).collect(), + vortex_mask::AllOr::None => Vec::new(), + vortex_mask::AllOr::Some(idxs) => idxs.to_vec(), + }; + for &i in &indices_owned { + let s = off[i]; + let e = off[i + 1]; + spans.push((usize::try_from(s)?, usize::try_from(e)?)); + cumulative += e - s; + new_off.push(cumulative); + } + + let range_start = spans[0].0; + let range_end = spans[spans.len() - 1].1; + + // Element-level mask: same length as the bounded elements range, true bits at positions + // that lie inside any kept span. `Mask::from_slices` rejects zero-width spans, so drop + // empty-list rows here. + let element_slices: Vec<(usize, usize)> = spans + .iter() + .filter(|(s, e)| s < e) + .map(|(s, e)| (s - range_start, e - range_start)) + .collect(); + let element_mask = Mask::from_slices(range_end - range_start, element_slices); + + let new_offsets = + Array::::new::(Buffer::::from(new_off), Validity::NonNullable) + .into_array(); + + Ok(ScatterGather { + elements_range: u64::try_from(range_start)?..u64::try_from(range_end)?, + element_mask, + new_offsets, + kept_count, + }) + }) +} + impl LayoutReader for ListReader { fn name(&self) -> &Arc { &self.name @@ -193,66 +310,86 @@ impl LayoutReader for ListReader { mask: MaskFuture, ) -> VortexResult { let offsets_fut = self.fetch_offsets(row_range)?; - // Read validity at full `row_range` length (all-true mask), not the caller's mask. - // The three children (validity, offsets, elements) must produce compositionally - // consistent outputs that `ListArray::try_new` can assemble. A list-row mask cannot - // be pushed down independently to each child — filtering validity to `popcount(mask)` - // rows would mismatch full-length offsets/elements, breaking the resulting array. The - // final `array.filter(mask)` below applies the mask via list-aware filter semantics - // (re-derives offsets, concatenates kept elements). - let validity_row_count = usize::try_from(row_range.end - row_range.start)?; - let validity_fut = self - .validity - .as_ref() - .map(|v| { - v.projection_evaluation(row_range, &root(), MaskFuture::new_true(validity_row_count)) - }) - .transpose()?; let elements_reader = Arc::clone(&self.elements); + let validity_reader = self.validity.clone(); let session = self.session.clone(); let nullability = self.layout.dtype().nullability(); let expr = expr.clone(); + let row_range = row_range.clone(); Ok(async move { - // Fetch offsets and validity in parallel. Elements waits until we know - // exactly which slice of the elements buffer it actually needs. - let (offsets, validity_array) = try_join!(offsets_fut, async move { - match validity_fut { - Some(fut) => fut.await.map(Some), - None => Ok(None), - } - },)?; - - // Bound the elements read using offsets[0] and offsets[-1] - let elements_range = calculate_elements_range(&offsets, &session)?; - - // Rebase the offsets so they start at zero - let rebased_offsets = rebase_offsets(offsets, elements_range.start)?; - - // Fetch only the elements we actually need. - let elements_len = elements_range.end - elements_range.start; - let elements = elements_reader - .projection_evaluation( + // Resolve offsets and the caller's mask in parallel. We need both before we can + // decide between the two paths below. + let (offsets, mask) = try_join!(offsets_fut, mask)?; + + // Path A: all-true mask — fetch the full bounded elements range, build a ListArray + // covering the entire row_range. The three children must be compositionally + // consistent here, so validity is read at full `row_range` length with an all-true + // mask too. + // + // Path B: sparse mask — scatter-gather. Bound the elements fetch to the tightest + // range covering the kept rows and pass an element-level mask so the elements child + // only materializes positions belonging to a kept row. Validity is fetched at + // `kept_count` length by pushing the caller mask down directly. + if mask.all_true() { + let elements_range = calculate_elements_range(&offsets, &session)?; + let rebased_offsets = rebase_offsets(offsets, elements_range.start)?; + let elements_len = elements_range.end - elements_range.start; + let validity_row_count = usize::try_from(row_range.end - row_range.start)?; + + let validity_fut = validity_reader + .as_ref() + .map(|v| { + v.projection_evaluation( + &row_range, + &root(), + MaskFuture::new_true(validity_row_count), + ) + }) + .transpose()?; + let elements_fut = elements_reader.projection_evaluation( &elements_range, &root(), MaskFuture::new_true(usize::try_from(elements_len)?), - )? - .await?; - - // Create ListArray - let validity = create_validity(validity_array, nullability); - let array = ListArray::try_new(elements, rebased_offsets, validity)?.into_array(); - - // Apply mask and expression - let mask = mask.await?; - let array = if !mask.all_true() { - array.filter(mask)? + )?; + + let (elements, validity_array) = try_join!(elements_fut, async move { + match validity_fut { + Some(fut) => fut.await.map(Some), + None => Ok(None), + } + })?; + + let validity = create_validity(validity_array, nullability); + let array = ListArray::try_new(elements, rebased_offsets, validity)?.into_array(); + array.apply(&expr) } else { - array - }; - - array.apply(&expr) + let sg = compute_scatter_gather(&offsets, &mask, &session)?; + + let validity_fut = validity_reader + .as_ref() + .map(|v| { + v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask)) + }) + .transpose()?; + let elements_fut = elements_reader.projection_evaluation( + &sg.elements_range, + &root(), + MaskFuture::ready(sg.element_mask), + )?; + + let (elements, validity_array) = try_join!(elements_fut, async move { + match validity_fut { + Some(fut) => fut.await.map(Some), + None => Ok(None), + } + })?; + + let validity = create_validity(validity_array, nullability); + let array = ListArray::try_new(elements, sg.new_offsets, validity)?.into_array(); + array.apply(&expr) + } } .boxed()) } @@ -343,6 +480,131 @@ mod tests { Ok(()) } + // ---- compute_scatter_gather -------------------------------------------------------------- + + /// Run `compute_scatter_gather` and unwrap the three derived fields plus the kept count. + /// Returns the raw `new_offsets` ArrayRef so callers with non-u32 offsets can materialize + /// the ptype themselves. + fn run_scatter_gather( + offsets: ArrayRef, + mask: Mask, + ) -> VortexResult<(Range, Vec, ArrayRef, usize)> { + let sg = compute_scatter_gather(&offsets, &mask, &SESSION)?; + let element_mask_bits: Vec = (0..sg.element_mask.len()) + .map(|i| sg.element_mask.value(i)) + .collect(); + Ok(( + sg.elements_range, + element_mask_bits, + sg.new_offsets, + sg.kept_count, + )) + } + + /// Source layout for these tests: 5 lists with offsets `[0, 2, 5, 5, 8, 10]`, i.e. + /// lengths `[2, 3, 0, 3, 2]`. Element positions for list i are `offsets[i]..offsets[i+1]`. + fn five_list_offsets() -> ArrayRef { + buffer![0u32, 2, 5, 5, 8, 10].into_array() + } + + #[test] + fn scatter_gather_single_middle_row() -> VortexResult<()> { + // Keep only list 1 (positions 2..5). + let mask = Mask::from_iter([false, true, false, false, false]); + let (range, elem_mask, new_off, kept) = run_scatter_gather(five_list_offsets(), mask)?; + assert_eq!(range, 2..5); + assert_eq!(elem_mask, vec![true; 3]); // entire bounded range is the kept span + assert_eq!(materialize_u32_array(new_off), vec![0, 3]); + assert_eq!(kept, 1); + Ok(()) + } + + #[test] + fn scatter_gather_two_adjacent_rows() -> VortexResult<()> { + // Keep lists 1 and 2 (positions 2..5 and 5..5 — second is empty). + let mask = Mask::from_iter([false, true, true, false, false]); + let (range, elem_mask, new_off, kept) = run_scatter_gather(five_list_offsets(), mask)?; + assert_eq!(range, 2..5); + assert_eq!(elem_mask, vec![true; 3]); + assert_eq!(materialize_u32_array(new_off), vec![0, 3, 3]); // second kept row has length 0 + assert_eq!(kept, 2); + Ok(()) + } + + #[test] + fn scatter_gather_two_far_apart_rows() -> VortexResult<()> { + // Keep lists 0 and 3 (positions 0..2 and 5..8). Element mask must skip position 2..5. + let mask = Mask::from_iter([true, false, false, true, false]); + let (range, elem_mask, new_off, kept) = run_scatter_gather(five_list_offsets(), mask)?; + assert_eq!(range, 0..8); + // positions 0..2 and 5..8 set, 2..5 unset. + assert_eq!( + elem_mask, + vec![true, true, false, false, false, true, true, true] + ); + assert_eq!(materialize_u32_array(new_off), vec![0, 2, 5]); // lengths 2 and 3 + assert_eq!(kept, 2); + Ok(()) + } + + #[test] + fn scatter_gather_at_boundaries() -> VortexResult<()> { + // Keep first and last list (positions 0..2 and 8..10). + let mask = Mask::from_iter([true, false, false, false, true]); + let (range, elem_mask, new_off, kept) = run_scatter_gather(five_list_offsets(), mask)?; + assert_eq!(range, 0..10); + let mut expected = vec![false; 10]; + expected[0] = true; + expected[1] = true; + expected[8] = true; + expected[9] = true; + assert_eq!(elem_mask, expected); + assert_eq!(materialize_u32_array(new_off), vec![0, 2, 4]); + assert_eq!(kept, 2); + Ok(()) + } + + #[test] + fn scatter_gather_empty_mask_returns_empty_plan() -> VortexResult<()> { + let mask = Mask::new_false(5); + let (range, elem_mask, new_off, kept) = run_scatter_gather(five_list_offsets(), mask)?; + assert_eq!(range, 0..0); + assert!(elem_mask.is_empty()); + // single zero, ready to be a 0-row ListArray's offsets (offsets.len() - 1 == 0 rows) + assert_eq!(materialize_u32_array(new_off), vec![0]); + assert_eq!(kept, 0); + Ok(()) + } + + #[test] + fn scatter_gather_kept_row_is_empty_list() -> VortexResult<()> { + // Keep only list 2, which has length 0 (offsets[2] == offsets[3] == 5). + let mask = Mask::from_iter([false, false, true, false, false]); + let (range, elem_mask, new_off, kept) = run_scatter_gather(five_list_offsets(), mask)?; + assert_eq!(range, 5..5); + assert!(elem_mask.is_empty()); + assert_eq!(materialize_u32_array(new_off), vec![0, 0]); + assert_eq!(kept, 1); + Ok(()) + } + + #[test] + fn scatter_gather_u64_offsets() -> VortexResult<()> { + // Verify the ptype-dispatch path works for u64 offsets, not just u32. + let offsets = buffer![0u64, 3, 7, 7, 12].into_array(); + let mask = Mask::from_iter([false, true, false, true]); + let (range, elem_mask, new_off, kept) = run_scatter_gather(offsets, mask)?; + assert_eq!(range, 3..12); + // positions 3..7 (4 bits) and 7..12 (5 bits) — middle "gap" at 7..7 is zero-width. + assert_eq!(elem_mask, vec![true; 9]); + // Walk the new_offsets slice as u64. + let mut ctx = SESSION.create_execution_ctx(); + let new_off_prim = new_off.execute::(&mut ctx)?; + assert_eq!(new_off_prim.as_slice::(), &[0u64, 4, 9]); + assert_eq!(kept, 2); + Ok(()) + } + fn create_basic_list_array(nullable: bool) -> ArrayRef { let validity = if nullable { Validity::Array(BoolArray::from_iter([true, false, true]).into_array()) @@ -419,4 +681,51 @@ mod tests { assert_arrays_eq!(result, expected); Ok(()) } + + /// Build a list with 5 rows and lengths [2, 3, 0, 3, 2]. Mirrors `five_list_offsets()`. + fn create_wider_list_array(nullable: bool) -> ArrayRef { + let validity = if nullable { + Validity::Array(BoolArray::from_iter([true, true, false, true, true]).into_array()) + } else { + Validity::NonNullable + }; + ListArray::try_new( + buffer![10i32, 11, 20, 21, 22, 30, 31, 32, 40, 41].into_array(), + buffer![0u32, 2, 5, 5, 8, 10].into_array(), + validity, + ) + .expect("array is valid") + .into_array() + } + + #[rstest] + // Single bit set far from start — exercises sparse path with tight elements range. + #[case::single_middle(Mask::from_iter([false, false, false, true, false]), false)] + // Two far-apart rows — element_mask has a gap between kept spans. + #[case::two_far_apart(Mask::from_iter([true, false, false, true, false]), false)] + // Boundary rows — first and last list. + #[case::boundaries(Mask::from_iter([true, false, false, false, true]), false)] + // Kept row is the empty list (zero-width span). + #[case::kept_empty_row(Mask::from_iter([false, false, true, false, false]), false)] + // Sparse with nullable elements/validity child — exercises validity push-down. + #[case::sparse_nullable(Mask::from_iter([true, false, true, false, true]), true)] + // No rows kept — degenerate empty output. + #[case::all_false(Mask::new_false(5), false)] + #[tokio::test] + async fn projection_evaluation_sparse_mask_round_trips( + #[case] mask: Mask, + #[case] nullable: bool, + ) -> VortexResult<()> { + let list = create_wider_list_array(nullable); + let (segments, layout) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &SESSION)?; + + let result = reader + .projection_evaluation(&(0..5), &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let expected = list.filter(mask)?; + assert_arrays_eq!(result, expected); + Ok(()) + } } From a6ebb030925fa1fe47953a5ee6bbb50b7cd16b5d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 28 May 2026 17:35:23 -0400 Subject: [PATCH 24/37] shortcut on whole-chunk unmasked reads Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 68 +++++++++++++++++++----- 1 file changed, 54 insertions(+), 14 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 946386211c9..52c158edbfd 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -315,28 +315,67 @@ impl LayoutReader for ListReader { let validity_reader = self.validity.clone(); let session = self.session.clone(); let nullability = self.layout.dtype().nullability(); + let layout_row_count = self.layout.row_count(); + let elements_row_count = self.elements.row_count(); let expr = expr.clone(); let row_range = row_range.clone(); Ok(async move { - // Resolve offsets and the caller's mask in parallel. We need both before we can - // decide between the two paths below. - let (offsets, mask) = try_join!(offsets_fut, mask)?; - - // Path A: all-true mask — fetch the full bounded elements range, build a ListArray - // covering the entire row_range. The three children must be compositionally - // consistent here, so validity is read at full `row_range` length with an all-true - // mask too. + // Await the caller mask first so we can decide the read shape. Offsets is already + // in flight from above and overlaps this wait. For statically-resolved masks + // (`MaskFuture::new_true`, `MaskFuture::ready` of an already-known mask — the + // common case for both full scans and filter pushdown) the await is free. + let mask = mask.await?; + let validity_row_count = usize::try_from(row_range.end - row_range.start)?; + let is_whole_chunk = + row_range.start == 0 && row_range.end == layout_row_count; + + // Path A1: whole-chunk read with all-true mask. The elements bound is the whole + // elements buffer (`0..elements.row_count()`) and `offsets[0] == 0` by construction + // within a chunk, so we don't need to read offsets to know the bound and we don't + // need to rebase. Fire elements + validity in parallel with the already-in-flight + // offsets — a single `try_join!` over all three children. // - // Path B: sparse mask — scatter-gather. Bound the elements fetch to the tightest - // range covering the kept rows and pass an element-level mask so the elements child - // only materializes positions belonging to a kept row. Validity is fetched at - // `kept_count` length by pushing the caller mask down directly. - if mask.all_true() { + // Path A2: partial range, all-true mask. The elements bound is + // `offsets[a]..offsets[b]` so we have to await offsets before firing elements. + // + // Path B: sparse mask. Bound the elements fetch to the tightest range covering the + // kept rows and pass an element-level mask so the elements child only materializes + // kept-row positions. Validity is fetched at `kept_count` length by pushing the + // caller mask down directly. + if mask.all_true() && is_whole_chunk { + let elements_fut = elements_reader.projection_evaluation( + &(0..elements_row_count), + &root(), + MaskFuture::new_true(usize::try_from(elements_row_count)?), + )?; + let validity_fut = validity_reader + .as_ref() + .map(|v| { + v.projection_evaluation( + &row_range, + &root(), + MaskFuture::new_true(validity_row_count), + ) + }) + .transpose()?; + + let (offsets, elements, validity_array) = + try_join!(offsets_fut, elements_fut, async move { + match validity_fut { + Some(fut) => fut.await.map(Some), + None => Ok(None), + } + })?; + + let validity = create_validity(validity_array, nullability); + let array = ListArray::try_new(elements, offsets, validity)?.into_array(); + array.apply(&expr) + } else if mask.all_true() { + let offsets = offsets_fut.await?; let elements_range = calculate_elements_range(&offsets, &session)?; let rebased_offsets = rebase_offsets(offsets, elements_range.start)?; let elements_len = elements_range.end - elements_range.start; - let validity_row_count = usize::try_from(row_range.end - row_range.start)?; let validity_fut = validity_reader .as_ref() @@ -365,6 +404,7 @@ impl LayoutReader for ListReader { let array = ListArray::try_new(elements, rebased_offsets, validity)?.into_array(); array.apply(&expr) } else { + let offsets = offsets_fut.await?; let sg = compute_scatter_gather(&offsets, &mask, &session)?; let validity_fut = validity_reader From 57a1bb0d5af1c9f9516c8e806441a5e1db8b39dc Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 16 Jun 2026 17:37:58 -0700 Subject: [PATCH 25/37] add required fallback to ListLayoutStrategy for non-list input Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 7 ++- vortex-layout/src/layouts/list/writer.rs | 56 ++++++++++++++++++++---- 2 files changed, 54 insertions(+), 9 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 52c158edbfd..d0ae795fbf9 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -461,7 +461,12 @@ mod tests { fn flat_list_strategy() -> ListLayoutStrategy { let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)) + ListLayoutStrategy::new( + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + ) } async fn write_layout( diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index a2dbee5c546..0cd2a32a5c6 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -33,15 +33,19 @@ use crate::sequence::SequentialStream; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; -/// Strategy for writing list-typed arrays. +/// Strategy for writing list-typed arrays, with a fallback for non-list dtypes. /// -/// Single-chunk only. The strategy: +/// Single-chunk only. For list-typed input the strategy: /// 1. Canonicalizes the input chunk into a [`ListViewArray`]. /// 2. Calls [`list_from_list_view`] to rebuild it into zero-copy-to-list form /// (sorted, gapless, non-overlapping offsets) and produce a [`ListArray`]. /// 3. Writes the `elements`, `offsets`, and (when nullable) `validity` columns into /// separately configurable downstream strategies, producing a single [`ListLayout`]. /// +/// For input whose dtype is not [`DType::List`], the stream is forwarded unchanged to the +/// configured `fallback` strategy. This lets `ListLayoutStrategy` slot in as a leaf strategy in +/// a heterogeneous column writer where some columns are lists and others are not. +/// /// # Chunking /// /// `ListLayoutStrategy` bails on empty or multi-chunk input, matching the convention used by @@ -52,6 +56,7 @@ pub struct ListLayoutStrategy { elements: Arc, offsets: Arc, validity: Arc, + fallback: Arc, } impl ListLayoutStrategy { @@ -59,11 +64,13 @@ impl ListLayoutStrategy { elements: Arc, offsets: Arc, validity: Arc, + fallback: Arc, ) -> Self { Self { elements, offsets, validity, + fallback, } } } @@ -80,7 +87,11 @@ impl LayoutStrategy for ListLayoutStrategy { ) -> VortexResult { let dtype = stream.dtype().clone(); if !dtype.is_list() { - vortex_bail!("ListLayoutStrategy requires a List dtype, got {dtype}"); + // Non-list input: route to the configured fallback strategy unchanged. + return self + .fallback + .write_stream(ctx, segment_sink, stream, eof, session) + .await; } // Writer wants exactly one chunk @@ -159,9 +170,12 @@ impl LayoutStrategy for ListLayoutStrategy { } fn buffered_bytes(&self) -> u64 { - self.elements.buffered_bytes() + // A given input stream takes either the list path (elements + offsets + validity) or the + // fallback, so the back-pressure budget is the max of the two — not the sum. + let list_bytes = self.elements.buffered_bytes() + self.offsets.buffered_bytes() - + self.validity.buffered_bytes() + + self.validity.buffered_bytes(); + list_bytes.max(self.fallback.buffered_bytes()) } } @@ -212,7 +226,12 @@ mod tests { fn flat_list_strategy() -> ListLayoutStrategy { let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - ListLayoutStrategy::new(Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat)) + ListLayoutStrategy::new( + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + ) } async fn write(strategy: &S, array: ArrayRef) -> VortexResult { @@ -278,6 +297,16 @@ mod tests { Ok(()) } + /// Non-list input dispatches to the fallback strategy unchanged. + #[tokio::test] + async fn non_list_input_routes_to_fallback() -> VortexResult<()> { + let primitive = buffer![1i32, 2, 3].into_array(); + let layout = write(&flat_list_strategy(), primitive).await?; + // Fallback is FlatLayoutStrategy, so the result is a flat layout (not a list layout). + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + Ok(()) + } + #[tokio::test] async fn empty_stream_errors() { let segments = Arc::new(TestSegments::default()); @@ -335,7 +364,12 @@ mod tests { let flat: Arc = Arc::new(FlatLayoutStrategy::default()); let table_strategy: Arc = Arc::new(TableStrategy::new(Arc::clone(&flat), Arc::clone(&flat))); - let writer = ListLayoutStrategy::new(table_strategy, Arc::clone(&flat), Arc::clone(&flat)); + let writer = ListLayoutStrategy::new( + table_strategy, + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + ); let layout = write(&writer, list).await?; insta::assert_snapshot!(layout.display_tree(), @" @@ -368,8 +402,14 @@ mod tests { Arc::clone(&flat), Arc::clone(&flat), Arc::clone(&flat), + Arc::clone(&flat), )); - let writer = ListLayoutStrategy::new(inner_strategy, Arc::clone(&flat), Arc::clone(&flat)); + let writer = ListLayoutStrategy::new( + inner_strategy, + Arc::clone(&flat), + Arc::clone(&flat), + Arc::clone(&flat), + ); let layout = write(&writer, list).await?; insta::assert_snapshot!(layout.display_tree(), @" From 49ef6ca55d936fde83a8af81f9c2d3600b91be88 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 16 Jun 2026 17:46:34 -0700 Subject: [PATCH 26/37] fmt Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index d0ae795fbf9..ead9d5096e2 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -327,8 +327,7 @@ impl LayoutReader for ListReader { // common case for both full scans and filter pushdown) the await is free. let mask = mask.await?; let validity_row_count = usize::try_from(row_range.end - row_range.start)?; - let is_whole_chunk = - row_range.start == 0 && row_range.end == layout_row_count; + let is_whole_chunk = row_range.start == 0 && row_range.end == layout_row_count; // Path A1: whole-chunk read with all-true mask. The elements bound is the whole // elements buffer (`0..elements.row_count()`) and `offsets[0] == 0` by construction @@ -409,9 +408,7 @@ impl LayoutReader for ListReader { let validity_fut = validity_reader .as_ref() - .map(|v| { - v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask)) - }) + .map(|v| v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask))) .transpose()?; let elements_fut = elements_reader.projection_evaluation( &sg.elements_range, From 0cb0f426aa9e6d275de682306e5b400d734449b3 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 16 Jun 2026 17:54:07 -0700 Subject: [PATCH 27/37] cleanup Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 31 +++++++++++++++--------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 0cd2a32a5c6..e0d78973386 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -8,6 +8,7 @@ use futures::StreamExt; use futures::stream; use vortex_array::ArrayContext; use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::List; @@ -100,23 +101,13 @@ impl LayoutStrategy for ListLayoutStrategy { }; let (sequence_id, array) = chunk?; - // Run the execution loop until we have a List or ListView; skip work when the input is - // already in one of those forms. If the input is already a ListArray we extract its - // parts directly; if it's a ListViewArray we rebuild via `list_from_list_view`. let mut exec_ctx = session.create_execution_ctx(); - let canonical = array.execute_until::(&mut exec_ctx)?; let ListDataParts { elements, offsets, validity, .. - } = if let Some(list) = canonical.as_opt::() { - list.into_owned().into_data_parts() - } else if let Some(view) = canonical.as_opt::() { - list_from_list_view(view.into_owned())?.into_data_parts() - } else { - unreachable!("AnyList matcher guarantees List or ListView"); - }; + } = canonicalize_to_list_parts(array, &mut exec_ctx)?; // There is one extra element in `offsets` let row_count = offsets.len().saturating_sub(1); @@ -179,6 +170,24 @@ impl LayoutStrategy for ListLayoutStrategy { } } +/// Canonicalize a list-dtype array into [`ListDataParts`]. Short-circuits when the input is +/// already a `List` or `ListView` array — otherwise drives the execution loop until one of +/// those forms appears. `ListView` is rebuilt into zero-copy-to-list form via +/// [`list_from_list_view`] before its parts are extracted. +fn canonicalize_to_list_parts( + array: ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let canonical = array.execute_until::(exec_ctx)?; + if let Some(list) = canonical.as_opt::() { + Ok(list.into_owned().into_data_parts()) + } else if let Some(view) = canonical.as_opt::() { + Ok(list_from_list_view(view.into_owned())?.into_data_parts()) + } else { + unreachable!("AnyList matcher guarantees List or ListView") + } +} + /// Wrap a single array as a one-shot [`SendableSequentialStream`] for handoff to a child writer. fn single_chunk_stream( dtype: DType, From f046808ec52be929e355776af5b5b94bab9eb3f7 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 16 Jun 2026 21:23:11 -0700 Subject: [PATCH 28/37] fix rebase conflicts Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 3 +++ vortex-layout/src/layouts/list/reader.rs | 21 +++++++++++++++------ vortex-layout/src/layouts/list/writer.rs | 2 +- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index b7627fff230..464e0e8885b 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -24,6 +24,7 @@ use vortex_session::registry::ReadContext; use crate::LayoutChildType; use crate::LayoutEncodingRef; use crate::LayoutId; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::LayoutRef; use crate::VTable; @@ -105,12 +106,14 @@ impl VTable for List { name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(ListReader::try_new( layout.clone(), name, segment_source, session.clone(), + ctx, )?)) } diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index ead9d5096e2..05f002660c6 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::collections::BTreeSet; use std::ops::Range; use std::sync::Arc; @@ -32,7 +31,9 @@ use vortex_session::VortexSession; use crate::ArrayFuture; use crate::LayoutReader; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; +use crate::RowSplits; use crate::SplitRange; use crate::layouts::list::ListLayout; use crate::segments::SegmentSource; @@ -53,16 +54,19 @@ impl ListReader { name: Arc, segment_source: Arc, session: VortexSession, + ctx: &LayoutReaderContext, ) -> VortexResult { let elements = layout.elements().new_reader( format!("{name}.elements").into(), Arc::clone(&segment_source), &session, + ctx, )?; let offsets = layout.offsets().new_reader( format!("{name}.offsets").into(), Arc::clone(&segment_source), &session, + ctx, )?; let validity = layout .validity() @@ -71,6 +75,7 @@ impl ListReader { format!("{name}.validity").into(), Arc::clone(&segment_source), &session, + ctx, ) }) .transpose()?; @@ -274,7 +279,7 @@ impl LayoutReader for ListReader { &self, field_mask: &[FieldMask], split_range: &SplitRange, - splits: &mut BTreeSet, + splits: &mut RowSplits, ) -> VortexResult<()> { self.offsets .register_splits(field_mask, split_range, splits)?; @@ -668,7 +673,8 @@ mod tests { let list = create_basic_list_array(false); let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; - let reader = layout.new_reader("".into(), segments, &SESSION)?; + let ctx = LayoutReaderContext::new(); + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; let reader = reader .as_any() .downcast_ref::() @@ -693,10 +699,11 @@ mod tests { #[case] nullable: bool, ) -> VortexResult<()> { let list = create_basic_list_array(nullable); + let ctx = LayoutReaderContext::new(); let len = usize::try_from(row_range.end - row_range.start)?; let (segments, layout) = write_layout(&flat_list_strategy(), list.clone()).await?; - let reader = layout.new_reader("".into(), segments, &SESSION)?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; let result = reader .projection_evaluation(&row_range, &root(), MaskFuture::new_true(len))? @@ -711,8 +718,9 @@ mod tests { #[tokio::test] async fn projection_evaluation_applies_mask() -> VortexResult<()> { let list = create_basic_list_array(false); + let ctx = LayoutReaderContext::new(); let (segments, layout) = write_layout(&flat_list_strategy(), list.clone()).await?; - let reader = layout.new_reader("".into(), segments, &SESSION)?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; let mask = Mask::from_iter([true, false, true]); let result = reader @@ -759,8 +767,9 @@ mod tests { #[case] nullable: bool, ) -> VortexResult<()> { let list = create_wider_list_array(nullable); + let ctx = LayoutReaderContext::new(); let (segments, layout) = write_layout(&flat_list_strategy(), list.clone()).await?; - let reader = layout.new_reader("".into(), segments, &SESSION)?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; let result = reader .projection_evaluation(&(0..5), &root(), MaskFuture::ready(mask.clone()))? diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index e0d78973386..75343434439 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -182,7 +182,7 @@ fn canonicalize_to_list_parts( if let Some(list) = canonical.as_opt::() { Ok(list.into_owned().into_data_parts()) } else if let Some(view) = canonical.as_opt::() { - Ok(list_from_list_view(view.into_owned())?.into_data_parts()) + Ok(list_from_list_view(view.into_owned(), exec_ctx)?.into_data_parts()) } else { unreachable!("AnyList matcher guarantees List or ListView") } From 1abc6c2a2c2d362a1a65add749d238e03d8182bb Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 16 Jun 2026 21:50:01 -0700 Subject: [PATCH 29/37] use ListLayoutStrategy as default leaf under unstable_encodings Signed-off-by: Matt Katz --- vortex-file/src/strategy.rs | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 804218779c8..67b5c386599 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -48,6 +48,8 @@ use vortex_layout::layouts::compressed::CompressingStrategy; use vortex_layout::layouts::compressed::CompressorPlugin; use vortex_layout::layouts::dict::writer::DictStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +#[cfg(feature = "unstable_encodings")] +use vortex_layout::layouts::list::writer::ListLayoutStrategy; use vortex_layout::layouts::repartition::RepartitionStrategy; use vortex_layout::layouts::repartition::RepartitionWriterOptions; use vortex_layout::layouts::table::TableStrategy; @@ -240,8 +242,21 @@ impl WriteStrategyBuilder { Arc::new(FlatLayoutStrategy::default()) }; - // 7. for each chunk create a flat layout - let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); + // 7. for each chunk create a layout. Under the `unstable_encodings` feature, list-typed + // chunks route through `ListLayoutStrategy` (separately-addressable elements/offsets/ + // validity sub-layouts; non-list chunks fall through its built-in fallback to `flat`). + // Otherwise everything goes through the flat strategy. + #[cfg(feature = "unstable_encodings")] + let leaf: Arc = Arc::new(ListLayoutStrategy::new( + Arc::clone(&flat), // elements + Arc::clone(&flat), // offsets + Arc::clone(&flat), // validity + Arc::clone(&flat), // fallback for non-list dtypes + )); + #[cfg(not(feature = "unstable_encodings"))] + let leaf: Arc = Arc::clone(&flat); + + let chunked = ChunkedLayoutStrategy::new(leaf); // 6. buffer chunks so they end up with closer segment ids physically let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB From 4cba2eca7d10feca3c39985d55d4cf7f2c8dccd1 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 17 Jun 2026 11:26:03 -0700 Subject: [PATCH 30/37] recurse into nested lists Signed-off-by: Matt Katz --- vortex-file/src/strategy.rs | 5 +- vortex-layout/src/layouts/list/writer.rs | 73 ++++++++++++++++++------ 2 files changed, 60 insertions(+), 18 deletions(-) diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 67b5c386599..aa4a6bf64df 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -245,10 +245,11 @@ impl WriteStrategyBuilder { // 7. for each chunk create a layout. Under the `unstable_encodings` feature, list-typed // chunks route through `ListLayoutStrategy` (separately-addressable elements/offsets/ // validity sub-layouts; non-list chunks fall through its built-in fallback to `flat`). - // Otherwise everything goes through the flat strategy. + // Nested lists (`list>`) recurse, shredding each level into its own + // `ListLayout`. Otherwise everything goes through the flat strategy. #[cfg(feature = "unstable_encodings")] let leaf: Arc = Arc::new(ListLayoutStrategy::new( - Arc::clone(&flat), // elements + Arc::clone(&flat), // elements (non-list); list elements recurse into a nested ListLayout Arc::clone(&flat), // offsets Arc::clone(&flat), // validity Arc::clone(&flat), // fallback for non-list dtypes diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 75343434439..16cd8066cb0 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -43,6 +43,15 @@ use crate::sequence::SequentialStreamExt; /// 3. Writes the `elements`, `offsets`, and (when nullable) `validity` columns into /// separately configurable downstream strategies, producing a single [`ListLayout`]. /// +/// # Nested lists +/// +/// When the `elements` column is itself list-typed (e.g. `list>`), the strategy +/// recurses into a clone of itself instead of using the configured `elements` strategy, so the +/// inner list is shredded into a nested [`ListLayout`] rather than written as a single opaque +/// chunk. This mirrors how [`TableStrategy`](crate::layouts::table::TableStrategy) descends into +/// nested struct fields. The configured `elements` strategy is therefore used for the innermost, +/// non-list values. +/// /// For input whose dtype is not [`DType::List`], the stream is forwarded unchanged to the /// configured `fallback` strategy. This lets `ListLayoutStrategy` slot in as a leaf strategy in /// a heterogeneous column writer where some columns are lists and others are not. @@ -53,6 +62,7 @@ use crate::sequence::SequentialStreamExt; /// [`FlatLayoutStrategy`](crate::layouts::flat::writer::FlatLayoutStrategy). /// /// [`ListArray`]: vortex_array::arrays::ListArray +#[derive(Clone)] pub struct ListLayoutStrategy { elements: Arc, offsets: Arc, @@ -120,6 +130,17 @@ impl LayoutStrategy for ListLayoutStrategy { }) .transpose()?; + // Recurse into a clone of ourselves when the elements are themselves list-typed, so that + // `list>` writes as `ListLayout { elements: ListLayout { .. } }` rather than + // collapsing the inner list into a single opaque chunk. Non-list elements use the + // configured `elements` strategy. Mirrors how `TableStrategy` descends into nested struct + // fields. + let elements_strategy: Arc = if elements.dtype().is_list() { + Arc::new(self.clone()) + } else { + Arc::clone(&self.elements) + }; + // Spawn each child write onto the runtime so they run concurrently let handle = session.handle(); let (elements_task, offsets_task, validity_task) = { @@ -138,7 +159,7 @@ impl LayoutStrategy for ListLayoutStrategy { }) }; ( - spawn_layout_writer(Arc::clone(&self.elements), elements), + spawn_layout_writer(elements_strategy, elements), spawn_layout_writer(Arc::clone(&self.offsets), offsets), validity_array.map(|arr| spawn_layout_writer(Arc::clone(&self.validity), arr)), ) @@ -406,21 +427,9 @@ mod tests { )? .into_array(); - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let inner_strategy: Arc = Arc::new(ListLayoutStrategy::new( - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - )); - let writer = ListLayoutStrategy::new( - inner_strategy, - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - ); - - let layout = write(&writer, list).await?; + // The all-flat strategy still recurses on list-typed elements: the inner `list` is + // shredded into a nested ListLayout rather than written as a single opaque flat chunk. + let layout = write(&flat_list_strategy(), list).await?; insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 @@ -431,6 +440,38 @@ mod tests { Ok(()) } + /// Recursion is unbounded: `list>>` produces three nested ListLayouts. + #[tokio::test] + async fn list_of_list_of_list_tree() -> VortexResult<()> { + let innermost = ListArray::try_new( + buffer![1i32, 2, 3, 4].into_array(), + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let middle = ListArray::try_new( + innermost, + buffer![0u32, 2].into_array(), + Validity::NonNullable, + )? + .into_array(); + let outer = + ListArray::try_new(middle, buffer![0u32, 1].into_array(), Validity::NonNullable)? + .into_array(); + + let layout = write(&flat_list_strategy(), outer).await?; + insta::assert_snapshot!(layout.display_tree(), @" + vortex.list, dtype: list(list(list(i32))), children: 2 + ├── elements: vortex.list, dtype: list(list(i32)), children: 2 + │ ├── elements: vortex.list, dtype: list(i32), children: 2 + │ │ ├── elements: vortex.flat, dtype: i32, segment: 2 + │ │ └── offsets: vortex.flat, dtype: u32, segment: 3 + │ └── offsets: vortex.flat, dtype: u32, segment: 1 + └── offsets: vortex.flat, dtype: u32, segment: 0 + "); + Ok(()) + } + #[tokio::test] async fn chunked_list_input_with_chunked_strategy_succeeds() -> VortexResult<()> { let chunk0 = ListArray::try_new( From a7614ea3dab4c189eb98411955edfcb56c7a7ed5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 22 Jun 2026 15:40:37 -0700 Subject: [PATCH 31/37] fix: update ListLayout::build to use LayoutBuildContext after trait change Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 464e0e8885b..899b1d547ee 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -19,8 +19,8 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; -use vortex_session::registry::ReadContext; +use crate::LayoutBuildContext; use crate::LayoutChildType; use crate::LayoutEncodingRef; use crate::LayoutId; @@ -124,7 +124,7 @@ impl VTable for List { metadata: &::Output, _segment_ids: Vec, children: &dyn LayoutChildren, - _ctx: &ReadContext, + _ctx: &LayoutBuildContext<'_>, ) -> VortexResult { validate_children(dtype, children.nchildren())?; From daed9301b1807d742ba2f858a0e8e0d067ae2123 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 22 Jun 2026 16:10:13 -0700 Subject: [PATCH 32/37] comments Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/mod.rs | 24 +++++++++--------------- vortex-layout/src/layouts/list/reader.rs | 3 ++- 2 files changed, 11 insertions(+), 16 deletions(-) diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 899b1d547ee..ba926e173f6 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -15,7 +15,7 @@ use vortex_array::dtype::PType; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; @@ -75,12 +75,11 @@ impl VTable for List { } fn nchildren(layout: &Self::Layout) -> usize { - let mut n = NUM_CHILDREN_NON_NULLABLE; if layout.dtype.is_nullable() { - n += 1; + NUM_CHILDREN_NON_NULLABLE + 1 + } else { + NUM_CHILDREN_NON_NULLABLE } - - n } fn child(layout: &Self::Layout, idx: usize) -> VortexResult { @@ -172,19 +171,14 @@ impl VTable for List { } /// Validates expected number of children based on `dtype` -#[inline] fn validate_children(dtype: &DType, n_children: usize) -> VortexResult<()> { - let mut expected = NUM_CHILDREN_NON_NULLABLE; - - if dtype.is_nullable() { - expected += 1; + let expected = if dtype.is_nullable() { + NUM_CHILDREN_NON_NULLABLE + 1 + } else { + NUM_CHILDREN_NON_NULLABLE }; - vortex_ensure!( - n_children == expected, - "ListLayout expects {expected} children, got {n_children}", - ); - + vortex_ensure_eq!(n_children, expected); Ok(()) } diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 05f002660c6..c9c79164adf 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -26,6 +26,7 @@ use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -295,7 +296,7 @@ impl LayoutReader for ListReader { _expr: &Expression, mask: Mask, ) -> VortexResult { - // All stats based pruning should already be done upstream + // All stats-based pruning should already be done upstream. Ok(MaskFuture::ready(mask)) } From 51d044e591642c2a75584772d34ac8325e00e55e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 22 Jun 2026 16:10:36 -0700 Subject: [PATCH 33/37] fix Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index c9c79164adf..a8ebe727b5b 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -26,7 +26,6 @@ use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_mask::Mask; use vortex_session::VortexSession; From 6f73444c1dad5e6a48c7505c50cd88ec08a43033 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 22 Jun 2026 18:20:13 -0700 Subject: [PATCH 34/37] clean up writer Signed-off-by: Matt Katz --- vortex-file/src/strategy.rs | 15 ++-- vortex-layout/src/layouts/list/reader.rs | 9 +-- vortex-layout/src/layouts/list/writer.rs | 99 +++++++++++------------- 3 files changed, 56 insertions(+), 67 deletions(-) diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index aa4a6bf64df..777846c553e 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -248,12 +248,15 @@ impl WriteStrategyBuilder { // Nested lists (`list>`) recurse, shredding each level into its own // `ListLayout`. Otherwise everything goes through the flat strategy. #[cfg(feature = "unstable_encodings")] - let leaf: Arc = Arc::new(ListLayoutStrategy::new( - Arc::clone(&flat), // elements (non-list); list elements recurse into a nested ListLayout - Arc::clone(&flat), // offsets - Arc::clone(&flat), // validity - Arc::clone(&flat), // fallback for non-list dtypes - )); + let leaf: Arc = Arc::new( + // Thread the configured `flat` (which carries `allow_encodings` / any custom flat + // override) through every child; list elements still recurse into a nested ListLayout. + ListLayoutStrategy::default() + .with_elements(Arc::clone(&flat)) + .with_offsets(Arc::clone(&flat)) + .with_validity(Arc::clone(&flat)) + .with_fallback(Arc::clone(&flat)), + ); #[cfg(not(feature = "unstable_encodings"))] let leaf: Arc = Arc::clone(&flat); diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index a8ebe727b5b..c282f43b39a 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -453,7 +453,6 @@ mod tests { use super::*; use crate::LayoutRef; use crate::LayoutStrategy; - use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::writer::ListLayoutStrategy; use crate::segments::SegmentSource; use crate::segments::TestSegments; @@ -462,13 +461,7 @@ mod tests { use crate::test::SESSION; fn flat_list_strategy() -> ListLayoutStrategy { - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - ListLayoutStrategy::new( - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - ) + ListLayoutStrategy::default() } async fn write_layout( diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 16cd8066cb0..6de14d5e1dd 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -25,6 +25,7 @@ use vortex_session::VortexSession; use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::ListLayout; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; @@ -43,15 +44,6 @@ use crate::sequence::SequentialStreamExt; /// 3. Writes the `elements`, `offsets`, and (when nullable) `validity` columns into /// separately configurable downstream strategies, producing a single [`ListLayout`]. /// -/// # Nested lists -/// -/// When the `elements` column is itself list-typed (e.g. `list>`), the strategy -/// recurses into a clone of itself instead of using the configured `elements` strategy, so the -/// inner list is shredded into a nested [`ListLayout`] rather than written as a single opaque -/// chunk. This mirrors how [`TableStrategy`](crate::layouts::table::TableStrategy) descends into -/// nested struct fields. The configured `elements` strategy is therefore used for the innermost, -/// non-list values. -/// /// For input whose dtype is not [`DType::List`], the stream is forwarded unchanged to the /// configured `fallback` strategy. This lets `ListLayoutStrategy` slot in as a leaf strategy in /// a heterogeneous column writer where some columns are lists and others are not. @@ -70,22 +62,46 @@ pub struct ListLayoutStrategy { fallback: Arc, } -impl ListLayoutStrategy { - pub fn new( - elements: Arc, - offsets: Arc, - validity: Arc, - fallback: Arc, - ) -> Self { +impl Default for ListLayoutStrategy { + /// Routes every child (elements, offsets, validity) and the non-list fallback through + /// [`FlatLayoutStrategy`]. Override individual children with the `with_*` builder methods. + fn default() -> Self { + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); Self { - elements, - offsets, - validity, - fallback, + elements: Arc::clone(&flat), + offsets: Arc::clone(&flat), + validity: Arc::clone(&flat), + fallback: flat, } } } +impl ListLayoutStrategy { + /// Strategy for the `elements` child. + pub fn with_elements(mut self, elements: Arc) -> Self { + self.elements = elements; + self + } + + /// Strategy for the `offsets` child. + pub fn with_offsets(mut self, offsets: Arc) -> Self { + self.offsets = offsets; + self + } + + /// Strategy for the `validity` child (written only when the list is nullable). + pub fn with_validity(mut self, validity: Arc) -> Self { + self.validity = validity; + self + } + + /// Strategy for non-list input, which is forwarded through this strategy unchanged. + pub fn with_fallback(mut self, fallback: Arc) -> Self { + self.fallback = fallback; + self + } +} + #[async_trait] impl LayoutStrategy for ListLayoutStrategy { async fn write_stream( @@ -130,17 +146,6 @@ impl LayoutStrategy for ListLayoutStrategy { }) .transpose()?; - // Recurse into a clone of ourselves when the elements are themselves list-typed, so that - // `list>` writes as `ListLayout { elements: ListLayout { .. } }` rather than - // collapsing the inner list into a single opaque chunk. Non-list elements use the - // configured `elements` strategy. Mirrors how `TableStrategy` descends into nested struct - // fields. - let elements_strategy: Arc = if elements.dtype().is_list() { - Arc::new(self.clone()) - } else { - Arc::clone(&self.elements) - }; - // Spawn each child write onto the runtime so they run concurrently let handle = session.handle(); let (elements_task, offsets_task, validity_task) = { @@ -159,7 +164,7 @@ impl LayoutStrategy for ListLayoutStrategy { }) }; ( - spawn_layout_writer(elements_strategy, elements), + spawn_layout_writer(Arc::clone(&self.elements), elements), spawn_layout_writer(Arc::clone(&self.offsets), offsets), validity_array.map(|arr| spawn_layout_writer(Arc::clone(&self.validity), arr)), ) @@ -182,8 +187,6 @@ impl LayoutStrategy for ListLayoutStrategy { } fn buffered_bytes(&self) -> u64 { - // A given input stream takes either the list path (elements + offsets + validity) or the - // fallback, so the back-pressure budget is the max of the two — not the sum. let list_bytes = self.elements.buffered_bytes() + self.offsets.buffered_bytes() + self.validity.buffered_bytes(); @@ -255,13 +258,7 @@ mod tests { use crate::test::SESSION; fn flat_list_strategy() -> ListLayoutStrategy { - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - ListLayoutStrategy::new( - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - ) + ListLayoutStrategy::default() } async fn write(strategy: &S, array: ArrayRef) -> VortexResult { @@ -332,7 +329,6 @@ mod tests { async fn non_list_input_routes_to_fallback() -> VortexResult<()> { let primitive = buffer![1i32, 2, 3].into_array(); let layout = write(&flat_list_strategy(), primitive).await?; - // Fallback is FlatLayoutStrategy, so the result is a flat layout (not a list layout). insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); Ok(()) } @@ -394,12 +390,7 @@ mod tests { let flat: Arc = Arc::new(FlatLayoutStrategy::default()); let table_strategy: Arc = Arc::new(TableStrategy::new(Arc::clone(&flat), Arc::clone(&flat))); - let writer = ListLayoutStrategy::new( - table_strategy, - Arc::clone(&flat), - Arc::clone(&flat), - Arc::clone(&flat), - ); + let writer = ListLayoutStrategy::default().with_elements(table_strategy); let layout = write(&writer, list).await?; insta::assert_snapshot!(layout.display_tree(), @" @@ -427,9 +418,9 @@ mod tests { )? .into_array(); - // The all-flat strategy still recurses on list-typed elements: the inner `list` is - // shredded into a nested ListLayout rather than written as a single opaque flat chunk. - let layout = write(&flat_list_strategy(), list).await?; + let writer = + ListLayoutStrategy::default().with_elements(Arc::new(ListLayoutStrategy::default())); + let layout = write(&writer, list).await?; insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 @@ -440,7 +431,6 @@ mod tests { Ok(()) } - /// Recursion is unbounded: `list>>` produces three nested ListLayouts. #[tokio::test] async fn list_of_list_of_list_tree() -> VortexResult<()> { let innermost = ListArray::try_new( @@ -459,7 +449,10 @@ mod tests { ListArray::try_new(middle, buffer![0u32, 1].into_array(), Validity::NonNullable)? .into_array(); - let layout = write(&flat_list_strategy(), outer).await?; + let writer = ListLayoutStrategy::default().with_elements(Arc::new( + ListLayoutStrategy::default().with_elements(Arc::new(ListLayoutStrategy::default())), + )); + let layout = write(&writer, outer).await?; insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(list(i32))), children: 2 ├── elements: vortex.list, dtype: list(list(i32)), children: 2 From 3bf85b7a5edee0a066cab3bb0982f3c8df596ca2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 23 Jun 2026 12:43:59 -0700 Subject: [PATCH 35/37] Improve list reader Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 611 ++++++++++++++++------- 1 file changed, 444 insertions(+), 167 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index c282f43b39a..2b514862fe7 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -5,6 +5,7 @@ use std::ops::Range; use std::sync::Arc; use futures::FutureExt; +use futures::future::BoxFuture; use futures::try_join; use vortex_array::Array; use vortex_array::ArrayRef; @@ -18,12 +19,18 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; +use vortex_array::dtype::IntegerPType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; +use vortex_array::expr::is_root; +use vortex_array::expr::not; use vortex_array::expr::root; +use vortex_array::scalar_fn::fns::is_not_null::IsNotNull; +use vortex_array::scalar_fn::fns::is_null::IsNull; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -38,7 +45,10 @@ use crate::SplitRange; use crate::layouts::list::ListLayout; use crate::segments::SegmentSource; +type OptionalArrayFuture = BoxFuture<'static, VortexResult>>; + /// Reader for [`ListLayout`]. +#[derive(Clone)] pub struct ListReader { layout: ListLayout, name: Arc, @@ -90,12 +100,83 @@ impl ListReader { }) } + /// Projection for [`ExprClass::Validity`] expressions (`is_null` / `is_not_null` of the list): + /// reads only the validity child — synthesizing all-valid for a non-nullable list — and never + /// touches the offsets or elements. + fn project_validity( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let validity_reader = self.validity.clone(); + let nullability = self.layout.dtype().nullability(); + let row_range = row_range.clone(); + // Evaluate the rewritten expression against the validity bool array (true == valid row). + let rewritten = rewrite_validity_expr(expr)?; + + Ok(async move { + let mask = mask.await?; + let row_count = usize::try_from(row_range.end - row_range.start)?; + let out_len = if mask.all_true() { + row_count + } else { + mask.true_count() + }; + + let validity_array = match validity_reader.as_ref() { + Some(v) => Some( + v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask))? + .await?, + ), + None => None, + }; + + let validity = create_validity(validity_array, nullability).to_array(out_len); + validity.apply(&rewritten) + } + .boxed()) + } + + /// Projection for [`ExprClass::Elements`] expressions (everything else): materializes the list + /// (offsets + elements + validity) and applies the expression. + fn project_elements( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + // Fire the offsets read before cloning so it overlaps the mask await below. + let projection = ElementsProjection { + reader: self.clone(), + expr: expr.clone(), + row_range: row_range.clone(), + offsets: self.fetch_offsets(row_range)?, + }; + + Ok(async move { + // Await the caller mask to decide the read shape. Offsets is already in flight and + // overlaps this wait; for statically-resolved masks the await is free. + let mask = mask.await?; + let is_whole_chunk = projection.row_range.start == 0 + && projection.row_range.end == projection.reader.layout.row_count(); + + if mask.all_true() && is_whole_chunk { + projection.project_whole_chunk().await + } else if mask.all_true() { + projection.project_full_range().await + } else { + projection.project_sparse(mask).await + } + } + .boxed()) + } + + /// Fire the offsets read for `row_range`. The offsets child has an extra entry, so reading + /// `row_range` maps to offsets in `[row_range.start..row_range.end + 1)`. fn fetch_offsets(&self, row_range: &Range) -> VortexResult { - // The offsets child has an extra entry, so reading `row_range` maps to offsets in - // `[row_range.start..row_range.end + 1)`. let offsets_range = row_range.start..(row_range.end + 1); let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; - self.offsets.projection_evaluation( &offsets_range, &root(), @@ -145,6 +226,68 @@ fn create_validity(validity_array: Option, nullability: Nullability) - } } +/// The deepest list child an expression needs, cheapest first. +/// +/// Drives "fetch as little as possible": a projection/filter that only inspects the list's +/// null-ness needs the validity child; everything else needs the element values. The ordering +/// `Validity < Elements` lets us take the max over the operands of a compound expression. +// TODO: have `filter_evaluation` use this too. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum ExprClass { + /// Only the list's validity is needed (`is_null` / `is_not_null` of the list itself). + Validity, + /// The element values are needed (everything else). + Elements, +} + +/// Classify `expr` by the deepest list child it touches, where `root()` is the list. +/// +/// Only the exact shapes `is_null(root())` / `is_not_null(root())` (validity) are recognized. Every +/// other access to the list, including a bare `root()`, falls through to [`ExprClass::Elements`], +/// which is always correct. +fn classify(expr: &Expression) -> ExprClass { + // `is_null(root())` / `is_not_null(root())` need only the list's own validity. Note this is + // the list's null-ness, not the validity of some derived value, so the child must be `root()`. + if (expr.is::() || expr.is::()) + && expr.children().len() == 1 + && is_root(expr.child(0)) + { + return ExprClass::Validity; + } + + // A bare reference to the list needs its elements. + if is_root(expr) { + return ExprClass::Elements; + } + + // Otherwise the requirement is the max over the operands. Operands that never touch the list + // (e.g. literals) contribute nothing, so an expression that never references `root()` is + // treated as the cheapest class. + expr.children() + .iter() + .map(classify) + .max() + .unwrap_or(ExprClass::Validity) +} + +/// Rewrite a validity-class expression so it can be evaluated against the list's validity bool +/// array (`true` == valid row): `is_not_null(root())` becomes `root()` and `is_null(root())` +/// becomes `not(root())`. All other nodes are rebuilt with rewritten children. +fn rewrite_validity_expr(expr: &Expression) -> VortexResult { + if expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) { + return Ok(root()); + } + if expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) { + return Ok(not(root())); + } + let children = expr + .children() + .iter() + .map(rewrite_validity_expr) + .collect::>>()?; + expr.clone().with_children(children) +} + /// Plan for fetching only the elements needed to materialize the kept list rows under a sparse /// row mask, plus the offsets array we'll hand to `ListArray::try_new` for those kept rows. /// @@ -206,55 +349,73 @@ fn compute_scatter_gather( }); } - // Within each macro arm, `O` is a concrete primitive integer type. Offsets are non-negative - // by construction; we materialize spans in `usize` so the element-mask construction is - // straightforward. vortex_array::match_each_integer_ptype!(ptype, |O| { - let off = prim_offsets.as_slice::(); + compute_scatter_gather_typed::(prim_offsets.as_slice::(), mask, kept_count) + }) +} - let mut spans: Vec<(usize, usize)> = Vec::with_capacity(kept_count); - let mut new_off: Vec = Vec::with_capacity(kept_count + 1); - new_off.push(O::default()); - let mut cumulative: O = O::default(); +fn compute_scatter_gather_typed( + offsets: &[O], + mask: &Mask, + kept_count: usize, +) -> VortexResult +where + O: IntegerPType, + usize: TryFrom, + VortexError: From<>::Error>, +{ + let mut new_off: Vec = Vec::with_capacity(kept_count + 1); + let mut element_slices: Vec<(usize, usize)> = Vec::with_capacity(kept_count); + new_off.push(O::default()); + let mut cumulative: O = O::default(); + let mut range_start: Option = None; + let mut range_end = 0usize; + + { + let mut keep_row = |i: usize| -> VortexResult<()> { + let start_offset = offsets[i]; + let end_offset = offsets[i + 1]; + cumulative += end_offset - start_offset; + new_off.push(cumulative); + + let start = usize::try_from(start_offset)?; + let end = usize::try_from(end_offset)?; + if start < end { + let start_base = *range_start.get_or_insert(start); + element_slices.push((start - start_base, end - start_base)); + range_end = end; + } + Ok(()) + }; // `mask.indices()` returns the set bit positions for `Values` masks; `AllTrue` is rare // here (caller checks density) but we handle it via fallback iteration. - let indices_owned: Vec = match mask.indices() { - vortex_mask::AllOr::All => (0..mask.len()).collect(), - vortex_mask::AllOr::None => Vec::new(), - vortex_mask::AllOr::Some(idxs) => idxs.to_vec(), - }; - for &i in &indices_owned { - let s = off[i]; - let e = off[i + 1]; - spans.push((usize::try_from(s)?, usize::try_from(e)?)); - cumulative += e - s; - new_off.push(cumulative); + match mask.indices() { + vortex_mask::AllOr::All => { + for i in 0..mask.len() { + keep_row(i)?; + } + } + vortex_mask::AllOr::None => {} + vortex_mask::AllOr::Some(idxs) => { + for &i in idxs { + keep_row(i)?; + } + } } + } - let range_start = spans[0].0; - let range_end = spans[spans.len() - 1].1; - - // Element-level mask: same length as the bounded elements range, true bits at positions - // that lie inside any kept span. `Mask::from_slices` rejects zero-width spans, so drop - // empty-list rows here. - let element_slices: Vec<(usize, usize)> = spans - .iter() - .filter(|(s, e)| s < e) - .map(|(s, e)| (s - range_start, e - range_start)) - .collect(); - let element_mask = Mask::from_slices(range_end - range_start, element_slices); - - let new_offsets = - Array::::new::(Buffer::::from(new_off), Validity::NonNullable) - .into_array(); - - Ok(ScatterGather { - elements_range: u64::try_from(range_start)?..u64::try_from(range_end)?, - element_mask, - new_offsets, - kept_count, - }) + let range_start = range_start.unwrap_or(0); + let element_mask = Mask::from_slices(range_end - range_start, element_slices); + let new_offsets = + Array::::new::(Buffer::::from(new_off), Validity::NonNullable) + .into_array(); + + Ok(ScatterGather { + elements_range: u64::try_from(range_start)?..u64::try_from(range_end)?, + element_mask, + new_offsets, + kept_count, }) } @@ -314,126 +475,164 @@ impl LayoutReader for ListReader { expr: &Expression, mask: MaskFuture, ) -> VortexResult { - let offsets_fut = self.fetch_offsets(row_range)?; - - let elements_reader = Arc::clone(&self.elements); - let validity_reader = self.validity.clone(); - let session = self.session.clone(); - let nullability = self.layout.dtype().nullability(); - let layout_row_count = self.layout.row_count(); - let elements_row_count = self.elements.row_count(); - let expr = expr.clone(); - let row_range = row_range.clone(); + // Read as little as possible based on which list children the expression needs. + match classify(expr) { + ExprClass::Validity => self.project_validity(row_range, expr, mask), + ExprClass::Elements => self.project_elements(row_range, expr, mask), + } + } +} - Ok(async move { - // Await the caller mask first so we can decide the read shape. Offsets is already - // in flight from above and overlaps this wait. For statically-resolved masks - // (`MaskFuture::new_true`, `MaskFuture::ready` of an already-known mask — the - // common case for both full scans and filter pushdown) the await is free. - let mask = mask.await?; - let validity_row_count = usize::try_from(row_range.end - row_range.start)?; - let is_whole_chunk = row_range.start == 0 && row_range.end == layout_row_count; - - // Path A1: whole-chunk read with all-true mask. The elements bound is the whole - // elements buffer (`0..elements.row_count()`) and `offsets[0] == 0` by construction - // within a chunk, so we don't need to read offsets to know the bound and we don't - // need to rebase. Fire elements + validity in parallel with the already-in-flight - // offsets — a single `try_join!` over all three children. - // - // Path A2: partial range, all-true mask. The elements bound is - // `offsets[a]..offsets[b]` so we have to await offsets before firing elements. - // - // Path B: sparse mask. Bound the elements fetch to the tightest range covering the - // kept rows and pass an element-level mask so the elements child only materializes - // kept-row positions. Validity is fetched at `kept_count` length by pushing the - // caller mask down directly. - if mask.all_true() && is_whole_chunk { - let elements_fut = elements_reader.projection_evaluation( - &(0..elements_row_count), - &root(), - MaskFuture::new_true(usize::try_from(elements_row_count)?), - )?; - let validity_fut = validity_reader - .as_ref() - .map(|v| { - v.projection_evaluation( - &row_range, - &root(), - MaskFuture::new_true(validity_row_count), - ) - }) - .transpose()?; - - let (offsets, elements, validity_array) = - try_join!(offsets_fut, elements_fut, async move { - match validity_fut { - Some(fut) => fut.await.map(Some), - None => Ok(None), - } - })?; - - let validity = create_validity(validity_array, nullability); - let array = ListArray::try_new(elements, offsets, validity)?.into_array(); - array.apply(&expr) - } else if mask.all_true() { - let offsets = offsets_fut.await?; - let elements_range = calculate_elements_range(&offsets, &session)?; - let rebased_offsets = rebase_offsets(offsets, elements_range.start)?; - let elements_len = elements_range.end - elements_range.start; - - let validity_fut = validity_reader - .as_ref() - .map(|v| { - v.projection_evaluation( - &row_range, - &root(), - MaskFuture::new_true(validity_row_count), - ) - }) - .transpose()?; - let elements_fut = elements_reader.projection_evaluation( - &elements_range, - &root(), - MaskFuture::new_true(usize::try_from(elements_len)?), - )?; - - let (elements, validity_array) = try_join!(elements_fut, async move { - match validity_fut { - Some(fut) => fut.await.map(Some), - None => Ok(None), - } - })?; - - let validity = create_validity(validity_array, nullability); - let array = ListArray::try_new(elements, rebased_offsets, validity)?.into_array(); - array.apply(&expr) - } else { - let offsets = offsets_fut.await?; - let sg = compute_scatter_gather(&offsets, &mask, &session)?; - - let validity_fut = validity_reader - .as_ref() - .map(|v| v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask))) - .transpose()?; - let elements_fut = elements_reader.projection_evaluation( - &sg.elements_range, - &root(), - MaskFuture::ready(sg.element_mask), - )?; - - let (elements, validity_array) = try_join!(elements_fut, async move { - match validity_fut { - Some(fut) => fut.await.map(Some), - None => Ok(None), - } - })?; - - let validity = create_validity(validity_array, nullability); - let array = ListArray::try_new(elements, sg.new_offsets, validity)?.into_array(); - array.apply(&expr) - } +/// Fetch the validity child for `row_range` under `mask`, yielding `None` for a non-nullable list +/// (which has no validity child). +fn fetch_validity( + validity: Option<&LayoutReaderRef>, + row_range: &Range, + mask: MaskFuture, +) -> VortexResult { + let fut = validity + .map(|v| v.projection_evaluation(row_range, &root(), mask)) + .transpose()?; + Ok(async move { + match fut { + Some(f) => f.await.map(Some), + None => Ok(None), } - .boxed()) + } + .boxed()) +} + +struct ListParts { + elements: ArrayRef, + offsets: ArrayRef, + validity: Option, +} + +/// Build the list array from its parts and apply the projection expression. +fn build_list( + parts: ListParts, + nullability: Nullability, + expr: &Expression, +) -> VortexResult { + let validity = create_validity(parts.validity, nullability); + ListArray::try_new(parts.elements, parts.offsets, validity)? + .into_array() + .apply(expr) +} + +struct ElementsProjection { + reader: ListReader, + expr: Expression, + row_range: Range, + offsets: ArrayFuture, +} + +impl ElementsProjection { + /// Path A1: whole-chunk read with an all-true mask. The elements bound is the whole elements + /// buffer (`0..elements_row_count`) and `offsets[0] == 0` within a chunk, so we don't need to + /// read offsets to know the bound and don't need to rebase. Fires elements + validity in + /// parallel with the already-in-flight offsets — a single `try_join!` over all three children. + async fn project_whole_chunk(self) -> VortexResult { + let Self { + reader, + expr, + row_range, + offsets, + } = self; + let validity_row_count = usize::try_from(row_range.end - row_range.start)?; + let elements_row_count = reader.elements.row_count(); + let elements_fut = reader.elements.projection_evaluation( + &(0..elements_row_count), + &root(), + MaskFuture::new_true(usize::try_from(elements_row_count)?), + )?; + let validity_fut = fetch_validity( + reader.validity.as_ref(), + &row_range, + MaskFuture::new_true(validity_row_count), + )?; + let (offsets, elements, validity) = try_join!(offsets, elements_fut, validity_fut)?; + build_list( + ListParts { + elements, + offsets, + validity, + }, + reader.layout.dtype().nullability(), + &expr, + ) + } + + /// Path A2: partial range with an all-true mask. The elements bound is + /// `offsets[a]..offsets[b]`, so we await offsets before firing the elements read and rebase the + /// offsets to start at zero. + async fn project_full_range(self) -> VortexResult { + let Self { + reader, + expr, + row_range, + offsets, + } = self; + let offsets = offsets.await?; + let elements_range = calculate_elements_range(&offsets, &reader.session)?; + let rebased_offsets = rebase_offsets(offsets, elements_range.start)?; + let elements_len = elements_range.end - elements_range.start; + let validity_row_count = usize::try_from(row_range.end - row_range.start)?; + + let elements_fut = reader.elements.projection_evaluation( + &elements_range, + &root(), + MaskFuture::new_true(usize::try_from(elements_len)?), + )?; + let validity_fut = fetch_validity( + reader.validity.as_ref(), + &row_range, + MaskFuture::new_true(validity_row_count), + )?; + let (elements, validity) = try_join!(elements_fut, validity_fut)?; + build_list( + ListParts { + elements, + offsets: rebased_offsets, + validity, + }, + reader.layout.dtype().nullability(), + &expr, + ) + } + + /// Path B: sparse mask. Bound the elements fetch to the tightest range covering the kept rows + /// and pass an element-level mask so the elements child only materializes kept-row positions; + /// validity is fetched for the kept rows by pushing the caller mask down directly. + async fn project_sparse(self, mask: Mask) -> VortexResult { + let Self { + reader, + expr, + row_range, + offsets, + } = self; + let validity_fut = fetch_validity( + reader.validity.as_ref(), + &row_range, + MaskFuture::ready(mask.clone()), + )?; + let offsets = offsets.await?; + let sg = compute_scatter_gather(&offsets, &mask, &reader.session)?; + let elements_fut = reader.elements.projection_evaluation( + &sg.elements_range, + &root(), + MaskFuture::ready(sg.element_mask), + )?; + let (elements, validity) = try_join!(elements_fut, validity_fut)?; + build_list( + ListParts { + elements, + offsets: sg.new_offsets, + validity, + }, + reader.layout.dtype().nullability(), + &expr, + ) } } @@ -448,6 +647,11 @@ mod tests { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::expr::eq; + use vortex_array::expr::is_not_null; + use vortex_array::expr::is_null; + use vortex_array::expr::lit; + use vortex_array::expr::not; use vortex_buffer::buffer; use super::*; @@ -460,6 +664,62 @@ mod tests { use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; + /// `classify` keys off the deepest list child an expression touches; `Elements` is the + /// always-correct default for anything not specifically recognized. + #[rstest] + // `is_null` / `is_not_null` of the list itself need only validity. + #[case::is_null(is_null(root()), ExprClass::Validity)] + #[case::is_not_null(is_not_null(root()), ExprClass::Validity)] + // Compound over validity-only operands stays validity. + #[case::not_is_null(not(is_null(root())), ExprClass::Validity)] + // A list-independent (constant) expression falls to the cheapest class. + #[case::constant(lit(5), ExprClass::Validity)] + // A bare list reference needs the elements. + #[case::bare_root(root(), ExprClass::Elements)] + // Any other fn over the list needs the elements. + #[case::not_root(not(root()), ExprClass::Elements)] + // `is_null` only short-circuits to validity when its argument is the list itself. + #[case::is_null_of_derived(is_null(not(root())), ExprClass::Elements)] + // Max over operands: validity + elements => elements. + #[case::validity_and_elements(eq(is_null(root()), root()), ExprClass::Elements)] + fn classify_expr_class(#[case] expr: Expression, #[case] expected: ExprClass) { + assert_eq!(classify(&expr), expected); + } + + /// Validity-class projections (`is_null` / `is_not_null` of the list) round-trip through the + /// validity-only read path, for both nullable and non-nullable lists. + #[rstest] + // `create_basic_list_array(true)` has validity `[true, false, true]`. + #[case::nullable(true, vec![true, false, true])] + #[case::non_nullable(false, vec![true, true, true])] + #[tokio::test] + async fn projection_validity_class( + #[case] nullable: bool, + #[case] valid: Vec, + ) -> VortexResult<()> { + let list = create_basic_list_array(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let not_null = reader + .projection_evaluation(&(0..3), &is_not_null(root()), MaskFuture::new_true(3))? + .await?; + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(not_null, BoolArray::from_iter(valid.clone()), &mut exec_ctx); + + let is_null_res = reader + .projection_evaluation(&(0..3), &is_null(root()), MaskFuture::new_true(3))? + .await?; + assert_arrays_eq!( + is_null_res, + BoolArray::from_iter(valid.iter().map(|v| !v).collect::>()), + &mut exec_ctx + ); + + Ok(()) + } + fn flat_list_strategy() -> ListLayoutStrategy { ListLayoutStrategy::default() } @@ -621,13 +881,27 @@ mod tests { // Keep only list 2, which has length 0 (offsets[2] == offsets[3] == 5). let mask = Mask::from_iter([false, false, true, false, false]); let (range, elem_mask, new_off, kept) = run_scatter_gather(five_list_offsets(), mask)?; - assert_eq!(range, 5..5); + assert_eq!(range, 0..0); assert!(elem_mask.is_empty()); assert_eq!(materialize_u32_array(new_off), vec![0, 0]); assert_eq!(kept, 1); Ok(()) } + #[test] + fn scatter_gather_ignores_empty_kept_boundary_rows() -> VortexResult<()> { + // The first and last kept rows are empty. The read range should be anchored to the one + // non-empty kept row, not widened across skipped rows. + let offsets = buffer![0u32, 0, 100, 102, 200, 200].into_array(); + let mask = Mask::from_iter([true, false, true, false, true]); + let (range, elem_mask, new_off, kept) = run_scatter_gather(offsets, mask)?; + assert_eq!(range, 100..102); + assert_eq!(elem_mask, vec![true, true]); + assert_eq!(materialize_u32_array(new_off), vec![0, 0, 2, 2]); + assert_eq!(kept, 3); + Ok(()) + } + #[test] fn scatter_gather_u64_offsets() -> VortexResult<()> { // Verify the ptype-dispatch path works for u64 offsets, not just u32. @@ -704,7 +978,8 @@ mod tests { let expected = list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)?; - assert_arrays_eq!(result, expected); + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); Ok(()) } @@ -721,7 +996,8 @@ mod tests { .await?; let expected = list.filter(mask)?; - assert_arrays_eq!(result, expected); + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); Ok(()) } @@ -769,7 +1045,8 @@ mod tests { .await?; let expected = list.filter(mask)?; - assert_arrays_eq!(result, expected); + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); Ok(()) } } From 3d431161cafa5eafc04ead08a74ba2f560865e02 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 23 Jun 2026 14:37:09 -0700 Subject: [PATCH 36/37] Fix list docs Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/writer.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 6de14d5e1dd..8d8998333bd 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -38,7 +38,7 @@ use crate::sequence::SequentialStreamExt; /// Strategy for writing list-typed arrays, with a fallback for non-list dtypes. /// /// Single-chunk only. For list-typed input the strategy: -/// 1. Canonicalizes the input chunk into a [`ListViewArray`]. +/// 1. Canonicalizes the input chunk into a [`ListView`]. /// 2. Calls [`list_from_list_view`] to rebuild it into zero-copy-to-list form /// (sorted, gapless, non-overlapping offsets) and produce a [`ListArray`]. /// 3. Writes the `elements`, `offsets`, and (when nullable) `validity` columns into @@ -51,7 +51,7 @@ use crate::sequence::SequentialStreamExt; /// # Chunking /// /// `ListLayoutStrategy` bails on empty or multi-chunk input, matching the convention used by -/// [`FlatLayoutStrategy`](crate::layouts::flat::writer::FlatLayoutStrategy). +/// [`FlatLayoutStrategy`]. /// /// [`ListArray`]: vortex_array::arrays::ListArray #[derive(Clone)] From fa581c42ee77452fb5d5cc02eaa0eb5b2d7f8837 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 23 Jun 2026 14:59:00 -0700 Subject: [PATCH 37/37] Add list filter evaluation Signed-off-by: Matt Katz --- vortex-layout/src/layouts/list/reader.rs | 117 ++++++++++++++++++++++- 1 file changed, 112 insertions(+), 5 deletions(-) diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 2b514862fe7..6c7b25c268d 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -47,6 +47,10 @@ use crate::segments::SegmentSource; type OptionalArrayFuture = BoxFuture<'static, VortexResult>>; +/// The threshold of mask density below which we push the input mask into projection evaluation, +/// and above which we evaluate the expression over all rows and intersect afterward. +const EXPR_EVAL_THRESHOLD: f64 = 0.2; + /// Reader for [`ListLayout`]. #[derive(Clone)] pub struct ListReader { @@ -231,7 +235,6 @@ fn create_validity(validity_array: Option, nullability: Nullability) - /// Drives "fetch as little as possible": a projection/filter that only inspects the list's /// null-ness needs the validity child; everything else needs the element values. The ordering /// `Validity < Elements` lets us take the max over the operands of a compound expression. -// TODO: have `filter_evaluation` use this too. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] enum ExprClass { /// Only the list's validity is needed (`is_null` / `is_not_null` of the list itself). @@ -462,11 +465,37 @@ impl LayoutReader for ListReader { fn filter_evaluation( &self, - _row_range: &Range, - _expr: &Expression, - _mask: MaskFuture, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, ) -> VortexResult { - todo!() + let len = mask.len(); + let reader = self.clone(); + let row_range = row_range.clone(); + let expr = expr.clone(); + let session = self.session.clone(); + + Ok(MaskFuture::new(len, async move { + let mask = mask.await?; + + if mask.all_false() { + return Ok(mask); + } + + if mask.density() < EXPR_EVAL_THRESHOLD { + let predicate = reader + .projection_evaluation(&row_range, &expr, MaskFuture::ready(mask.clone()))? + .await?; + let predicate_mask = predicate_array_to_mask(predicate, &session)?; + Ok(mask.intersect_by_rank(&predicate_mask)) + } else { + let predicate = reader + .projection_evaluation(&row_range, &expr, MaskFuture::new_true(len))? + .await?; + let predicate_mask = predicate_array_to_mask(predicate, &session)?; + Ok(mask & &predicate_mask) + } + })) } fn projection_evaluation( @@ -520,6 +549,11 @@ fn build_list( .apply(expr) } +fn predicate_array_to_mask(array: ArrayRef, session: &VortexSession) -> VortexResult { + let mut ctx = session.create_execution_ctx(); + array.null_as_false().execute(&mut ctx) +} + struct ElementsProjection { reader: ListReader, expr: Expression, @@ -720,6 +754,65 @@ mod tests { Ok(()) } + #[rstest] + #[case::is_not_null_nullable(true, is_not_null(root()), Mask::from_iter([true, false, true]))] + #[case::is_not_null_non_nullable(false, is_not_null(root()), Mask::new_true(3))] + #[case::is_null_nullable(true, is_null(root()), Mask::from_iter([false, true, false]))] + #[case::is_null_non_nullable(false, is_null(root()), Mask::new_false(3))] + #[tokio::test] + async fn filter_evaluation_validity_class( + #[case] nullable: bool, + #[case] expr: Expression, + #[case] expected: Mask, + ) -> VortexResult<()> { + let list = create_basic_list_array(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let result = reader + .filter_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? + .await?; + + assert_eq!(result, expected); + Ok(()) + } + + #[tokio::test] + async fn filter_evaluation_intersects_with_input_mask() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let input_mask = Mask::from_iter([true, true, false]); + let result = reader + .filter_evaluation(&(0..3), &is_not_null(root()), MaskFuture::ready(input_mask))? + .await?; + + assert_eq!(result, Mask::from_iter([true, false, false])); + Ok(()) + } + + #[tokio::test] + async fn filter_evaluation_sparse_mask_maps_by_rank() -> VortexResult<()> { + let list = create_six_list_array(); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let input_mask = Mask::from_iter([false, false, false, false, true, false]); + let result = reader + .filter_evaluation(&(0..6), &is_not_null(root()), MaskFuture::ready(input_mask))? + .await?; + + assert_eq!( + result, + Mask::from_iter([false, false, false, false, true, false]) + ); + Ok(()) + } + fn flat_list_strategy() -> ListLayoutStrategy { ListLayoutStrategy::default() } @@ -935,6 +1028,20 @@ mod tests { .into_array() } + fn create_six_list_array() -> ArrayRef { + let validity = Validity::Array( + BoolArray::from_iter([true, false, true, false, true, true]).into_array(), + ); + + ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 1, 2, 3, 4, 5, 6].into_array(), + validity, + ) + .expect("array is valid") + .into_array() + } + #[tokio::test] async fn fetch_offsets_includes_extra_endpoint() -> VortexResult<()> { let list = create_basic_list_array(false);