diff --git a/Cargo.lock b/Cargo.lock index 459b58dded4..f63ea0bbd60 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10502,6 +10502,7 @@ dependencies = [ "bit-vec", "flatbuffers", "futures", + "insta", "itertools 0.14.0", "kanal", "moka", diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 980082b5802..769462db060 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -50,6 +50,7 @@ use vortex_layout::layouts::compressed::CompressingStrategy; use vortex_layout::layouts::compressed::CompressorPlugin; use vortex_layout::layouts::dict::writer::DictStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::layouts::list::writer::ListLayoutStrategy; use vortex_layout::layouts::repartition::RepartitionStrategy; use vortex_layout::layouts::repartition::RepartitionWriterOptions; use vortex_layout::layouts::table::TableStrategy; @@ -71,6 +72,14 @@ use vortex_zstd::ZstdBuffers; const ONE_MEG: u64 = 1 << 20; +/// Upper bound on the number of `elements` rows in a single element-chunk of a list column. +/// +/// The `elements` child of a list is repartitioned into chunks so a reader can fetch only the +/// element-chunks a row range overlaps. `block_size_target` keeps each chunk near 1 MiB for +/// fixed-width elements; this bounds the row count for narrow/variable-width elements where the +/// byte target alone would produce very large chunks. +const ELEMENTS_PER_CHUNK: usize = 1 << 20; + /// Static registry of all allowed array encodings for file writing. /// /// This includes all canonical encodings from vortex-array plus all compressed @@ -242,12 +251,8 @@ impl WriteStrategyBuilder { Arc::new(FlatLayoutStrategy::default()) }; - // 7. for each chunk create a flat layout - let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); - // 6. buffer chunks so they end up with closer segment ids physically - let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB - - // 5. compress each chunk. + // The data compressor is built up front: the list leaf (step 7) reuses it to compress each + // elements chunk, and the outer chain reuses it at step 5. // Exclude IntDictScheme from the data compressor because DictStrategy (step 3) already // dictionary-encodes columns. Allowing IntDictScheme here would redundantly // dictionary-encode the integer codes produced by that earlier step. @@ -260,7 +265,42 @@ impl WriteStrategyBuilder { ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; - let compressing = CompressingStrategy::new(buffered, data_compressor); + + // The `elements` child of a list is repartitioned into ~1 MiB chunks, each compressed + // independently, producing a `ChunkedLayout`. This lets the reader fetch only the + // element-chunks a row range overlaps (bounded read) instead of the whole elements buffer. + let chunked_elements: Arc = Arc::new(RepartitionStrategy::new( + CompressingStrategy::new( + ChunkedLayoutStrategy::new(Arc::clone(&flat)), + Arc::clone(&data_compressor), + ), + RepartitionWriterOptions { + block_size_minimum: ONE_MEG, + block_len_multiple: ELEMENTS_PER_CHUNK, + block_size_target: Some(ONE_MEG), + canonicalize: true, + }, + )); + + // 7. for each chunk create a layout. List-typed chunks route through + // `ListLayoutStrategy` (separately-addressable elements/offsets/validity sub-layouts; + // non-list chunks fall through its built-in fallback to `flat`). Offsets, validity, and the + // fallback thread the configured `flat` (which carries `allow_encodings` / any custom flat + // override); elements route through the chunked+compressed strategy above. + let leaf: Arc = Arc::new( + ListLayoutStrategy::default() + .with_elements(chunked_elements) + .with_offsets(Arc::clone(&flat)) + .with_validity(Arc::clone(&flat)) + .with_fallback(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 + + // 5. compress each chunk. + let compressing = CompressingStrategy::new(buffered, Arc::clone(&data_compressor)); // 4. prior to compression, coalesce up to a minimum size let coalescing = RepartitionStrategy::new( diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index b196f68dc96..be33caf7fcb 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2004,6 +2004,50 @@ async fn test_segment_ordering_dict_codes_before_values() -> VortexResult<()> { Ok(()) } +/// A list column whose per-chunk elements exceed the ~1 MiB element-chunk target is written with a +/// chunked `elements` layout, so a reader can fetch only the element-chunks a row range overlaps. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn list_layout_chunks_large_elements() -> VortexResult<()> { + // 20k lists of length 64 => ~1.28M i32 elements. After the 8k-row repartition each list-chunk + // still holds > 1 MiB of elements, so `elements` is written as a chunked layout (rather than a + // single flat segment, which `ChunkedLayoutStrategy` would collapse to for < 1 MiB). + let list_len = 64u32; + let n_rows = 20_000usize; + let n_elems = n_rows * list_len as usize; + let elements = PrimitiveArray::from_iter(0..n_elems as i32).into_array(); + let offsets = PrimitiveArray::from_iter((0..=n_rows as u32).map(|i| i * list_len)).into_array(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array(); + let st = StructArray::from_fields(&[("l", list)])?.into_array(); + + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .write(&mut buf, st.to_array_stream()) + .await?; + + fn any_list_has_chunked_elements(layout: &dyn Layout) -> bool { + if layout.encoding_id().as_ref() == "vortex.list" + && layout.child(0).unwrap().encoding_id().as_ref() == "vortex.chunked" + { + return true; + } + for child in layout.children().unwrap() { + if any_list_has_chunked_elements(child.as_ref()) { + return true; + } + } + false + } + + let root = summary.footer().layout(); + assert!( + any_list_has_chunked_elements(root.as_ref()), + "expected a list layout with chunked elements" + ); + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> { diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 61b1253ef43..daa96d4e138 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -55,6 +55,7 @@ vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] futures = { workspace = true, features = ["executor"] } +insta = { workspace = true } rstest = { workspace = true } temp-env = { workspace = true } tokio = { workspace = true, features = ["rt", "macros"] } diff --git a/vortex-layout/src/layouts/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/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs new file mode 100644 index 00000000000..d955b8a92af --- /dev/null +++ b/vortex-layout/src/layouts/list/expr.rs @@ -0,0 +1,151 @@ +// 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 deepest list child an expression needs, cheapest first, where I/O cost order is defined as +/// `Validity < OffsetsAndValidity < All`. +/// +/// 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, + /// The offsets and validity children are needed, but not the element values (`list_length`). + OffsetsAndValidity, + /// All children are needed. + All, +} + +/// The minimal 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..a0ffeeb7425 --- /dev/null +++ b/vortex-layout/src/layouts/list/reader.rs @@ -0,0 +1,1015 @@ +// 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 (`is_null` / `is_not_null` of + /// the list): reads only the validity child, synthesizing all-valid for a non-nullable list, + /// and never touches the offsets or elements. + fn project_validity( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let validity_reader = self.validity.clone(); + let nullability = self.layout.dtype().nullability(); + let row_range = row_range.clone(); + // Evaluate the rewritten expression against the validity bool array (true == valid row). + let rewritten = rewrite_validity_expr(expr)?; + + Ok(async move { + let mask = mask.await?; + let row_count = usize::try_from(row_range.end - row_range.start)?; + let out_len = if mask.all_true() { + row_count + } else { + mask.true_count() + }; + + let validity_array = match validity_reader.as_ref() { + Some(v) => Some( + v.projection_evaluation(&row_range, &root(), MaskFuture::ready(mask))? + .await?, + ), + None => None, + }; + + let validity = create_validity(validity_array, nullability).to_array(out_len); + validity.apply(&rewritten) + } + .boxed()) + } + + /// Projection for [`ListChildrenNeeded::All`] expressions: materializes the list and applies + /// the expression. + /// + /// Dispatches between two reads of the same chunk: + /// + /// * [`project_elements_bounded`](Self::project_elements_bounded) — used for a strict sub-range + /// when `elements` is a chunked layout. It decodes the first and last offset of the range and + /// reads only the elements they bound, so the underlying [`ChunkedReader`] fetches just the + /// element-chunks the range overlaps. This costs one offsets→elements round-trip. + /// * [`project_elements_whole_chunk`](Self::project_elements_whole_chunk) — used for a + /// full-range read, or when `elements` is a single flat segment (nothing to skip). It fetches + /// `offsets`, `elements`, and `validity` concurrently with no round-trip. + fn project_elements( + &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_elements_whole_chunk(row_range, expr, mask) + } else { + self.project_elements_bounded(row_range, expr, mask) + } + } + + /// Returns `true` if the `elements` child is a chunked layout, so a bounded elements read can + /// skip the chunks a row range does not overlap. + fn elements_are_chunked(&self) -> bool { + self.elements + .as_any() + .downcast_ref::() + .is_some() + } + + /// Whole-chunk read: the entire `elements` buffer, `offsets`, and `validity` for the chunk are + /// fetched concurrently — there is no offsets→elements round-trip because the elements bound is + /// the whole buffer (`offsets[0] == 0` within a chunk). The assembled list is sliced to + /// `row_range` and filtered by the caller mask in memory. + fn project_elements_whole_chunk( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let chunk_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. + // Offsets has one extra entry (`n + 1`). + let offsets_fut = self.offsets.projection_evaluation( + &(0..chunk_row_count + 1), + &root(), + MaskFuture::new_true(usize::try_from(chunk_row_count + 1)?), + )?; + let elements_fut = self.elements.projection_evaluation( + &(0..elements_row_count), + &root(), + MaskFuture::new_true(usize::try_from(elements_row_count)?), + )?; + let validity_fut = fetch_validity( + self.validity.as_ref(), + &(0..chunk_row_count), + MaskFuture::new_true(usize::try_from(chunk_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 == chunk_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 of a chunked-elements list. 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 elements buffer. With chunked elements, reading only + /// `[first..last)` fetches just the element-chunks the range overlaps. The offsets read is on + /// the critical path (the elements bound depends on it); validity is read concurrently. + fn project_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 session = self.session.clone(); + let elements_reader = Arc::clone(&self.elements); + + // Offsets are on the critical path (they bound the elements read); validity is independent. + let offsets_fut = self.fetch_offsets(row_range)?; + let validity_fut = fetch_validity( + self.validity.as_ref(), + row_range, + MaskFuture::new_true(row_count), + )?; + + Ok(async move { + let offsets = offsets_fut.await?; + let elements_range = calculate_elements_range(&offsets, &session)?; + let elements_len = usize::try_from(elements_range.end - elements_range.start)?; + + // Read only the elements this range covers; a chunked elements layout skips the rest. + let elements = elements_reader + .projection_evaluation( + &elements_range, + &root(), + MaskFuture::new_true(elements_len), + )? + .await?; + + // Rebase offsets to index into the sliced elements buffer. + let offsets = rebase_offsets(offsets, elements_range.start)?; + let validity = validity_fut.await?; + let list = + ListArray::try_new(elements, offsets, create_validity(validity, nullability))? + .into_array(); + + // 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()) + } + + /// Projection for [`ListChildrenNeeded::OffsetsAndValidity`] expressions (`list_length(root())` + /// and expressions composed from it): reads offsets and list validity, but never touches + /// element values. + fn project_offsets( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let offsets = self.fetch_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`. The offsets child has an extra entry, so reading + /// `row_range` maps to offsets in `[row_range.start..row_range.end + 1)`. + fn fetch_offsets(&self, row_range: &Range) -> VortexResult { + let offsets_range = row_range.start..(row_range.end + 1); + let offsets_count = usize::try_from(offsets_range.end - offsets_range.start)?; + self.offsets.projection_evaluation( + &offsets_range, + &root(), + MaskFuture::new_true(offsets_count), + ) + } +} + +/// Read `offsets[0]` and `offsets[-1]` and return the elements-buffer range they describe. +fn calculate_elements_range( + offsets: &ArrayRef, + session: &VortexSession, +) -> VortexResult> { + if offsets.is_empty() { + return Ok(0..0); + } + let mut exec_ctx = session.create_execution_ctx(); + let start = offsets + .execute_scalar(0, &mut exec_ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value fits in u64"); + let end = offsets + .execute_scalar(offsets.len() - 1, &mut exec_ctx)? + .as_primitive() + .as_::() + .vortex_expect("offset value fits in u64"); + Ok(start..end) +} + +/// Subtract `first` from every offset so the resulting offsets index into a sliced +/// `elements[first..]` buffer starting at zero. The constant array is cast to the offsets' dtype. +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) +} + +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 { + // All stats-based pruning should already be done upstream. + 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 { + // 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(row_range, expr, mask), + ListChildrenNeeded::All => self.project_elements(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()) +} + +/// 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_u32_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_offsets(&(1..3))?.await?; + assert_eq!(materialize_u32_array(offsets), vec![2, 4, 5]); + + Ok(()) + } + + #[rstest] + #[case::full_range(0..3, false)] + #[case::partial_start(0..2, false)] + #[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. + 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 fall back to 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..34a86612083 --- /dev/null +++ b/vortex-layout/src/layouts/list/writer.rs @@ -0,0 +1,511 @@ +// 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::stream; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::List; +use vortex_array::arrays::ListView; +use vortex_array::arrays::list::ListDataParts; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::dtype::DType; +use vortex_array::matcher::Matcher; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_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; + +/// Strategy for writing list-typed arrays, with a fallback for non-list dtypes. +/// +/// Single-chunk only. For list-typed input the strategy: +/// 1. Canonicalizes the input chunk into a [`ListView`]. +/// 2. Calls [`list_from_list_view`] to rebuild it into zero-copy-to-list form +/// (sorted, gapless, non-overlapping offsets) and produce a [`ListArray`]. +/// 3. Writes the `elements`, `offsets`, and (when nullable) `validity` columns into +/// separately configurable downstream strategies, producing a single [`ListLayout`]. +/// +/// For input whose dtype is not [`DType::List`], the stream is forwarded unchanged to the +/// configured `fallback` strategy. This lets `ListLayoutStrategy` slot in as a leaf strategy in +/// a heterogeneous column writer where some columns are lists and others are not. +/// +/// # Chunking +/// +/// `ListLayoutStrategy` bails on empty or multi-chunk input, matching the convention used by +/// [`FlatLayoutStrategy`]. +/// +/// [`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, + mut stream: SendableSequentialStream, + mut eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + if !dtype.is_list() { + // Non-list input: route to the configured fallback strategy unchanged. + return self + .fallback + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + + // Writer wants exactly one chunk + let Some(chunk) = stream.next().await else { + vortex_bail!("ListLayoutStrategy needs a single chunk"); + }; + let (sequence_id, array) = chunk?; + + let mut exec_ctx = session.create_execution_ctx(); + let ListDataParts { + elements, + offsets, + validity, + .. + } = canonicalize_to_list_parts(array, &mut exec_ctx)?; + + // There is one extra element in `offsets` + let row_count = offsets.len().saturating_sub(1); + let validity_array = dtype + .is_nullable() + .then(|| { + validity + .execute_mask(row_count, &mut exec_ctx) + .map(|m| m.into_array()) + }) + .transpose()?; + + // Spawn each child write onto the runtime so they run concurrently + let handle = session.handle(); + let (elements_task, offsets_task, validity_task) = { + let mut sp = sequence_id.descend(); + let mut spawn_layout_writer = |strategy: Arc, array: ArrayRef| { + let stream = single_chunk_stream(array.dtype().clone(), sp.advance(), array); + let child_eof = eof.split_off(); + let ctx = ctx.clone(); + let segment_sink = 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, stream, child_eof, &session) + .await + }) + }; + ( + spawn_layout_writer(Arc::clone(&self.elements), elements), + spawn_layout_writer(Arc::clone(&self.offsets), offsets), + validity_array.map(|arr| spawn_layout_writer(Arc::clone(&self.validity), arr)), + ) + }; + + // Should not have more than one chunk + if stream.next().await.is_some() { + vortex_bail!("ListLayoutStrategy received more than a single chunk"); + } + + let (elements_layout, offsets_layout, validity_layout) = + futures::try_join!(elements_task, offsets_task, async move { + match validity_task { + Some(t) => t.await.map(Some), + None => Ok(None), + } + },)?; + + Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) + } + + 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()) + } +} + +/// Canonicalize a list-dtype array into [`ListDataParts`]. Short-circuits when the input is +/// already a `List` or `ListView` array — otherwise drives the execution loop until one of +/// those forms appears. `ListView` is rebuilt into zero-copy-to-list form via +/// [`list_from_list_view`] before its parts are extracted. +fn canonicalize_to_list_parts( + array: ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let canonical = array.execute_until::(exec_ctx)?; + if let Some(list) = canonical.as_opt::() { + Ok(list.into_owned().into_data_parts()) + } else if let Some(view) = canonical.as_opt::() { + Ok(list_from_list_view(view.into_owned(), exec_ctx)?.into_data_parts()) + } else { + unreachable!("AnyList matcher guarantees List or ListView") + } +} + +/// Wrap a single array as a one-shot [`SendableSequentialStream`] for handoff to a child writer. +fn single_chunk_stream( + dtype: DType, + sequence_id: SequenceId, + array: ArrayRef, +) -> SendableSequentialStream { + SequentialStreamAdapter::new( + dtype, + stream::once(async move { Ok((sequence_id, array)) }).boxed(), + ) + .sendable() +} + +/// Matcher for `Array` or `Array`. Used to short-circuit the execution loop +/// when the input is already in (or directly produces) a list form, avoiding a redundant +/// `ListView` round-trip when the writer already has the parts it needs. +struct AnyList; + +impl Matcher for AnyList { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option> { + (array.as_opt::().is_some() || array.as_opt::().is_some()).then_some(()) + } +} + +#[cfg(test)] +mod tests { + use vortex_array::arrays::BoolArray; + 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: u32, 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: u32, 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 chunked_list_input_without_chunked_strategy_fails() -> VortexResult<()> { + let chunk0 = ListArray::try_new( + buffer![1i32, 2].into_array(), + buffer![0u32, 2].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let chunk1 = ListArray::try_new( + buffer![3i32, 4, 5].into_array(), + buffer![0u32, 3].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let chunked = + ChunkedArray::try_new(vec![chunk0, chunk1], i32_list_dtype(false))?.into_array(); + + let res = write(&flat_list_strategy(), chunked).await; + assert!(res.is_err()); + Ok(()) + } + + #[tokio::test] + async fn list_of_struct_tree() -> VortexResult<()> { + let struct_array = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3, 4, 5].into_array()), + ("b", buffer![10i32, 20, 30, 40, 50].into_array()), + ] + .as_slice(), + )? + .into_array(); + let list = ListArray::try_new( + struct_array, + buffer![0u32, 2, 5, 5].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let table_strategy: Arc = + Arc::new(TableStrategy::new(Arc::clone(&flat), Arc::clone(&flat))); + let writer = ListLayoutStrategy::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: u32, segment: 0 + "); + Ok(()) + } + + #[tokio::test] + async fn list_of_list_tree() -> VortexResult<()> { + let inner_list = ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 2, 5, 5, 6].into_array(), + Validity::NonNullable, + )? + .into_array(); + let list = ListArray::try_new( + inner_list, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let 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: u32, segment: 2 + └── offsets: vortex.flat, dtype: u32, 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: u32, segment: 3 + │ └── offsets: vortex.flat, dtype: u32, segment: 1 + └── offsets: vortex.flat, dtype: u32, segment: 0 + "); + Ok(()) + } + + #[tokio::test] + async fn chunked_list_input_with_chunked_strategy_succeeds() -> VortexResult<()> { + let chunk0 = ListArray::try_new( + 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: u32, segment: 1 + └── [1]: vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.flat, dtype: i32, segment: 2 + └── offsets: vortex.flat, dtype: u32, segment: 3 + "); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 18df5b8f347..47fa31aa3d9 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -16,6 +16,7 @@ pub mod dict; pub mod file_stats; pub mod flat; pub(crate) mod foreign; +pub mod list; pub(crate) mod partitioned; pub mod repartition; pub mod row_idx; diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index be938bdc192..6bb52b1d1e8 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -12,6 +12,7 @@ use crate::LayoutEncodingRef; use crate::layouts::chunked::ChunkedLayoutEncoding; use crate::layouts::dict::DictLayoutEncoding; use crate::layouts::flat::FlatLayoutEncoding; +use crate::layouts::list::ListLayoutEncoding; use crate::layouts::struct_::StructLayoutEncoding; use crate::layouts::zoned::LegacyStatsLayoutEncoding; use crate::layouts::zoned::ZonedLayoutEncoding; @@ -57,6 +58,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 } }