ListLayout leaf version

vortex-file/src/strategy.rs

@@ -50,6 +50,8 @@ 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::use_experimental_list_layout;
+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;
@@ -242,8 +244,20 @@ 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 layout
+        let leaf: Arc<dyn LayoutStrategy> = if use_experimental_list_layout() {
+            Arc::new(
+                ListLayoutStrategy::default()
+                    .with_elements(Arc::clone(&flat))
+                    .with_offsets(Arc::clone(&flat))
+                    .with_validity(Arc::clone(&flat))
+                    .with_fallback(Arc::clone(&flat)),
+            )
+        } else {
+            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
 

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