diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 980082b5802..7488827d9db 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -53,6 +53,7 @@ use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::repartition::RepartitionStrategy; use vortex_layout::layouts::repartition::RepartitionWriterOptions; use vortex_layout::layouts::table::TableStrategy; +use vortex_layout::layouts::table::use_experimental_list_layout; use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; use vortex_layout::layouts::zoned::writer::ZonedStrategy; #[cfg(feature = "unstable_encodings")] @@ -156,6 +157,10 @@ pub struct WriteStrategyBuilder { field_writers: HashMap>, allow_encodings: Option>, flat_strategy: Option>, + /// Whether to write list-like fields using structural list layouts. + /// + /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy + use_list_layout: bool, } impl Default for WriteStrategyBuilder { @@ -168,6 +173,7 @@ impl Default for WriteStrategyBuilder { field_writers: HashMap::new(), allow_encodings: Some(ALLOWED_ENCODINGS.clone()), flat_strategy: None, + use_list_layout: use_experimental_list_layout(), } } } @@ -182,6 +188,14 @@ impl WriteStrategyBuilder { self } + /// Enable writing list and fixed-size-list fields with structural list layouts. + /// + /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy + pub fn with_list_layout(mut self) -> Self { + self.use_list_layout = true; + self + } + /// Override the write layout for a specific field somewhere in the nested schema tree. /// /// The field path is matched after the root struct is split into columns. This is useful when a @@ -242,8 +256,9 @@ 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 leaf layout. + 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 @@ -321,8 +336,13 @@ impl WriteStrategyBuilder { let validity_strategy = CollectStrategy::new(compress_then_flat); // Take any field overrides from the builder and apply them to the final strategy. - let table_strategy = TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) - .with_field_writers(self.field_writers); + let mut table_strategy = + TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) + .with_field_writers(self.field_writers); + + if self.use_list_layout { + table_strategy = table_strategy.with_list_layout(); + } Arc::new(table_strategy) } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 7ceb7f21ba9..67a378c65ea 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -15,10 +15,12 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::Dict; +use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -1668,6 +1670,95 @@ async fn test_writer_with_complex_types() -> VortexResult<()> { Ok(()) } +/// Write `array` with list decomposition forced on (through the full compress/zone pipeline) and +/// read the whole thing back. +async fn write_read_roundtrip(array: ArrayRef) -> VortexResult { + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_list_layout() + .build(); + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, array.to_array_stream()) + .await?; + SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .read_all() + .await +} + +/// A `list>` column round-trips through the `TableStrategy` dispatcher, exercising list +/// decomposition recursing into itself (the outer list's `elements` are themselves lists). +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn nested_list_of_list_roundtrip() -> VortexResult<()> { + let inner = 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 outer = ListArray::try_new( + inner, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let st = StructArray::from_fields(&[("nested", outer)])?.into_array(); + + let result = write_read_roundtrip(st.clone()).await?; + assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +/// A `fixed_size_list>` column round-trips through the `TableStrategy` +/// dispatcher, exercising fixed-size-list decomposition recursing into itself. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn nested_fixed_size_list_of_fixed_size_list_roundtrip() -> VortexResult<()> { + let inner = FixedSizeListArray::new( + buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(), + 2, + Validity::NonNullable, + 4, + ) + .into_array(); + let outer = FixedSizeListArray::new(inner, 2, Validity::NonNullable, 2).into_array(); + let st = StructArray::from_fields(&[("nested", outer)])?.into_array(); + + let result = write_read_roundtrip(st.clone()).await?; + assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +/// A `struct<{ items: list>? }>` column round-trips, exercising list decomposition +/// recursing into struct decomposition (list `elements` are structs) plus a nullable list validity +/// child. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn nested_struct_list_struct_roundtrip() -> VortexResult<()> { + let inner_struct = StructArray::from_fields(&[ + ("a", buffer![1i32, 2, 3, 4, 5].into_array()), + ("b", buffer![10i32, 20, 30, 40, 50].into_array()), + ])? + .into_array(); + let items = ListArray::try_new( + inner_struct, + buffer![0u32, 2, 5, 5].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + )? + .into_array(); + let st = StructArray::from_fields(&[("items", items)])?.into_array(); + + let result = write_read_roundtrip(st.clone()).await?; + assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); + Ok(()) +} + #[tokio::test] async fn test_writer_with_statistics() -> VortexResult<()> { let array = StructArray::from_fields(&[("numbers", buffer![1u32, 2, 3, 4, 5].into_array())])? diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index a21b5605b31..336de3235b1 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod reader; +pub(crate) mod reader; pub mod writer; use std::sync::Arc; diff --git a/vortex-layout/src/layouts/fixed_size_list/expr.rs b/vortex-layout/src/layouts/fixed_size_list/expr.rs new file mode 100644 index 00000000000..d7b95466c23 --- /dev/null +++ b/vortex-layout/src/layouts/fixed_size_list/expr.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::list_length::ListLength; +use vortex_error::VortexResult; + +/// The minimal set of fixed-size-list children an expression needs for evaluation. +/// +/// For example: +/// - `is_null(root())` only needs validity. +/// - `list_length(root())` needs the fixed list size and validity, but not elements. +/// - `root()` needs elements and validity. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(super) enum FixedSizeListChildrenNeeded { + /// Only validity is needed (`is_null` / `is_not_null`). + Validity, + /// Only the fixed list size and validity are needed (`list_length`). + ListLengthAndValidity, + /// Elements and validity are needed. + Elements, +} + +/// The minimal set of fixed-size-list children needed to evaluate `expr`, where `root()` is a +/// field with fixed-size-list dtype. +pub(super) fn get_necessary_fixed_size_list_children( + expr: &Expression, +) -> FixedSizeListChildrenNeeded { + if is_null_root(expr) { + return FixedSizeListChildrenNeeded::Validity; + } + + if is_list_length_root(expr) { + return FixedSizeListChildrenNeeded::ListLengthAndValidity; + } + + if is_root(expr) { + return FixedSizeListChildrenNeeded::Elements; + } + + expr.children() + .iter() + .map(get_necessary_fixed_size_list_children) + .max() + .unwrap_or(FixedSizeListChildrenNeeded::Validity) +} + +fn is_null_root(expr: &Expression) -> bool { + (expr.is::() || expr.is::()) + && expr.children().len() == 1 + && is_root(expr.child(0)) +} + +fn is_list_length_root(expr: &Expression) -> bool { + expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) +} + +/// Rewrite a validity-class expression so it can be evaluated against the fixed-size-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. +pub(super) 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) +} + +/// Rewrite a list-length-class expression so it can be evaluated against an array of list lengths. +/// `list_length(root())` becomes `root()`. Other references to `root()` are left intact: for +/// list-length-class expressions they can only be validity checks, and the lengths array carries +/// the same validity as the original fixed-size-list. +pub(super) fn rewrite_list_length_expr(expr: &Expression) -> VortexResult { + if is_list_length_root(expr) { + return Ok(root()); + } + + let children = expr + .children() + .iter() + .map(rewrite_list_length_expr) + .collect::>>()?; + expr.clone().with_children(children) +} diff --git a/vortex-layout/src/layouts/fixed_size_list/mod.rs b/vortex-layout/src/layouts/fixed_size_list/mod.rs new file mode 100644 index 00000000000..657ec97d478 --- /dev/null +++ b/vortex-layout/src/layouts/fixed_size_list/mod.rs @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod expr; +mod reader; +pub mod writer; + +use std::sync::Arc; + +use reader::FixedSizeListReader; +use vortex_array::DeserializeMetadata; +use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +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::CachedId; + +use crate::LayoutBuildContext; +use crate::LayoutChildType; +use crate::LayoutEncodingRef; +use crate::LayoutId; +use crate::LayoutReaderContext; +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 `validity` layout, only present when the fixed-size list dtype is nullable. +pub const VALIDITY_CHILD_INDEX: usize = 1; + +vtable!(FixedSizeList); + +impl VTable for FixedSizeList { + type Layout = FixedSizeListLayout; + type Encoding = FixedSizeListLayoutEncoding; + type Metadata = EmptyMetadata; + + fn id(_encoding: &Self::Encoding) -> LayoutId { + static ID: CachedId = CachedId::new("vortex.fixed_size_list"); + *ID + } + + fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef { + LayoutEncodingRef::new_ref(FixedSizeListLayoutEncoding.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 { + EmptyMetadata + } + + fn segment_ids(_layout: &Self::Layout) -> Vec { + vec![] + } + + fn nchildren(layout: &Self::Layout) -> usize { + 1 + usize::from(layout.dtype.is_nullable()) + } + + fn child(layout: &Self::Layout, idx: usize) -> VortexResult { + match (idx, layout.validity.as_ref()) { + (ELEMENTS_CHILD_INDEX, _) => Ok(Arc::clone(&layout.elements)), + (VALIDITY_CHILD_INDEX, Some(validity)) => Ok(Arc::clone(validity)), + _ => vortex_bail!("Invalid child index {idx} for FixedSizeListLayout"), + } + } + + fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType { + match (idx, layout.validity.is_some()) { + (ELEMENTS_CHILD_INDEX, _) => LayoutChildType::Auxiliary("elements".into()), + (VALIDITY_CHILD_INDEX, true) => LayoutChildType::Auxiliary("validity".into()), + _ => vortex_panic!("Invalid child index {idx} for FixedSizeListLayout"), + } + } + + fn new_reader( + layout: &Self::Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + Ok(Arc::new(FixedSizeListReader::try_new( + layout.clone(), + name, + segment_source, + session.clone(), + ctx, + )?)) + } + + fn build( + _encoding: &Self::Encoding, + dtype: &DType, + row_count: u64, + _metadata: &::Output, + _segment_ids: Vec, + children: &dyn LayoutChildren, + _ctx: &LayoutBuildContext<'_>, + ) -> VortexResult { + validate_children(dtype, row_count, children)?; + + let element_dtype = dtype + .as_fixed_size_list_element_opt() + .ok_or_else(|| vortex_err!("FixedSizeListLayout requires a FixedSizeList dtype"))?; + let elements = children.child(ELEMENTS_CHILD_INDEX, element_dtype)?; + let validity = dtype + .is_nullable() + .then(|| children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))) + .transpose()?; + + Ok(FixedSizeListLayout { + row_count, + dtype: dtype.clone(), + elements, + validity, + }) + } + + fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { + validate_child_count(layout.dtype(), children.len())?; + + let mut iter = children.into_iter(); + layout.elements = iter + .next() + .ok_or_else(|| vortex_err!("missing elements child"))?; + layout.validity = layout + .dtype + .is_nullable() + .then(|| { + iter.next() + .ok_or_else(|| vortex_err!("missing validity child")) + }) + .transpose()?; + Ok(()) + } +} + +fn validate_child_count(dtype: &DType, nchildren: usize) -> VortexResult<()> { + let expected = 1 + usize::from(dtype.is_nullable()); + vortex_ensure!( + nchildren == expected, + "FixedSizeListLayout expects {expected} children, got {nchildren}" + ); + Ok(()) +} + +fn validate_children( + dtype: &DType, + row_count: u64, + children: &dyn LayoutChildren, +) -> VortexResult<()> { + validate_child_count(dtype, children.nchildren())?; + let DType::FixedSizeList(_, list_size, _) = dtype else { + vortex_bail!("FixedSizeListLayout requires a FixedSizeList dtype, got {dtype}"); + }; + let expected_elements = row_count + .checked_mul(u64::from(*list_size)) + .ok_or_else(|| vortex_err!("fixed-size list elements row count overflow"))?; + let actual_elements = children.child_row_count(ELEMENTS_CHILD_INDEX); + vortex_ensure!( + actual_elements == expected_elements, + "FixedSizeListLayout elements row count {actual_elements} does not match expected {expected_elements}" + ); + if dtype.is_nullable() { + let validity_rows = children.child_row_count(VALIDITY_CHILD_INDEX); + vortex_ensure!( + validity_rows == row_count, + "FixedSizeListLayout validity row count {validity_rows} does not match row count {row_count}" + ); + } + Ok(()) +} + +#[derive(Debug)] +pub struct FixedSizeListLayoutEncoding; + +/// Stores a fixed-size list by shredding elements and optional list validity into child layouts. +#[derive(Clone, Debug)] +pub struct FixedSizeListLayout { + row_count: u64, + dtype: DType, + elements: LayoutRef, + validity: Option, +} + +impl FixedSizeListLayout { + /// Construct a fixed-size-list layout from its components. + /// + /// # Invariants + /// + /// - `dtype` must be a [`DType::FixedSizeList`]. + /// - `elements.row_count() == row_count * list_size`. + /// - `validity` is present iff `dtype.is_nullable()`. + pub fn new( + row_count: u64, + dtype: DType, + elements: LayoutRef, + validity: Option, + ) -> Self { + Self { + row_count, + dtype, + elements, + validity, + } + } + + /// Number of fixed-size-list rows in this layout. + #[inline] + pub fn row_count(&self) -> u64 { + self.row_count + } + + #[inline] + pub fn elements(&self) -> &LayoutRef { + &self.elements + } + + #[inline] + pub fn validity(&self) -> Option<&LayoutRef> { + self.validity.as_ref() + } + + /// The fixed number of elements in each list row. + #[inline] + pub fn list_size(&self) -> u32 { + match &self.dtype { + DType::FixedSizeList(_, list_size, _) => *list_size, + _ => vortex_panic!("FixedSizeListLayout dtype must be FixedSizeList"), + } + } + + /// The dtype of the inner elements column. + pub fn elements_dtype(&self) -> &DType { + self.dtype + .as_fixed_size_list_element_opt() + .vortex_expect("FixedSizeListLayout dtype must be FixedSizeList") + } +} diff --git a/vortex-layout/src/layouts/fixed_size_list/reader.rs b/vortex-layout/src/layouts/fixed_size_list/reader.rs new file mode 100644 index 00000000000..12de80f2969 --- /dev/null +++ b/vortex-layout/src/layouts/fixed_size_list/reader.rs @@ -0,0 +1,738 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::try_join; +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::FixedSizeListArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::root; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +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::fixed_size_list::FixedSizeListLayout; +use crate::layouts::fixed_size_list::expr::FixedSizeListChildrenNeeded; +use crate::layouts::fixed_size_list::expr::get_necessary_fixed_size_list_children; +use crate::layouts::fixed_size_list::expr::rewrite_list_length_expr; +use crate::layouts::fixed_size_list::expr::rewrite_validity_expr; +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; + +#[derive(Clone)] +pub(super) struct FixedSizeListReader { + layout: FixedSizeListLayout, + name: Arc, + session: VortexSession, + elements: LayoutReaderRef, + validity: Option, +} + +impl FixedSizeListReader { + pub(super) fn try_new( + layout: FixedSizeListLayout, + 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 validity = layout + .validity() + .map(|v| { + v.new_reader( + format!("{name}.validity").into(), + Arc::clone(&segment_source), + &session, + ctx, + ) + }) + .transpose()?; + + Ok(Self { + layout, + name, + session, + elements, + validity, + }) + } + + 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(); + 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, + }; + + create_validity(validity_array, nullability) + .to_array(out_len) + .apply(&rewritten) + } + .boxed()) + } + + fn project_list_length( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let list_size = u64::from(self.layout.list_size()); + let nullability = self.layout.dtype().nullability(); + let row_count = usize::try_from(row_range.end - row_range.start)?; + let rewritten = rewrite_list_length_expr(expr)?; + let validity_fut = fetch_validity( + self.validity.as_ref(), + row_range, + MaskFuture::new_true(row_count), + )?; + + Ok(async move { + let validity = validity_fut.await?; + let lengths = ConstantArray::new(list_size, row_count) + .into_array() + .cast(DType::Primitive(PType::U64, nullability))?; + let lengths = apply_validity(lengths, validity, nullability)?; + + let mask = mask.await?; + let lengths = if mask.all_true() { + lengths + } else { + lengths.filter(mask)? + }; + lengths.apply(&rewritten) + } + .boxed()) + } + + fn project_elements( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let reader = self.clone(); + let expr = expr.clone(); + let row_range = row_range.clone(); + + Ok(async move { + let row_count = usize::try_from(row_range.end - row_range.start)?; + let list_size = u64::from(reader.layout.list_size()); + let elements_range = element_range(&row_range, list_size)?; + let elements_len = usize::try_from(elements_range.end - elements_range.start)?; + + let elements_fut = reader.elements.projection_evaluation( + &elements_range, + &root(), + MaskFuture::new_true(elements_len), + )?; + + let validity_fut = fetch_validity( + reader.validity.as_ref(), + &row_range, + MaskFuture::new_true(row_count), + )?; + + let (elements, validity) = try_join!(elements_fut, validity_fut)?; + let fsl = build_fixed_size_list(elements, validity, reader.layout.dtype(), row_count)?; + + let mask = mask.await?; + let fsl = if mask.all_true() { + fsl + } else { + fsl.filter(mask)? + }; + fsl.apply(&expr) + } + .boxed()) + } +} + +impl LayoutReader for FixedSizeListReader { + 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 RowSplits, + ) -> VortexResult<()> { + split_range.check_bounds(self.layout.row_count())?; + splits.push(split_range.root_row_range().end); + + let list_size = u64::from(self.layout.list_size()); + if list_size != 0 { + let element_range = element_range(split_range.row_range(), list_size)?; + + // The elements reader reports natural splits in element coordinates. Keep these + // temporary splits in element space, then convert them back to fixed-size-list row + // coordinates below, rounding forward so a split never cuts through a parent row. + let mut element_splits = RowSplits::new_capacity(8); + self.elements.register_splits( + field_mask, + &SplitRange::try_new(0, element_range.clone())?, + &mut element_splits, + )?; + + for element_split in element_splits.into_sorted_deduped() { + if element_split <= element_range.start || element_split >= element_range.end { + continue; + } + let element_offset = element_split - element_range.start; + let row_offset = element_offset.div_ceil(list_size); + let root_row = split_range + .root_row_range() + .start + .checked_add(row_offset) + .ok_or_else(|| vortex_err!("fixed-size-list split offset overflow"))?; + if root_row < split_range.root_row_range().end { + splits.push(root_row); + } + } + } + + if let Some(validity) = &self.validity { + validity.register_splits(field_mask, split_range, splits)?; + } + Ok(()) + } + + // TODO(mk): either have zone pruning upstream or implement here + 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 { + 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( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + match get_necessary_fixed_size_list_children(expr) { + FixedSizeListChildrenNeeded::Validity => self.project_validity(row_range, expr, mask), + FixedSizeListChildrenNeeded::ListLengthAndValidity => { + self.project_list_length(row_range, expr, mask) + } + FixedSizeListChildrenNeeded::Elements => self.project_elements(row_range, expr, mask), + } + } +} + +fn element_range(row_range: &Range, list_size: u64) -> VortexResult> { + let start = row_range + .start + .checked_mul(list_size) + .ok_or_else(|| vortex_err!("fixed-size-list element range overflow"))?; + let end = row_range + .end + .checked_mul(list_size) + .ok_or_else(|| vortex_err!("fixed-size-list element range overflow"))?; + Ok(start..end) +} + +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()) +} + +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, + }, + } +} + +fn apply_validity( + array: ArrayRef, + validity_array: Option, + nullability: Nullability, +) -> VortexResult { + if matches!(nullability, Nullability::Nullable) { + let len = array.len(); + array.mask(create_validity(validity_array, nullability).to_array(len)) + } else { + Ok(array) + } +} + +fn build_fixed_size_list( + elements: ArrayRef, + validity_array: Option, + dtype: &DType, + len: usize, +) -> VortexResult { + let DType::FixedSizeList(_, list_size, nullability) = dtype else { + return Err(vortex_err!( + "FixedSizeListLayout requires FixedSizeList dtype, got {dtype}" + )); + }; + let validity = create_validity(validity_array, *nullability); + Ok(FixedSizeListArray::try_new(elements, *list_size, validity, len)?.into_array()) +} + +fn predicate_array_to_mask(array: ArrayRef, session: &VortexSession) -> VortexResult { + let mut ctx = session.create_execution_ctx(); + array.null_as_false().execute(&mut ctx) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::FieldPath; + use vortex_array::expr::gt; + use vortex_array::expr::is_not_null; + use vortex_array::expr::is_null; + use vortex_array::expr::list_length; + use vortex_array::expr::lit; + use vortex_array::validity::Validity; + + use super::*; + use crate::LayoutStrategy; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::layouts::fixed_size_list::writer::FixedSizeListLayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::repartition::RepartitionStrategy; + use crate::layouts::repartition::RepartitionWriterOptions; + 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; + + async fn write_layout( + array: ArrayRef, + ) -> VortexResult<(Arc, crate::LayoutRef)> { + write_layout_with_strategy(&FixedSizeListLayoutStrategy::default(), array).await + } + + async fn write_layout_with_strategy( + strategy: &S, + array: ArrayRef, + ) -> VortexResult<(Arc, crate::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 repartitioned_chunk_strategy(block_len_multiple: usize) -> Arc { + Arc::new(RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple, + block_size_target: None, + canonicalize: true, + }, + )) + } + + fn create_fsl(nullable: bool) -> ArrayRef { + let validity = if nullable { + Validity::Array(BoolArray::from_iter([true, false, true, true]).into_array()) + } else { + Validity::NonNullable + }; + FixedSizeListArray::new( + PrimitiveArray::from_iter(0i32..8).into_array(), + 2, + validity, + 4, + ) + .into_array() + } + + #[rstest] + #[case::non_nullable(false)] + #[case::nullable(true)] + #[tokio::test] + async fn projection_full_range(#[case] nullable: bool) -> VortexResult<()> { + let fsl = create_fsl(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl.clone()).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let result = reader + .projection_evaluation(&(0..4), &root(), MaskFuture::new_true(4))? + .await?; + + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, fsl, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_partial_range() -> VortexResult<()> { + let fsl = create_fsl(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl.clone()).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let result = reader + .projection_evaluation(&(1..4), &root(), MaskFuture::new_true(3))? + .await?; + let expected = fsl.slice(1..4)?; + + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_sparse_mask() -> VortexResult<()> { + let fsl = create_fsl(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl.clone()).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + let mask = Mask::from_iter([true, false, false, true]); + + let result = reader + .projection_evaluation(&(0..4), &root(), MaskFuture::ready(mask.clone()))? + .await?; + let expected = fsl.filter(mask)?; + + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_degenerate_list_size_zero() -> VortexResult<()> { + let fsl = FixedSizeListArray::new( + PrimitiveArray::empty::(Nullability::NonNullable).into_array(), + 0, + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + 3, + ) + .into_array(); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl.clone()).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + let mask = Mask::from_iter([false, true, true]); + + let result = reader + .projection_evaluation(&(0..3), &root(), MaskFuture::ready(mask.clone()))? + .await?; + let expected = fsl.filter(mask)?; + + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn layout_splits_are_list_row_coordinates() -> VortexResult<()> { + let fsl = create_fsl(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let splits = SplitBy::Layout.splits( + reader.as_ref(), + &(0..4), + &[FieldMask::Exact(FieldPath::root())], + )?; + + assert_eq!(splits, vec![0, 4]); + Ok(()) + } + + #[tokio::test] + async fn layout_splits_round_chunked_element_boundaries_to_parent_rows() -> VortexResult<()> { + let fsl = FixedSizeListArray::new( + PrimitiveArray::from_iter(0i32..15).into_array(), + 3, + Validity::NonNullable, + 5, + ) + .into_array(); + let strategy = + FixedSizeListLayoutStrategy::default().with_elements(repartitioned_chunk_strategy(4)); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout_with_strategy(&strategy, fsl).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let splits = SplitBy::Layout.splits( + reader.as_ref(), + &(0..5), + &[FieldMask::Exact(FieldPath::root())], + )?; + + // Element boundaries 4, 8, and 12 map to parent rows ceil(n / 3): 2, 3, and 4. + assert_eq!(splits, vec![0, 2, 3, 4, 5]); + Ok(()) + } + + #[tokio::test] + async fn layout_splits_map_through_nested_fixed_size_lists() -> VortexResult<()> { + let inner = FixedSizeListArray::new( + PrimitiveArray::from_iter(0i32..24).into_array(), + 3, + Validity::NonNullable, + 8, + ) + .into_array(); + let outer = FixedSizeListArray::new(inner, 2, Validity::NonNullable, 4).into_array(); + + let inner_strategy = + FixedSizeListLayoutStrategy::default().with_elements(repartitioned_chunk_strategy(7)); + let outer_strategy = + FixedSizeListLayoutStrategy::default().with_elements(Arc::new(inner_strategy)); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout_with_strategy(&outer_strategy, outer).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let splits = SplitBy::Layout.splits( + reader.as_ref(), + &(0..4), + &[FieldMask::Exact(FieldPath::root())], + )?; + + // Primitive boundaries 7, 14, and 21 first map to inner rows 3, 5, and 7, + // then to outer rows 2, 3, and 4. The range end already supplies row 4. + assert_eq!(splits, vec![0, 2, 3, 4]); + Ok(()) + } + + #[tokio::test] + async fn layout_splits_preserve_chunked_validity_row_boundaries() -> VortexResult<()> { + let fsl = FixedSizeListArray::new( + PrimitiveArray::from_iter(0i32..15).into_array(), + 3, + Validity::Array(BoolArray::from_iter([true, false, true, true, false]).into_array()), + 5, + ) + .into_array(); + let strategy = + FixedSizeListLayoutStrategy::default().with_validity(repartitioned_chunk_strategy(2)); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout_with_strategy(&strategy, fsl).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let splits = SplitBy::Layout.splits( + reader.as_ref(), + &(0..5), + &[FieldMask::Exact(FieldPath::root())], + )?; + + assert_eq!(splits, vec![0, 2, 4, 5]); + Ok(()) + } + + #[rstest] + #[case::nullable(true, vec![true, false, true, true])] + #[case::non_nullable(false, vec![true, true, true, true])] + #[tokio::test] + async fn projection_validity_class( + #[case] nullable: bool, + #[case] valid: Vec, + ) -> VortexResult<()> { + let fsl = create_fsl(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let not_null = reader + .projection_evaluation(&(0..4), &is_not_null(root()), MaskFuture::new_true(4))? + .await?; + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(not_null, BoolArray::from_iter(valid.clone()), &mut exec_ctx); + + let null = reader + .projection_evaluation(&(0..4), &is_null(root()), MaskFuture::new_true(4))? + .await?; + assert_arrays_eq!( + null, + BoolArray::from_iter(valid.iter().map(|v| !v).collect::>()), + &mut exec_ctx + ); + Ok(()) + } + + #[tokio::test] + async fn projection_list_length_preserves_validity() -> VortexResult<()> { + let fsl = create_fsl(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let result = reader + .projection_evaluation(&(0..4), &list_length(root()), MaskFuture::new_true(4))? + .await?; + + let expected = + PrimitiveArray::from_option_iter::([Some(2), None, Some(2), Some(2)]) + .into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_list_length_applies_sparse_mask() -> VortexResult<()> { + let fsl = create_fsl(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + let mask = Mask::from_iter([false, true, false, true]); + + let result = reader + .projection_evaluation(&(0..4), &list_length(root()), MaskFuture::ready(mask))? + .await?; + + let expected = PrimitiveArray::from_option_iter::([None, Some(2)]).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn filter_evaluation_list_length() -> VortexResult<()> { + let fsl = create_fsl(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout) = write_layout(fsl).await?; + let reader = layout.new_reader("".into(), segments, &SESSION, &ctx)?; + + let result = reader + .filter_evaluation( + &(0..4), + >(list_length(root()), lit(1u64)), + MaskFuture::new_true(4), + )? + .await?; + + assert_eq!(result, Mask::from_iter([true, false, true, true])); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/fixed_size_list/writer.rs b/vortex-layout/src/layouts/fixed_size_list/writer.rs new file mode 100644 index 00000000000..d824d9cbc21 --- /dev/null +++ b/vortex-layout/src/layouts/fixed_size_list/writer.rs @@ -0,0 +1,396 @@ +// 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; +use futures::future::try_join_all; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::fixed_size_list::FixedSizeListDataParts; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::kanal_ext::KanalExt; +use vortex_io::session::RuntimeSessionExt; +use vortex_session::VortexSession; + +use crate::IntoLayout; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::fixed_size_list::FixedSizeListLayout; +use crate::layouts::flat::writer::FlatLayoutStrategy; +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; + +/// Item carried on each child sub-stream: a sequenced, materialized chunk. +type ChildChunk = VortexResult<(SequenceId, ArrayRef)>; + +/// Strategy for writing fixed-size-list arrays, with a fallback for other dtypes. +/// +/// For fixed-size-list input the strategy transposes the whole column stream into `elements` +/// and optional `validity` sub-streams. Each input chunk is canonicalized to a +/// [`FixedSizeListArray`], then its children are streamed concurrently to independently +/// configurable child strategies. +#[derive(Clone)] +pub struct FixedSizeListLayoutStrategy { + elements: Arc, + validity: Arc, + fallback: Arc, +} + +impl Default for FixedSizeListLayoutStrategy { + fn default() -> Self { + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + Self { + elements: Arc::clone(&flat), + validity: Arc::clone(&flat), + fallback: flat, + } + } +} + +impl FixedSizeListLayoutStrategy { + /// Strategy for the `elements` child. + pub fn with_elements(mut self, elements: Arc) -> Self { + self.elements = elements; + self + } + + /// Strategy for the `validity` child, written only when the list dtype is nullable. + pub fn with_validity(mut self, validity: Arc) -> Self { + self.validity = validity; + self + } + + /// Strategy for non-fixed-size-list input, which is forwarded unchanged. + pub fn with_fallback(mut self, fallback: Arc) -> Self { + self.fallback = fallback; + self + } +} + +#[async_trait] +impl LayoutStrategy for FixedSizeListLayoutStrategy { + 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_fixed_size_list() { + return self + .fallback + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + + let is_nullable = dtype.is_nullable(); + let element_dtype = dtype + .as_fixed_size_list_element_opt() + .vortex_expect("DType is FixedSizeList") + .as_ref() + .clone(); + + let (elements_tx, elements_rx) = kanal::bounded_async::(1); + let (validity_tx, validity_rx) = if is_nullable { + let (tx, rx) = kanal::bounded_async::(1); + (Some(tx), Some(rx)) + } else { + (None, None) + }; + + let fanout_fut = + transpose_fixed_size_list_column(stream, session.clone(), elements_tx, validity_tx); + + let handle = session.handle(); + let mut child_specs: Vec<( + DType, + Arc, + kanal::AsyncReceiver, + )> = vec![(element_dtype, Arc::clone(&self.elements), elements_rx)]; + if let Some(validity_rx) = validity_rx { + child_specs.push(( + DType::Bool(Nullability::NonNullable), + Arc::clone(&self.validity), + validity_rx, + )); + } + + let layout_futures: Vec<_> = child_specs + .into_iter() + .map(|(child_dtype, strategy, rx)| { + let child_stream = + SequentialStreamAdapter::new(child_dtype, rx.into_stream().boxed()).sendable(); + let child_eof = eof.split_off(); + let ctx = ctx.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); + strategy + .write_stream(ctx, segment_sink, child_stream, child_eof, &session) + .await + }) + }) + .collect(); + + let (row_count, layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; + let mut layouts = layouts.into_iter(); + let elements_layout = layouts.next().vortex_expect("elements layout present"); + let validity_layout = + is_nullable.then(|| layouts.next().vortex_expect("validity layout present")); + + Ok( + FixedSizeListLayout::new(row_count, dtype, elements_layout, validity_layout) + .into_layout(), + ) + } + + fn buffered_bytes(&self) -> u64 { + let fsl_bytes = self.elements.buffered_bytes() + self.validity.buffered_bytes(); + fsl_bytes.max(self.fallback.buffered_bytes()) + } +} + +/// Transpose a fixed-size-list column into `elements` and (when present) `validity` child +/// sub-streams. Errors surface to the caller, which joins this against the child writers, rather +/// than being hidden as an early channel close. +async fn transpose_fixed_size_list_column( + mut stream: SendableSequentialStream, + session: VortexSession, + elements_tx: kanal::AsyncSender, + validity_tx: Option>, +) -> VortexResult { + let mut exec_ctx = session.create_execution_ctx(); + let mut row_count = 0u64; + let mut saw_chunk = false; + + while let Some(chunk) = stream.next().await { + let (sequence_id, array) = chunk?; + saw_chunk = true; + let len = array.len(); + let FixedSizeListDataParts { + elements, validity, .. + } = canonicalize_to_fixed_size_list_parts(array, &mut exec_ctx)?; + row_count += u64::try_from(len)?; + + let mut sp = sequence_id.descend(); + if elements_tx + .send(Ok((sp.advance(), elements))) + .await + .is_err() + { + vortex_bail!("fixed-size-list elements writer finished before all chunks were sent"); + }; + + if let Some(validity_tx) = &validity_tx { + let validity = validity.execute_mask(len, &mut exec_ctx)?.into_array(); + if validity_tx + .send(Ok((sp.advance(), validity))) + .await + .is_err() + { + vortex_bail!( + "fixed-size-list validity writer finished before all chunks were sent" + ); + } + } + } + + if !saw_chunk { + vortex_bail!("FixedSizeListLayoutStrategy needs at least one chunk"); + } + + Ok(row_count) +} + +fn canonicalize_to_fixed_size_list_parts( + array: ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + Ok(array + .execute::(exec_ctx)? + .into_data_parts()) +} + +#[cfg(test)] +mod tests { + use futures::stream; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + + use super::*; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::segments::TestSegments; + use crate::sequence::SequentialArrayStreamExt; + use crate::test::SESSION; + + async fn write(strategy: &S, array: ArrayRef) -> VortexResult { + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .await + } + + async fn write_chunks( + strategy: &S, + dtype: DType, + chunks: Vec, + ) -> VortexResult { + let segments = Arc::new(TestSegments::default()); + let (mut ptr, eof) = SequenceId::root().split(); + let chunks = chunks + .into_iter() + .map(move |chunk| Ok((ptr.advance(), chunk))); + let stream = SequentialStreamAdapter::new(dtype, stream::iter(chunks).boxed()).sendable(); + strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .await + } + + fn i32_fsl_dtype(nullable: bool) -> DType { + DType::FixedSizeList( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + 2, + if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }, + ) + } + + fn create_fsl(validity: Validity) -> ArrayRef { + FixedSizeListArray::new(buffer![1i32, 2, 3, 4, 5, 6].into_array(), 2, validity, 3) + .into_array() + } + + #[tokio::test] + async fn basic_non_nullable_input() -> VortexResult<()> { + let layout = write( + &FixedSizeListLayoutStrategy::default(), + create_fsl(Validity::NonNullable), + ) + .await?; + + assert_eq!( + layout.display_tree().to_string(), + "vortex.fixed_size_list, dtype: fixed_size_list(i32)[2], children: 1\n\ + └── elements: vortex.flat, dtype: i32, segment: 0\n" + ); + Ok(()) + } + + #[tokio::test] + async fn basic_nullable_input() -> VortexResult<()> { + let layout = write( + &FixedSizeListLayoutStrategy::default(), + create_fsl(Validity::Array( + BoolArray::from_iter([true, false, true]).into_array(), + )), + ) + .await?; + + assert_eq!( + layout.display_tree().to_string(), + "vortex.fixed_size_list, dtype: fixed_size_list(i32)[2]?, children: 2\n\ + ├── elements: vortex.flat, dtype: i32, segment: 0\n\ + └── validity: vortex.flat, dtype: bool, segment: 1\n" + ); + Ok(()) + } + + #[tokio::test] + async fn chunked_input_with_chunked_child_strategies_succeeds() -> VortexResult<()> { + let chunk0 = FixedSizeListArray::new( + buffer![1i32, 2, 3, 4].into_array(), + 2, + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + 2, + ) + .into_array(); + let chunk1 = FixedSizeListArray::new( + buffer![5i32, 6].into_array(), + 2, + Validity::Array(BoolArray::from_iter([true]).into_array()), + 1, + ) + .into_array(); + + let child_strategy: Arc = + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())); + let strategy = FixedSizeListLayoutStrategy::default() + .with_elements(Arc::clone(&child_strategy)) + .with_validity(child_strategy); + + let layout = write_chunks(&strategy, i32_fsl_dtype(true), vec![chunk0, chunk1]).await?; + + assert_eq!(layout.row_count(), 3); + insta::assert_snapshot!(layout.display_tree(), @" + vortex.fixed_size_list, dtype: fixed_size_list(i32)[2]?, children: 2 + ├── elements: vortex.chunked, dtype: i32, children: 2 + │ ├── [0]: vortex.flat, dtype: i32, segment: 0 + │ └── [1]: vortex.flat, dtype: i32, segment: 1 + └── validity: vortex.chunked, dtype: bool, children: 2 + ├── [0]: vortex.flat, dtype: bool, segment: 2 + └── [1]: vortex.flat, dtype: bool, segment: 3 + "); + Ok(()) + } + + #[tokio::test] + async fn non_fixed_size_list_input_routes_to_fallback() -> VortexResult<()> { + let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let layout = write(&FixedSizeListLayoutStrategy::default(), primitive).await?; + assert_eq!( + layout.display_tree().to_string(), + "vortex.flat, dtype: i32, segment: 0\n" + ); + Ok(()) + } + + #[tokio::test] + async fn empty_stream_errors() { + let segments = Arc::new(TestSegments::default()); + let (_, eof) = SequenceId::root().split(); + let empty = stream::empty::>().boxed(); + let stream = SequentialStreamAdapter::new( + DType::FixedSizeList( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + 2, + Nullability::NonNullable, + ), + empty, + ) + .sendable(); + + let res = FixedSizeListLayoutStrategy::default() + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .await; + assert!(res.is_err()); + } +} diff --git a/vortex-layout/src/layouts/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs new file mode 100644 index 00000000000..152b7ea5dce --- /dev/null +++ b/vortex-layout/src/layouts/list/expr.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::list_length::ListLength; +use vortex_error::VortexResult; + +/// The minimal set of list children an expression needs for evaluation. +/// +/// For example: +/// - `is_null(root())` only needs the validity child. +/// - `list_length(root())` only needs the offsets and validity children. +/// - `root()` needs elements, offsets, and validity children. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(super) enum ListChildrenNeeded { + /// Only the validity child is needed (`is_null` / `is_not_null`). + Validity, + /// Only the offsets and validity children are needed (`list_length`). + OffsetsAndValidity, + /// All children are needed. + All, +} + +/// The minimal set of list children needed to evaluate `expr`, where `root()` is a field with list dtype. +pub(super) fn get_necessary_list_children(expr: &Expression) -> ListChildrenNeeded { + if is_null_root(expr) { + return ListChildrenNeeded::Validity; + } + + if is_list_length_root(expr) { + return ListChildrenNeeded::OffsetsAndValidity; + } + + if is_root(expr) { + return ListChildrenNeeded::All; + } + + // Otherwise the requirement is the max over the operands. Childless expressions that never + // touch the list, such as literals, fall back to the cheapest usable child. + expr.children() + .iter() + .map(get_necessary_list_children) + .max() + .unwrap_or(ListChildrenNeeded::Validity) +} + +fn is_null_root(expr: &Expression) -> bool { + (expr.is::() || expr.is::()) + && expr.children().len() == 1 + && is_root(expr.child(0)) +} + +fn is_list_length_root(expr: &Expression) -> bool { + expr.is::() && expr.children().len() == 1 && is_root(expr.child(0)) +} + +/// 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. +pub(super) 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) +} + +/// Rewrite an offsets-class expression so it can be evaluated against an array of list lengths. +/// `list_length(root())` becomes `root()`. Other references to `root()` are left intact: for +/// offsets-class expressions they can only be validity checks, and the lengths array carries the +/// same validity as the original list. +pub(super) fn rewrite_offsets_expr(expr: &Expression) -> VortexResult { + if is_list_length_root(expr) { + return Ok(root()); + } + + let children = expr + .children() + .iter() + .map(rewrite_offsets_expr) + .collect::>>()?; + expr.clone().with_children(children) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::cast; + use vortex_array::expr::eq; + use vortex_array::expr::gt; + use vortex_array::expr::is_not_null; + use vortex_array::expr::is_null; + use vortex_array::expr::list_length; + use vortex_array::expr::lit; + use vortex_array::expr::not; + use vortex_array::expr::root; + + use super::*; + + /// `get_necessary_list_children` keys off the deepest list child an expression touches; `All` + /// 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()), ListChildrenNeeded::Validity)] + #[case::is_not_null(is_not_null(root()), ListChildrenNeeded::Validity)] + // Compound over validity-only operands stays validity. + #[case::not_is_null(not(is_null(root())), ListChildrenNeeded::Validity)] + // A list-independent (constant) expression falls to the cheapest usable child. + #[case::constant(lit(5), ListChildrenNeeded::Validity)] + // `list_length(root())` needs offsets and validity, but not elements. + #[case::list_length(list_length(root()), ListChildrenNeeded::OffsetsAndValidity)] + // Compound over offsets-only operands stays offsets. + #[case::list_length_filter( + gt(list_length(root()), lit(1u64)), + ListChildrenNeeded::OffsetsAndValidity + )] + #[case::cast_list_length( + cast( + list_length(root()), + DType::Primitive(PType::I64, Nullability::Nullable), + ), + ListChildrenNeeded::OffsetsAndValidity + )] + // A bare list reference needs the elements. + #[case::bare_root(root(), ListChildrenNeeded::All)] + // Any other fn over the list needs the elements. + #[case::not_root(not(root()), ListChildrenNeeded::All)] + // `is_null` only short-circuits to validity when its argument is the list itself. + #[case::is_null_of_derived(is_null(not(root())), ListChildrenNeeded::All)] + // Max over operands: validity + elements => elements. + #[case::validity_and_elements(eq(is_null(root()), root()), ListChildrenNeeded::All)] + fn classify_expr_class(#[case] expr: Expression, #[case] expected: ListChildrenNeeded) { + assert_eq!(get_necessary_list_children(&expr), expected); + } +} diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs new file mode 100644 index 00000000000..2d020ab231e --- /dev/null +++ b/vortex-layout/src/layouts/list/mod.rs @@ -0,0 +1,271 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod expr; +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_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::LayoutBuildContext; +use crate::LayoutChildType; +use crate::LayoutEncodingRef; +use crate::LayoutId; +use crate::LayoutReaderContext; +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; + +vtable!(List); + +impl VTable for List { + type Layout = ListLayout; + type Encoding = ListLayoutEncoding; + type Metadata = ProstMetadata; + + fn id(_encoding: &Self::Encoding) -> LayoutId { + static ID: CachedId = CachedId::new("vortex.list"); + *ID + } + + 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_NON_NULLABLE + 1 + } 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::Auxiliary("validity".into()), + _ => vortex_panic!("Invalid child index {idx} for ListLayout"), + } + } + + fn new_reader( + layout: &Self::Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + Ok(Arc::new(ListReader::try_new( + layout.clone(), + name, + segment_source, + session.clone(), + ctx, + )?)) + } + + fn build( + _encoding: &Self::Encoding, + dtype: &DType, + _row_count: u64, + metadata: &::Output, + _segment_ids: Vec, + children: &dyn LayoutChildren, + _ctx: &LayoutBuildContext<'_>, + ) -> VortexResult { + validate_children(dtype, 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 = dtype + .is_nullable() + .then(|| children.child(VALIDITY_CHILD_INDEX, &DType::Bool(Nullability::NonNullable))) + .transpose()?; + + Ok(ListLayout { + dtype: dtype.clone(), + elements, + offsets, + validity, + }) + } + + fn with_children(layout: &mut Self::Layout, children: Vec) -> VortexResult<()> { + validate_children(layout.dtype(), 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 = layout + .dtype + .is_nullable() + .then(|| { + iter.next() + .ok_or_else(|| vortex_err!("missing validity child")) + }) + .transpose()?; + Ok(()) + } +} + +/// Validates expected number of children based on `dtype` +fn validate_children(dtype: &DType, n_children: usize) -> VortexResult<()> { + let expected = if dtype.is_nullable() { + NUM_CHILDREN_NON_NULLABLE + 1 + } else { + NUM_CHILDREN_NON_NULLABLE + }; + + vortex_ensure_eq!(n_children, expected); + Ok(()) +} + +#[derive(Debug)] +pub struct ListLayoutEncoding; + +/// Stores a list-typed array by shredding `elements`, `offsets`, and optional `validity` children. +#[derive(Clone, Debug)] +pub struct ListLayout { + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, +} + +impl ListLayout { + /// Construct a new `ListLayout` from its components. + /// + /// # 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)`. + pub fn new( + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, + ) -> Self { + Self { + dtype, + elements, + offsets, + validity, + } + } + + /// Number of lists in this layout. + #[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() + } + + /// The integer type used for the `offsets` child layout. + #[inline] + pub fn offsets_ptype(&self) -> PType { + self.offsets.dtype().as_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..0b9b1e27ac5 --- /dev/null +++ b/vortex-layout/src/layouts/list/reader.rs @@ -0,0 +1,1010 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::try_join; +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::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +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_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; +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::chunked::reader::ChunkedReader; +use crate::layouts::list::ListLayout; +use crate::layouts::list::expr::ListChildrenNeeded; +use crate::layouts::list::expr::get_necessary_list_children; +use crate::layouts::list::expr::rewrite_offsets_expr; +use crate::layouts::list::expr::rewrite_validity_expr; +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 { + 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, + 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() + .map(|v| { + v.new_reader( + format!("{name}.validity").into(), + Arc::clone(&segment_source), + &session, + ctx, + ) + }) + .transpose()?; + + Ok(Self { + layout, + name, + session, + elements, + offsets, + validity, + }) + } + + /// Projection for [`ListChildrenNeeded::Validity`] expressions. 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 [`ListChildrenNeeded::All`] expressions. Materializes the list and applies + /// the expression. + /// + /// Dispatches between a bounded read (a strict sub-range over chunked elements, fetching only + /// the element-chunks the range overlaps) and a whole-chunk read (full range, or a single flat + /// elements segment where there is nothing to skip). + fn project_all( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let is_full_range = row_range.start == 0 && row_range.end == self.layout.row_count(); + if is_full_range || !self.elements_are_chunked() { + self.project_all_concurrent(row_range, expr, mask) + } else { + self.project_all_elements_bounded(row_range, expr, mask) + } + } + + /// Whether the `elements` child is a chunked layout, so a bounded read can skip the element + /// chunks a row range does not overlap. For a single flat elements segment there is nothing to + /// skip, so the whole-chunk read (which avoids the offsets→elements round-trip) is preferred. + fn elements_are_chunked(&self) -> bool { + self.elements + .as_any() + .downcast_ref::() + .is_some() + } + + /// Whole-chunk read: fetches the entire `elements`, `offsets`, and `validity` children + /// concurrently — no offsets→elements round-trip, since the elements bound is the whole buffer. + /// The materialized list is sliced to `row_range` and filtered by the caller mask. + fn project_all_concurrent( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_count = self.layout.row_count(); + let elements_row_count = self.elements.row_count(); + let nullability = self.layout.dtype().nullability(); + let expr = expr.clone(); + let row_range = row_range.clone(); + + // Fire all three child reads up front so they run concurrently and overlap the mask await. + let offsets_fut = self.fetch_raw_offsets(&(0..row_count))?; + let elements_fut = self.fetch_raw_elements(&(0..elements_row_count))?; + let validity_fut = fetch_validity( + self.validity.as_ref(), + &(0..row_count), + MaskFuture::new_true(usize::try_from(row_count)?), + )?; + + Ok(async move { + let (offsets, elements, validity) = try_join!(offsets_fut, elements_fut, validity_fut)?; + let list = + ListArray::try_new(elements, offsets, create_validity(validity, nullability))? + .into_array(); + + // Slice the whole-chunk list down to the requested row range. + let list = if row_range.start == 0 && row_range.end == row_count { + list + } else { + list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)? + }; + + // Filter before applying the expression: the expression may depend on the filtered + // rows being removed (e.g. `cast(a, u8) where a < 256`). + let mask = mask.await?; + let list = if mask.all_true() { + list + } else { + list.filter(mask)? + }; + list.apply(&expr) + } + .boxed()) + } + + /// Bounded read for a strict sub-range over chunked elements. Reads `offsets[row_range]`, + /// decodes the first and last offset to bound the elements read to `[first..last)`, then + /// rebases the offsets to index into that sliced buffer. With chunked elements this fetches + /// only the element-chunks the range overlaps, at the cost of one offsets→elements round-trip. + fn project_all_elements_bounded( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let nullability = self.layout.dtype().nullability(); + let expr = expr.clone(); + let row_count = usize::try_from(row_range.end - row_range.start)?; + let reader = self.clone(); + + let offsets_fut = self.fetch_raw_offsets(row_range)?; + let validity_fut = fetch_validity( + self.validity.as_ref(), + row_range, + MaskFuture::new_true(row_count), + )?; + + Ok(async move { + let (offsets, validity) = try_join!(offsets_fut, validity_fut)?; + let elements_range = elements_range_from_offsets(&offsets, &reader.session)?; + + // Read only the elements this range covers. + let elements = reader.fetch_raw_elements(&elements_range)?.await?; + + // Rebase the offsets to index into the sliced elements buffer. + let offsets = rebase_offsets(offsets, elements_range.start)?; + let list = + ListArray::try_new(elements, offsets, create_validity(validity, nullability))? + .into_array(); + + // Filter before applying the expression (see `project_elements_whole_chunk`). + let mask = mask.await?; + let list = if mask.all_true() { + list + } else { + list.filter(mask)? + }; + list.apply(&expr) + } + .boxed()) + } + + /// Projection for [`ListChildrenNeeded::OffsetsAndValidity`] expressions. Only reads offsets and validity children. + fn project_offsets_validity( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let offsets = self.fetch_raw_offsets(row_range)?; + let reader = self.clone(); + let row_range = row_range.clone(); + let rewritten = rewrite_offsets_expr(expr)?; + + Ok(async move { + let mask = mask.await?; + let row_count = usize::try_from(row_range.end - row_range.start)?; + let nullability = reader.layout.dtype().nullability(); + + let validity_mask = if mask.all_true() { + MaskFuture::new_true(row_count) + } else { + MaskFuture::ready(mask.clone()) + }; + let validity_fut = fetch_validity(reader.validity.as_ref(), &row_range, validity_mask)?; + + let offsets = offsets.await?; + let lengths = list_lengths_from_offsets(offsets)?; + let lengths = if mask.all_true() { + lengths + } else { + lengths.filter(mask)? + }; + let validity = validity_fut.await?; + let lengths = apply_lengths_validity(lengths, validity, nullability)?; + + lengths.apply(&rewritten) + } + .boxed()) + } + + /// Fire the offsets read for `row_range` in list row space. The offsets child has an extra entry, so reading + /// `row_range` maps to offsets in `[row_range.start..row_range.end + 1)`. + /// + /// No mask or expression is applied. + fn fetch_raw_offsets(&self, row_range: &Range) -> VortexResult { + 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), + ) + } + + /// Fire the elements read for `row_range` in element space. + /// + /// No mask or expression is applied. + fn fetch_raw_elements(&self, row_range: &Range) -> VortexResult { + let row_count = usize::try_from(row_range.end - row_range.start)?; + self.elements + .projection_evaluation(row_range, &root(), MaskFuture::new_true(row_count)) + } +} + +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 + } + + 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 RowSplits, + ) -> VortexResult<()> { + self.offsets + .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( + &self, + row_range: &Range, + expr: &Expression, + mask: Mask, + ) -> VortexResult { + // Only validity-class predicates (`is_null` / `is_not_null` of the list) can be pruned via + // a child zone map: they rewrite to a predicate over the validity bool child, which prunes + // against its own zones. `list_length` is *not* prunable from the offsets child (its zones + // cover offset values, not per-row lengths), and element-value predicates don't map to + // row-space zones — both pass through unpruned. + if get_necessary_list_children(expr) != ListChildrenNeeded::Validity { + return Ok(MaskFuture::ready(mask)); + } + let Some(validity) = self.validity.as_ref() else { + // Non-nullable list: no nulls, so nothing to prune on. + return Ok(MaskFuture::ready(mask)); + }; + let rewritten = rewrite_validity_expr(expr)?; + validity.pruning_evaluation(row_range, &rewritten, mask) + } + + fn filter_evaluation( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + 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( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + // Read as little as possible based on which list children the expression needs. + match get_necessary_list_children(expr) { + ListChildrenNeeded::Validity => self.project_validity(row_range, expr, mask), + ListChildrenNeeded::OffsetsAndValidity => { + self.project_offsets_validity(row_range, expr, mask) + } + ListChildrenNeeded::All => self.project_all(row_range, expr, mask), + } + } +} + +/// 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()) +} + +/// Read `offsets[0]` and `offsets[-1]` and return the elements range they bound. +fn elements_range_from_offsets( + 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 they index into a sliced `elements[first..]` buffer that +/// starts at zero. +fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { + if first == 0 { + return Ok(offsets); + } + let constant = ConstantArray::new(first, offsets.len()) + .into_array() + .cast(offsets.dtype().clone())?; + offsets.binary(constant, Operator::Sub) +} + +/// Compute `offsets[i + 1] - offsets[i]` as the unmasked list length values. +fn list_lengths_from_offsets(offsets: ArrayRef) -> VortexResult { + let len = offsets.len().saturating_sub(1); + offsets + .slice(1..offsets.len())? + .binary(offsets.slice(0..len)?, Operator::Sub) +} + +fn apply_lengths_validity( + lengths: ArrayRef, + validity: Option, + nullability: Nullability, +) -> VortexResult { + let len = lengths.len(); + let lengths = lengths.cast(DType::Primitive(PType::U64, nullability))?; + + if matches!(nullability, Nullability::Nullable) { + lengths.mask(create_validity(validity, nullability).to_array(len)) + } else { + Ok(lengths) + } +} + +fn predicate_array_to_mask(array: ArrayRef, session: &VortexSession) -> VortexResult { + let mut ctx = session.create_execution_ctx(); + array.null_as_false().execute(&mut ctx) +} + +#[cfg(test)] +mod tests { + use std::ops::Range; + + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::expr::cast; + use vortex_array::expr::gt; + use vortex_array::expr::is_not_null; + use vortex_array::expr::is_null; + use vortex_array::expr::list_length; + use vortex_array::expr::lit; + use vortex_buffer::buffer; + use vortex_io::session::RuntimeSession; + use vortex_io::session::RuntimeSessionExt; + + 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::layouts::repartition::RepartitionStrategy; + use crate::layouts::repartition::RepartitionWriterOptions; + use crate::segments::SegmentSource; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialArrayStreamExt; + use crate::session::LayoutSession; + use crate::test::SESSION; + + /// 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, session) = 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(()) + } + + #[tokio::test] + async fn projection_list_length_reads_offsets() -> VortexResult<()> { + let list = create_basic_list_array(false); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::new_true(3))? + .await?; + + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, buffer![2u64, 2, 1].into_array(), &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_list_length_preserves_validity() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::new_true(3))? + .await?; + + let expected = + PrimitiveArray::from_option_iter::([Some(2), None, Some(1)]).into_array(); + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_list_length_applies_sparse_mask() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let mask = Mask::from_iter([false, true, true]); + let result = reader + .projection_evaluation(&(0..3), &list_length(root()), MaskFuture::ready(mask))? + .await?; + + let expected = PrimitiveArray::from_option_iter::([None, Some(1)]).into_array(); + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_cast_list_length() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let expr = cast( + list_length(root()), + DType::Primitive(PType::I64, Nullability::Nullable), + ); + let result = reader + .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? + .await?; + + let expected = + PrimitiveArray::from_option_iter::([Some(2), None, Some(1)]).into_array(); + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn filter_evaluation_list_length() -> VortexResult<()> { + let list = create_basic_list_array(true); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .filter_evaluation( + &(0..3), + >(list_length(root()), lit(1u64)), + MaskFuture::new_true(3), + )? + .await?; + + assert_eq!(result, Mask::from_iter([true, false, false])); + 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, session) = 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, session) = 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, session) = 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() + } + + fn layout_test_session() -> VortexSession { + vortex_array::array_session() + .with::() + .with::() + } + + async fn write_layout( + strategy: &S, + array: ArrayRef, + ) -> VortexResult<(Arc, LayoutRef, VortexSession)> { + let session = layout_test_session().with_tokio(); + 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, session)) + } + + fn materialize_u64_array(array: ArrayRef) -> Vec { + let mut ctx = SESSION.create_execution_ctx(); + array + .execute::(&mut ctx) + .unwrap() + .as_slice::() + .to_vec() + } + + 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() + } + + 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); + + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let ctx = LayoutReaderContext::new(); + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + let reader = reader + .as_any() + .downcast_ref::() + .expect("ListReader"); + + let offsets = reader.fetch_raw_offsets(&(1..3))?.await?; + assert_eq!(materialize_u64_array(offsets), vec![2u64, 4, 5]); + + Ok(()) + } + + #[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)] + #[tokio::test] + async fn projection_evaluation_round_trips( + #[case] row_range: Range, + #[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, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + 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)?)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + #[tokio::test] + async fn projection_evaluation_applies_mask() -> VortexResult<()> { + let list = create_basic_list_array(false); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + 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)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + /// Build a list with 5 rows and lengths [2, 3, 0, 3, 2]. + 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() + } + + /// Sparse row masks over a whole chunk round-trip: the whole-chunk read is filtered by the + /// mask in memory. + #[rstest] + #[case::single_middle(Mask::from_iter([false, false, false, true, false]), false)] + #[case::two_far_apart(Mask::from_iter([true, false, false, true, false]), false)] + #[case::boundaries(Mask::from_iter([true, false, false, false, true]), false)] + #[case::kept_empty_row(Mask::from_iter([false, false, true, false, false]), false)] + #[case::sparse_nullable(Mask::from_iter([true, false, true, false, true]), true)] + #[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 ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(0..5), &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let expected = list.filter(mask)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + /// A partial (sub-chunk) row range still returns exactly that range: the whole chunk is read, + /// then sliced. + #[tokio::test] + async fn projection_evaluation_partial_range_slices_whole_chunk() -> VortexResult<()> { + let list = create_wider_list_array(false); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&(1..4), &root(), MaskFuture::new_true(3))? + .await?; + + let expected = list.slice(1..4)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } + + /// A list strategy whose `elements` child is repartitioned into two-element chunks, so the + /// reader takes the bounded (chunk-skipping) path for strict sub-ranges. Offsets stay flat. + fn chunked_elements_list_strategy() -> ListLayoutStrategy { + let chunked_elements: Arc = Arc::new(RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: 2, + block_size_target: None, + canonicalize: true, + }, + )); + ListLayoutStrategy::default().with_elements(chunked_elements) + } + + /// The chunked-elements strategy must actually produce a chunked `elements` layout, otherwise + /// the reader would silently take the whole-chunk path and the bounded read would be untested. + #[tokio::test] + async fn chunked_elements_produces_chunked_layout() -> VortexResult<()> { + let list = create_wider_list_array(false); + let (_segments, layout, _session) = + write_layout(&chunked_elements_list_strategy(), list).await?; + let tree = layout.display_tree().to_string(); + assert!( + tree.contains("elements: vortex.chunked"), + "elements should be chunked:\n{tree}" + ); + Ok(()) + } + + /// With chunked elements, sub-range projections take the bounded read path. Every + /// range/mask/nullability combination must match the same projection over the ground-truth + /// array (`list.slice(range).filter(mask)`). + #[rstest] + #[case::full_all_true(0..5, Mask::new_true(5), false)] + #[case::subrange_all_true(1..4, Mask::new_true(3), false)] + #[case::subrange_sparse(1..4, Mask::from_iter([true, false, true]), false)] + #[case::partial_start(0..2, Mask::new_true(2), false)] + #[case::partial_end(2..5, Mask::new_true(3), false)] + #[case::single_non_empty(0..1, Mask::new_true(1), false)] + #[case::single_empty_row(2..3, Mask::from_iter([true]), false)] + #[case::empty_range(2..2, Mask::new_true(0), false)] + #[case::subrange_all_false(1..4, Mask::new_false(3), false)] + #[case::subrange_sparse_nullable(1..4, Mask::from_iter([true, false, true]), true)] + #[case::partial_end_nullable(2..5, Mask::new_true(3), true)] + #[tokio::test] + async fn chunked_elements_round_trips( + #[case] row_range: Range, + #[case] mask: Mask, + #[case] nullable: bool, + ) -> VortexResult<()> { + let list = create_wider_list_array(nullable); + let ctx = LayoutReaderContext::new(); + let (segments, layout, session) = + write_layout(&chunked_elements_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let result = reader + .projection_evaluation(&row_range, &root(), MaskFuture::ready(mask.clone()))? + .await?; + + let sliced = + list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)?; + let expected = sliced.filter(mask)?; + let mut exec_ctx = session.create_execution_ctx(); + assert_arrays_eq!(result, expected, &mut exec_ctx); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs new file mode 100644 index 00000000000..9f33c1eee3e --- /dev/null +++ b/vortex-layout/src/layouts/list/writer.rs @@ -0,0 +1,578 @@ +// 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; +use futures::future::try_join_all; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListView; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::list::ListDataParts; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::matcher::Matcher; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::kanal_ext::KanalExt; +use vortex_io::session::RuntimeSessionExt; +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; +use crate::sequence::SequenceId; +use crate::sequence::SequencePointer; +use crate::sequence::SequentialStream; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; + +/// Item carried on each child sub-stream: a sequenced, materialized chunk. +type ChildChunk = VortexResult<(SequenceId, ArrayRef)>; + +/// Strategy for writing list-typed arrays, with a fallback for non-list dtypes. +/// +/// This is a *structural* writer that decomposes a list column into independent `elements`, +/// `offsets`, and (when nullable) `validity` sub-columns, each written through its own downstream +/// strategy, producing a single [`ListLayout`]. +/// +/// For list-typed input the strategy transposes the whole column stream into three sub-streams: +/// 1. Each chunk is canonicalized to a [`ListArray`] (rebuilding a [`ListView`] via +/// [`list_from_list_view`] when necessary). +/// 2. `offsets` are rebased to global `u64` positions (cumulative across chunks) so the single +/// `offsets` child indexes into the concatenated `elements` child. +/// 3. `elements`, `offsets`, and `validity` are streamed to their child strategies concurrently. +/// +/// For input whose dtype is not [`DType::List`], the stream is forwarded unchanged to the +/// configured `fallback` strategy. +/// +/// [`ListArray`]: vortex_array::arrays::ListArray +#[derive(Clone)] +pub struct ListLayoutStrategy { + elements: Arc, + offsets: Arc, + validity: Arc, + fallback: Arc, +} + +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: 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( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + mut eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + if !dtype.is_list() { + return self + .fallback + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + + let is_nullable = dtype.is_nullable(); + let element_dtype = dtype + .as_list_element_opt() + .vortex_expect("DType is List") + .as_ref() + .clone(); + // Global (whole-column) offsets are cumulative and may exceed the input offset width, + // so definsively widen. + let offsets_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); + + // One bounded sub-stream per child: elements, offsets, and (when nullable) validity. + let (elements_tx, elements_rx) = kanal::bounded_async::(1); + let (offsets_tx, offsets_rx) = kanal::bounded_async::(1); + let (validity_tx, validity_rx) = if is_nullable { + let (tx, rx) = kanal::bounded_async::(1); + (Some(tx), Some(rx)) + } else { + (None, None) + }; + + // Transpose the list column into its child sub-streams and rebase offsets to global + // positions. Kept joined with the child writers below so producer errors surface rather + // than being hidden as an early channel close. + let fanout_fut = transpose_list_column( + stream, + session.clone(), + elements_tx, + offsets_tx, + validity_tx, + ); + + // Spawn a writer per child sub-stream, concurrently. + let handle = session.handle(); + let mut child_specs: Vec<( + DType, + Arc, + kanal::AsyncReceiver, + )> = vec![ + (element_dtype, Arc::clone(&self.elements), elements_rx), + (offsets_dtype, Arc::clone(&self.offsets), offsets_rx), + ]; + if let Some(validity_rx) = validity_rx { + child_specs.push(( + DType::Bool(Nullability::NonNullable), + Arc::clone(&self.validity), + validity_rx, + )); + } + + let layout_futures: Vec<_> = child_specs + .into_iter() + .map(|(child_dtype, strategy, rx)| { + let child_stream = + SequentialStreamAdapter::new(child_dtype, rx.into_stream().boxed()).sendable(); + let child_eof = eof.split_off(); + let ctx = ctx.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); + strategy + .write_stream(ctx, segment_sink, child_stream, child_eof, &session) + .await + }) + }) + .collect(); + + let (_, layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; + let mut layouts = layouts.into_iter(); + let elements_layout = layouts.next().vortex_expect("elements layout present"); + let offsets_layout = layouts.next().vortex_expect("offsets layout present"); + let validity_layout = + is_nullable.then(|| layouts.next().vortex_expect("validity layout present")); + + Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) + } + + fn buffered_bytes(&self) -> u64 { + let list_bytes = self.elements.buffered_bytes() + + self.offsets.buffered_bytes() + + self.validity.buffered_bytes(); + list_bytes.max(self.fallback.buffered_bytes()) + } +} + +/// Transpose a list column into its `elements`, `offsets`, and (when present) `validity` child +/// sub-streams, rebasing each chunk's local `offsets` to global `u64` positions so the single +/// `offsets` child indexes into the concatenated `elements` child. +/// +/// `validity_tx` is `Some` exactly when the list is nullable. Errors surface to the caller, which +/// joins this against the child writers, rather than being hidden as an early channel close. +async fn transpose_list_column( + mut stream: SendableSequentialStream, + session: VortexSession, + elements_tx: kanal::AsyncSender, + offsets_tx: kanal::AsyncSender, + validity_tx: Option>, +) -> VortexResult<()> { + let mut exec_ctx = session.create_execution_ctx(); + let mut element_base: u64 = 0; + let mut first = true; + let mut saw_chunk = false; + while let Some(chunk) = stream.next().await { + let (sequence_id, array) = chunk?; + saw_chunk = true; + let mut sp = sequence_id.descend(); + let ListDataParts { + elements, + offsets, + validity, + .. + } = canonicalize_to_list_parts(array, &mut exec_ctx)?; + let n_elements = elements.len() as u64; + let row_count = offsets.len().saturating_sub(1); + let offsets = global_offsets(offsets, element_base, first, &mut exec_ctx)?; + element_base += n_elements; + first = false; + + if elements_tx + .send(Ok((sp.advance(), elements))) + .await + .is_err() + || offsets_tx.send(Ok((sp.advance(), offsets))).await.is_err() + { + vortex_bail!("list child writer finished before all chunks were sent"); + } + if let Some(validity_tx) = &validity_tx { + let validity = validity + .execute_mask(row_count, &mut exec_ctx)? + .into_array(); + if validity_tx + .send(Ok((sp.advance(), validity))) + .await + .is_err() + { + vortex_bail!("list validity writer finished before all chunks were sent"); + } + } + } + if !saw_chunk { + vortex_bail!("ListLayoutStrategy needs at least one chunk"); + } + Ok(()) +} + +/// Canonicalize a list-dtype array into [`ListDataParts`]. +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(), exec_ctx)?.into_data_parts()) + } else { + unreachable!("AnyList matcher guarantees List or ListView") + } +} + +/// Rebase a chunk's local `offsets` into global `u64` positions for the whole-column `offsets` +/// child. Each chunk's offsets are shifted by `element_base` (the number of elements already +/// emitted) so they index into the concatenated `elements`. The duplicated boundary offset is +/// dropped on every chunk after the first, so the concatenation of all chunks' contributions is a +/// single monotonic `[0, .., total_elements]` array of length `row_count + 1`. +fn global_offsets( + offsets: ArrayRef, + element_base: u64, + first: bool, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let widened = offsets.cast(DType::Primitive(PType::U64, Nullability::NonNullable))?; + let based = if element_base == 0 { + widened + } else { + let base = ConstantArray::new(element_base, widened.len()).into_array(); + widened.binary(base, Operator::Add)? + }; + let based = if first { + based + } else { + based.slice(1..based.len())? + }; + // Materialize so the child sub-stream carries a concrete array rather than a lazy expression. + Ok(based.execute::(exec_ctx)?.into_array()) +} + +/// Matcher for `Array` or `Array`. +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 futures::stream; + 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; + use vortex_buffer::buffer; + use vortex_io::session::RuntimeSession; + + 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::session::LayoutSession; + + fn layout_test_session() -> VortexSession { + vortex_array::array_session() + .with::() + .with::() + .with_tokio() + } + + fn flat_list_strategy() -> ListLayoutStrategy { + ListLayoutStrategy::default() + } + + async fn write(strategy: &S, array: ArrayRef) -> VortexResult { + let session = layout_test_session(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .await + } + + fn i32_list_dtype(nullable: bool) -> DType { + DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }, + ) + } + + 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 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: u64, segment: 1 + "); + Ok(()) + } + + #[tokio::test] + 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: u64, segment: 1 + └── validity: vortex.flat, dtype: bool, segment: 2 + "); + 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?; + 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()); + let (_, eof) = SequenceId::root().split(); + let empty = stream::empty::>().boxed(); + let stream = SequentialStreamAdapter::new(i32_list_dtype(false), empty).sendable(); + let session = layout_test_session(); + + let res = flat_list_strategy() + .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .await; + assert!(res.is_err()) + } + + #[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::default().with_elements(table_strategy); + + 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: u64, 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 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 + │ ├── elements: vortex.flat, dtype: i32, segment: 1 + │ └── offsets: vortex.flat, dtype: u64, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 0 + "); + Ok(()) + } + + #[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 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 + │ ├── elements: vortex.list, dtype: list(i32), children: 2 + │ │ ├── elements: vortex.flat, dtype: i32, segment: 2 + │ │ └── offsets: vortex.flat, dtype: u64, segment: 3 + │ └── offsets: vortex.flat, dtype: u64, segment: 1 + └── offsets: vortex.flat, dtype: u64, segment: 0 + "); + Ok(()) + } + + #[tokio::test] + 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(), + 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 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: u64, segment: 1 + └── [1]: vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.flat, dtype: i32, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 3 + "); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 18df5b8f347..83d171c9e42 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -14,8 +14,10 @@ pub mod collect; pub mod compressed; pub mod dict; pub mod file_stats; +pub mod fixed_size_list; 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/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 8e7256d44aa..f51b3ed77be 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -4,14 +4,18 @@ //! A configurable writer strategy for tabular data. //! //! [`TableStrategy`] is a *dispatcher*: it inspects the dtype of the stream it is handed and -//! routes struct columns to [`StructStrategy`] and everything else to the configured leaf -//! strategy. Because it hands *itself* (suitably descended) to [`StructStrategy`] as the strategy -//! for the struct's children, arbitrarily nested struct trees are written with no manual wiring. +//! routes struct columns to [`StructStrategy`], list columns to [`ListLayoutStrategy`], +//! fixed-size-list columns to [`FixedSizeListLayoutStrategy`], and everything else to the +//! configured leaf strategy. Because it hands *itself* (suitably descended) to those structural +//! writers as the strategy for their children, arbitrarily nested struct/list trees are written +//! with no manual wiring. //! //! The dispatcher also owns field-path overrides, letting callers force a specific leaf field — //! at any depth — onto a custom strategy. +use std::env; use std::sync::Arc; +use std::sync::LazyLock; use async_trait::async_trait; use vortex_array::ArrayContext; @@ -25,17 +29,38 @@ use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::layouts::fixed_size_list::writer::FixedSizeListLayoutStrategy; +use crate::layouts::list::writer::ListLayoutStrategy; use crate::layouts::struct_::StructStrategy; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; +/// Whether [`TableStrategy`] writes list-like fields using structural list layouts by default. +/// Disabled unless the environment variable `VORTEX_EXPERIMENTAL_LIST_LAYOUT` is set to `1`. +/// +/// [`FixedSizeListLayoutStrategy`]: crate::layouts::fixed_size_list::writer::FixedSizeListLayoutStrategy +/// [`ListLayoutStrategy`]: crate::layouts::list::writer::ListLayoutStrategy +pub fn use_experimental_list_layout() -> bool { + static USE_EXPERIMENTAL_LIST_LAYOUT: LazyLock = + LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_LIST_LAYOUT").is_ok_and(|v| v == "1")); + *USE_EXPERIMENTAL_LIST_LAYOUT +} + /// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the /// structural writer for its dtype. /// /// Dispatch rules, applied to the dtype of the stream handed to [`write_stream`]: /// - **struct** → [`StructStrategy`], with each field written by its override (if any) or by a /// descended copy of this dispatcher. +/// - **list** → [`ListLayoutStrategy`], with `elements` written by a descended copy of this +/// dispatcher (so nested structs/lists recurse) and `offsets`/`validity` by the leaf/validity +/// strategies. Gated: only when list decomposition is enabled via +/// [`with_list_layout`][Self::with_list_layout] (off by default); otherwise a list falls through +/// to the leaf strategy. +/// - **fixed-size-list** → [`FixedSizeListLayoutStrategy`], with `elements` written by a +/// descended copy of this dispatcher and `validity` by the validity strategy. Gated by the same +/// [`with_list_layout`][Self::with_list_layout] switch as list. /// - **anything else** → the leaf strategy. /// /// [`write_stream`]: LayoutStrategy::write_stream @@ -47,6 +72,11 @@ pub struct TableStrategy { validity: Arc, /// The writer for leaf fields, i.e. anything that is not a struct. leaf: Arc, + /// Whether to write list-like fields using structural list layouts. + /// + /// [`FixedSizeListLayoutStrategy`]: crate::layouts::fixed_size_list::writer::FixedSizeListLayoutStrategy + /// [`ListLayoutStrategy`]: crate::layouts::list::writer::ListLayoutStrategy + use_list_layout: bool, } impl TableStrategy { @@ -73,6 +103,7 @@ impl TableStrategy { leaf_writers: Default::default(), validity, leaf: fallback, + use_list_layout: false, } } @@ -139,6 +170,12 @@ impl TableStrategy { self.validity = validity; self } + + /// Enable writing list and fixed-size-list fields with structural list layouts. + pub fn with_list_layout(mut self) -> Self { + self.use_list_layout = true; + self + } } impl TableStrategy { @@ -174,6 +211,31 @@ impl TableStrategy { .with_field_writers(field_writers) } + /// Build the [`ListLayoutStrategy`] used to write a list field stream at this level. + /// + /// The `elements` sub-column is routed back through a clean descended dispatcher so nested + /// structs/lists recurse; `offsets` go straight to the leaf (they are always a primitive + /// column); and `validity` uses the shared validity strategy. + fn list_strategy(&self) -> ListLayoutStrategy { + ListLayoutStrategy::default() + .with_elements(Arc::new(self.descend_clean())) + .with_offsets(Arc::clone(&self.leaf)) + .with_validity(Arc::clone(&self.validity)) + .with_fallback(Arc::clone(&self.leaf)) + } + + /// Build the [`FixedSizeListLayoutStrategy`] used to write a fixed-size-list field stream at + /// this level. + /// + /// The `elements` sub-column is routed back through a clean descended dispatcher so nested + /// structs/lists/fixed-size-lists recurse; `validity` uses the shared validity strategy. + fn fixed_size_list_strategy(&self) -> FixedSizeListLayoutStrategy { + FixedSizeListLayoutStrategy::default() + .with_elements(Arc::new(self.descend_clean())) + .with_validity(Arc::clone(&self.validity)) + .with_fallback(Arc::clone(&self.leaf)) + } + /// Descend into a subfield, retaining only the overrides that apply beneath it (rebased to be /// relative to the child). fn descend(&self, field: &Field) -> Self { @@ -192,6 +254,7 @@ impl TableStrategy { leaf_writers: new_writers, validity: Arc::clone(&self.validity), leaf: Arc::clone(&self.leaf), + use_list_layout: self.use_list_layout, } } @@ -202,6 +265,7 @@ impl TableStrategy { leaf_writers: HashMap::default(), validity: Arc::clone(&self.validity), leaf: Arc::clone(&self.leaf), + use_list_layout: self.use_list_layout, } } @@ -244,6 +308,20 @@ impl LayoutStrategy for TableStrategy { .await; } + if dtype.is_list() && self.use_list_layout { + return self + .list_strategy() + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + + if dtype.is_fixed_size_list() && self.use_list_layout { + return self + .fixed_size_list_strategy() + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + // Leaf: hand off to the leaf strategy. self.leaf .write_stream(ctx, segment_sink, stream, eof, session) @@ -259,7 +337,11 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::array_session; + use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; @@ -268,10 +350,13 @@ mod tests { use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; use vortex_array::field_path; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; + use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; @@ -283,14 +368,22 @@ mod tests { use crate::sequence::SequentialArrayStreamExt; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; - use crate::test::SESSION; + use crate::session::LayoutSession; + + fn layout_test_session() -> VortexSession { + array_session() + .with::() + .with::() + .with_tokio() + } async fn write(strategy: &S, array: ArrayRef) -> VortexResult { + let session = layout_test_session(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let stream = array.to_array_stream().sequenced(ptr); strategy - .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .write_stream(ArrayContext::empty(), segments, stream, eof, &session) .await } @@ -322,6 +415,212 @@ mod tests { Ok(()) } + /// A `list>` column: the dispatcher recurses into itself so the outer list's + /// `elements` are decomposed as a nested `ListLayout`. + #[tokio::test] + async fn dispatches_nested_list() -> VortexResult<()> { + let inner = 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 outer = ListArray::try_new( + inner, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let layout = write(&flat_table().with_list_layout(), outer).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + 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: u64, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 0 + "); + Ok(()) + } + + /// A `struct<{ items: list>? }>` column: list decomposition recurses into struct + /// decomposition for the elements, and a nullable list writes a validity child. + #[tokio::test] + async fn dispatches_struct_list_struct() -> VortexResult<()> { + let inner_struct = 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 items = ListArray::try_new( + inner_struct, + buffer![0u32, 2, 5, 5].into_array(), + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + )? + .into_array(); + let st = StructArray::from_fields([("items", items)].as_slice())?.into_array(); + + let layout = write(&flat_table().with_list_layout(), st).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1 + └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3 + ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 + │ ├── a: vortex.flat, dtype: i32, segment: 2 + │ └── b: vortex.flat, dtype: i32, segment: 3 + ├── offsets: vortex.flat, dtype: u64, segment: 0 + └── validity: vortex.flat, dtype: bool, segment: 1 + "); + Ok(()) + } + + /// Fixed-size-list decomposition is gated by the same switch as list decomposition. Without + /// it, fixed-size-list columns fall through to the configured leaf strategy. + #[tokio::test] + async fn fixed_size_list_uses_leaf_without_list_layout() -> VortexResult<()> { + let fsl = FixedSizeListArray::new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + 2, + Validity::NonNullable, + 3, + ) + .into_array(); + + let layout = write(&flat_table(), fsl).await?; + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: fixed_size_list(i32)[2], segment: 0"); + Ok(()) + } + + /// A `fixed_size_list>` column: the dispatcher recurses into itself so + /// the outer fixed-size-list's `elements` are decomposed as a nested fixed-size-list layout. + #[tokio::test] + async fn dispatches_nested_fixed_size_list() -> VortexResult<()> { + let inner = FixedSizeListArray::new( + buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(), + 2, + Validity::NonNullable, + 4, + ) + .into_array(); + let outer = FixedSizeListArray::new(inner, 2, Validity::NonNullable, 2).into_array(); + + let layout = write(&flat_table().with_list_layout(), outer).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.fixed_size_list, dtype: fixed_size_list(fixed_size_list(i32)[2])[2], children: 1 + └── elements: vortex.fixed_size_list, dtype: fixed_size_list(i32)[2], children: 1 + └── elements: vortex.flat, dtype: i32, segment: 0 + "); + Ok(()) + } + + /// A `struct<{ items: fixed_size_list>? }>` column: fixed-size-list + /// decomposition recurses into struct decomposition for the elements, and a nullable + /// fixed-size-list writes a validity child. + #[tokio::test] + async fn dispatches_struct_fixed_size_list_struct() -> VortexResult<()> { + let inner_struct = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3, 4, 5, 6].into_array()), + ("b", buffer![10i32, 20, 30, 40, 50, 60].into_array()), + ] + .as_slice(), + )? + .into_array(); + let items = FixedSizeListArray::new( + inner_struct, + 2, + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + 3, + ) + .into_array(); + let st = StructArray::from_fields([("items", items)].as_slice())?.into_array(); + + let layout = write(&flat_table().with_list_layout(), st).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {items=fixed_size_list({a=i32, b=i32})[2]?}, children: 1 + └── items: vortex.fixed_size_list, dtype: fixed_size_list({a=i32, b=i32})[2]?, 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 + └── validity: vortex.flat, dtype: bool, segment: 0 + "); + Ok(()) + } + + /// A multi-chunk `list` written with a chunked leaf: each sub-column (`elements`, + /// `offsets`) becomes its own `ChunkedLayout`, so elements are chunked independently of rows. + /// This is the "list-of-chunkeds" topology top-level decomposition unlocks. + #[tokio::test] + async fn dispatches_chunked_list() -> VortexResult<()> { + 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 flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let dispatcher = TableStrategy::new( + Arc::clone(&flat), + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())), + ) + .with_list_layout(); + let layout = write(&dispatcher, chunked).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.chunked, dtype: i32, children: 2 + │ ├── [0]: vortex.flat, dtype: i32, segment: 0 + │ └── [1]: vortex.flat, dtype: i32, segment: 1 + └── offsets: vortex.chunked, dtype: u64, children: 2 + ├── [0]: vortex.flat, dtype: u64, segment: 2 + └── [1]: vortex.flat, dtype: u64, segment: 3 + "); + Ok(()) + } + + /// A multi-chunk `fixed_size_list` written with a chunked leaf becomes a single + /// fixed-size-list layout whose elements child is independently chunked. + #[tokio::test] + async fn dispatches_chunked_fixed_size_list() -> VortexResult<()> { + let chunk0 = FixedSizeListArray::new( + buffer![1i32, 2, 3, 4].into_array(), + 2, + Validity::NonNullable, + 2, + ) + .into_array(); + let chunk1 = + FixedSizeListArray::new(buffer![5i32, 6].into_array(), 2, Validity::NonNullable, 1) + .into_array(); + let dtype = chunk0.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array(); + + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let dispatcher = TableStrategy::new( + Arc::clone(&flat), + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())), + ) + .with_list_layout(); + let layout = write(&dispatcher, chunked).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.fixed_size_list, dtype: fixed_size_list(i32)[2], children: 1 + └── elements: vortex.chunked, dtype: i32, children: 2 + ├── [0]: vortex.flat, dtype: i32, segment: 0 + └── [1]: vortex.flat, dtype: i32, segment: 1 + "); + Ok(()) + } + /// A non-struct stream is not shredded; it is handed straight to the leaf strategy. #[tokio::test] async fn non_struct_input_uses_leaf() -> VortexResult<()> { @@ -448,7 +747,7 @@ mod tests { ); block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = layout_test_session().with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index be938bdc192..5b9769455d9 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -11,7 +11,9 @@ use vortex_session::registry::Registry; use crate::LayoutEncodingRef; use crate::layouts::chunked::ChunkedLayoutEncoding; use crate::layouts::dict::DictLayoutEncoding; +use crate::layouts::fixed_size_list::FixedSizeListLayoutEncoding; 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; @@ -49,6 +51,10 @@ impl Default for LayoutSession { // Register the built-in layout encodings. layouts.register(ChunkedLayoutEncoding.id(), ChunkedLayoutEncoding.as_ref()); + layouts.register( + FixedSizeListLayoutEncoding.id(), + FixedSizeListLayoutEncoding.as_ref(), + ); layouts.register(FlatLayoutEncoding.id(), FlatLayoutEncoding.as_ref()); layouts.register(StructLayoutEncoding.id(), StructLayoutEncoding.as_ref()); layouts.register(ZonedLayoutEncoding.id(), ZonedLayoutEncoding.as_ref()); @@ -57,6 +63,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 } }