Add fixed-size-list layout

vortex-file/src/strategy.rs

@@ -53,6 +53,7 @@ use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
 use vortex_layout::layouts::repartition::RepartitionStrategy;
 use vortex_layout::layouts::repartition::RepartitionWriterOptions;
 use vortex_layout::layouts::table::TableStrategy;
+use vortex_layout::layouts::table::use_experimental_list_layout;
 use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions;
 use vortex_layout::layouts::zoned::writer::ZonedStrategy;
 #[cfg(feature = "unstable_encodings")]
@@ -156,6 +157,10 @@ pub struct WriteStrategyBuilder {
     field_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
     allow_encodings: Option<HashSet<ArrayId>>,
     flat_strategy: Option<Arc<dyn LayoutStrategy>>,
+    /// Whether to write list-like fields using structural list layouts.
+    ///
+    /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy
+    use_list_layout: bool,
 }
 
 impl Default for WriteStrategyBuilder {
@@ -168,6 +173,7 @@ impl Default for WriteStrategyBuilder {
             field_writers: HashMap::new(),
             allow_encodings: Some(ALLOWED_ENCODINGS.clone()),
             flat_strategy: None,
+            use_list_layout: use_experimental_list_layout(),
         }
     }
 }
@@ -182,6 +188,14 @@ impl WriteStrategyBuilder {
         self
     }
 
+    /// Enable writing list and fixed-size-list fields with structural list layouts.
+    ///
+    /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy
+    pub fn with_list_layout(mut self) -> Self {
+        self.use_list_layout = true;
+        self
+    }
+
     /// Override the write layout for a specific field somewhere in the nested schema tree.
     ///
     /// The field path is matched after the root struct is split into columns. This is useful when a
@@ -242,8 +256,9 @@ impl WriteStrategyBuilder {
             Arc::new(FlatLayoutStrategy::default())
         };
 
-        // 7. for each chunk create a flat layout
-        let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat));
+        // 7. for each chunk create a leaf layout.
+        let leaf: Arc<dyn LayoutStrategy> = Arc::clone(&flat);
+        let chunked = ChunkedLayoutStrategy::new(leaf);
         // 6. buffer chunks so they end up with closer segment ids physically
         let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB
 
@@ -321,8 +336,13 @@ impl WriteStrategyBuilder {
         let validity_strategy = CollectStrategy::new(compress_then_flat);
 
         // Take any field overrides from the builder and apply them to the final strategy.
-        let table_strategy = TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition))
-            .with_field_writers(self.field_writers);
+        let mut table_strategy =
+            TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition))
+                .with_field_writers(self.field_writers);
+
+        if self.use_list_layout {
+            table_strategy = table_strategy.with_list_layout();
+        }
 
         Arc::new(table_strategy)
     }

vortex-file/src/tests.rs

@@ -15,10 +15,12 @@ use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
 use vortex_array::array_session;
+use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ChunkedArray;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::DecimalArray;
 use vortex_array::arrays::Dict;
+use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::ListArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
@@ -1668,6 +1670,95 @@ async fn test_writer_with_complex_types() -> VortexResult<()> {
     Ok(())
 }
 
+/// Write `array` with list decomposition forced on (through the full compress/zone pipeline) and
+/// read the whole thing back.
+async fn write_read_roundtrip(array: ArrayRef) -> VortexResult<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 `fixed_size_list<fixed_size_list<i32>>` column round-trips through the `TableStrategy`
+/// dispatcher, exercising fixed-size-list decomposition recursing into itself.
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn nested_fixed_size_list_of_fixed_size_list_roundtrip() -> VortexResult<()> {
+    let inner = FixedSizeListArray::new(
+        buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(),
+        2,
+        Validity::NonNullable,
+        4,
+    )
+    .into_array();
+    let outer = FixedSizeListArray::new(inner, 2, Validity::NonNullable, 2).into_array();
+    let st = StructArray::from_fields(&[("nested", outer)])?.into_array();
+
+    let result = write_read_roundtrip(st.clone()).await?;
+    assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx());
+    Ok(())
+}
+
+/// A `struct<{ items: list<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/fixed_size_list/expr.rs

@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::expr::Expression;
+use vortex_array::expr::is_root;
+use vortex_array::expr::not;
+use vortex_array::expr::root;
+use vortex_array::scalar_fn::fns::is_not_null::IsNotNull;
+use vortex_array::scalar_fn::fns::is_null::IsNull;
+use vortex_array::scalar_fn::fns::list_length::ListLength;
+use vortex_error::VortexResult;
+
+/// The minimal set of fixed-size-list children an expression needs for evaluation.
+///
+/// For example:
+///     - `is_null(root())` only needs validity.
+///     - `list_length(root())` needs the fixed list size and validity, but not elements.
+///     - `root()` needs elements and validity.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub(super) enum FixedSizeListChildrenNeeded {
+    /// Only validity is needed (`is_null` / `is_not_null`).
+    Validity,
+    /// Only the fixed list size and validity are needed (`list_length`).
+    ListLengthAndValidity,
+    /// Elements and validity are needed.
+    Elements,
+}
+
+/// The minimal set of fixed-size-list children needed to evaluate `expr`, where `root()` is a
+/// field with fixed-size-list dtype.
+pub(super) fn get_necessary_fixed_size_list_children(
+    expr: &Expression,
+) -> FixedSizeListChildrenNeeded {
+    if is_null_root(expr) {
+        return FixedSizeListChildrenNeeded::Validity;
+    }
+
+    if is_list_length_root(expr) {
+        return FixedSizeListChildrenNeeded::ListLengthAndValidity;
+    }
+
+    if is_root(expr) {
+        return FixedSizeListChildrenNeeded::Elements;
+    }
+
+    expr.children()
+        .iter()
+        .map(get_necessary_fixed_size_list_children)
+        .max()
+        .unwrap_or(FixedSizeListChildrenNeeded::Validity)
+}
+
+fn is_null_root(expr: &Expression) -> bool {
+    (expr.is::<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 fixed-size-list's
+/// validity bool array (`true` == valid row): `is_not_null(root())` becomes `root()` and
+/// `is_null(root())` becomes `not(root())`. All other nodes are rebuilt with rewritten children.
+pub(super) fn rewrite_validity_expr(expr: &Expression) -> VortexResult<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 a list-length-class expression so it can be evaluated against an array of list lengths.
+/// `list_length(root())` becomes `root()`. Other references to `root()` are left intact: for
+/// list-length-class expressions they can only be validity checks, and the lengths array carries
+/// the same validity as the original fixed-size-list.
+pub(super) fn rewrite_list_length_expr(expr: &Expression) -> VortexResult<Expression> {
+    if is_list_length_root(expr) {
+        return Ok(root());
+    }
+
+    let children = expr
+        .children()
+        .iter()
+        .map(rewrite_list_length_expr)
+        .collect::<VortexResult<Vec<_>>>()?;
+    expr.clone().with_children(children)
+}