ListLayout baseline

vortex-file/src/strategy.rs

@@ -50,9 +50,11 @@ 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;
+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")]
@@ -157,6 +159,10 @@ pub struct WriteStrategyBuilder {
     allow_encodings: Option<HashSet<ArrayId>>,
     flat_strategy: Option<Arc<dyn LayoutStrategy>>,
     probe_compressor: Option<Arc<dyn CompressorPlugin>>,
+    /// Whether to write list fields using [`ListLayoutStrategy`].
+    ///
+    /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy
+    use_list_layout: bool,
 }
 
 impl Default for WriteStrategyBuilder {
@@ -170,6 +176,7 @@ impl Default for WriteStrategyBuilder {
             allow_encodings: Some(ALLOWED_ENCODINGS.clone()),
             flat_strategy: None,
             probe_compressor: None,
+            use_list_layout: use_experimental_list_layout(),
         }
     }
 }
@@ -184,6 +191,17 @@ impl WriteStrategyBuilder {
         self
     }
 
+    /// Enable writing list fields with [`ListLayoutStrategy`].
+    ///
+    /// **Note**: this is an unstable and experimental layout that is expected to change.
+    /// Using it may lead to unreadable files in the future.
+    ///
+    /// [`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
@@ -308,12 +326,14 @@ impl WriteStrategyBuilder {
             probe_compressor,
         );
 
+        let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0");
+
         // 2. calculate stats for each row group
         let stats = ZonedStrategy::new(
             dict,
             compress_then_flat.clone(),
             ZonedLayoutOptions {
-                block_size: NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"),
+                block_size: row_block_size,
                 ..Default::default()
             },
         );
@@ -332,11 +352,37 @@ impl WriteStrategyBuilder {
         );
 
         // 0. start with splitting columns
-        let validity_strategy = CollectStrategy::new(compress_then_flat);
+        let validity_strategy = CollectStrategy::new(compress_then_flat.clone());
 
         // 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 {
+            // We need a closure here to enable recursive application of list layout.
+            table_strategy = table_strategy.with_list_layout_factory(
+                move |list_layout: ListLayoutStrategy| -> Arc<dyn LayoutStrategy> {
+                    let zoned = ZonedStrategy::new(
+                        list_layout,
+                        compress_then_flat.clone(),
+                        ZonedLayoutOptions {
+                            block_size: row_block_size,
+                            ..Default::default()
+                        },
+                    );
+                    Arc::new(RepartitionStrategy::new(
+                        zoned,
+                        RepartitionWriterOptions {
+                            block_size_minimum: 0,
+                            block_len_multiple: row_block_size.get(),
+                            block_size_target: None,
+                            canonicalize: false,
+                        },
+                    ))
+                },
+            );
+        }
 
         Arc::new(table_strategy)
     }

vortex-file/src/tests.rs

@@ -15,6 +15,7 @@ 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;
@@ -1671,6 +1672,75 @@ 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<ArrayRef> {
+    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<list<i32>>` 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 `struct<{ items: list<struct<{a,b}>>? }>` 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())])?

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,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::<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);
+    }
+}