feat: chunk list elements + bounded sub-range reader

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<dyn LayoutStrategy> = 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<dyn LayoutStrategy> = 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(

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<()> {

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"] }

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;

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::<IsNull>() || expr.is::<IsNotNull>())
+        && expr.children().len() == 1
+        && is_root(expr.child(0))
+}
+
+fn is_list_length_root(expr: &Expression) -> bool {
+    expr.is::<ListLength>() && 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<Expression> {
+    if expr.is::<IsNotNull>() && expr.children().len() == 1 && is_root(expr.child(0)) {
+        return Ok(root());
+    }
+    if expr.is::<IsNull>() && expr.children().len() == 1 && is_root(expr.child(0)) {
+        return Ok(not(root()));
+    }
+    let children = expr
+        .children()
+        .iter()
+        .map(rewrite_validity_expr)
+        .collect::<VortexResult<Vec<_>>>()?;
+    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<Expression> {
+    if is_list_length_root(expr) {
+        return Ok(root());
+    }
+
+    let children = expr
+        .children()
+        .iter()
+        .map(rewrite_offsets_expr)
+        .collect::<VortexResult<Vec<_>>>()?;
+    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);
+    }
+}