perf(layout): prune multi-field boolean predicates in the struct reader

vortex-layout/src/layouts/struct_/reader.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::ops::BitAnd;
 use std::ops::Range;
 use std::sync::Arc;
 use std::sync::OnceLock;
@@ -9,9 +10,11 @@ use futures::FutureExt;
 use futures::future::try_join_all;
 use futures::try_join;
 use itertools::Itertools;
+use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
 use vortex_array::MaskFuture;
 use vortex_array::arrays::StructArray;
+use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldMask;
 use vortex_array::dtype::FieldName;
@@ -33,10 +36,12 @@ use vortex_array::expr::transform::partition_bound;
 use vortex_array::expr::traversal::NodeExt;
 use vortex_array::expr::traversal::Transformed;
 use vortex_array::expr::traversal::TraversalOrder;
+use vortex_array::scalar_fn::fns::binary::Binary;
 use vortex_array::scalar_fn::fns::get_item::GetItem;
 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::mask::Mask as MaskFn;
+use vortex_array::scalar_fn::fns::operators::Operator;
 use vortex_array::scalar_fn::fns::pack::Pack;
 use vortex_array::scalar_fn::fns::select::FieldSelection;
 use vortex_array::scalar_fn::fns::select::Select;
@@ -602,9 +607,29 @@ impl LayoutReader for StructReader {
                         ))
                     })
             }
-            Partitioned::Multi(_) | Partitioned::Direct(_) => {
-                // TODO(ngates): if all partitions are boolean, we can use a pruning evaluation. Otherwise
-                //  there's not much we can do? Maybe... it's complicated...
+            Partitioned::Multi(_) => {
+                // Only split after single-field delegation so field readers see whole predicates.
+                // AND/OR compose conservative "may match" masks, even for nullable predicates;
+                // other multi-field operators cannot generally combine these masks (e.g. NOT).
+                let Some(operator @ (Operator::And | Operator::Or)) =
+                    expr.as_opt::<Binary>().copied()
+                else {
+                    return Ok(MaskFuture::ready(mask));
+                };
+                let lhs = self.pruning_evaluation(row_range, &expr.children()[0], mask.clone())?;
+                let rhs = self.pruning_evaluation(row_range, &expr.children()[1], mask.clone())?;
+                Ok(MaskFuture::new(mask.len(), async move {
+                    let (lhs, rhs) = try_join!(lhs, rhs)?;
+                    let combined = if operator == Operator::And {
+                        &lhs & &rhs
+                    } else {
+                        &lhs | &rhs
+                    };
+                    Ok(mask.bitand(&combined))
+                }))
+            }
+            Partitioned::Direct(_) => {
+                // A pass-through projection has no predicate to prune.
                 Ok(MaskFuture::ready(mask))
             }
         }
@@ -698,6 +723,7 @@ impl LayoutReader for StructReader {
 
 #[cfg(test)]
 mod tests {
+    use std::num::NonZeroUsize;
     use std::sync::Arc;
 
     use rstest::fixture;
@@ -709,6 +735,7 @@ mod tests {
     use vortex_array::VortexSessionExecute;
     use vortex_array::array_session;
     use vortex_array::arrays::BoolArray;
+    use vortex_array::arrays::ChunkedArray;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::StructArray;
     use vortex_array::arrays::struct_::StructArrayExt;
@@ -729,6 +756,7 @@ mod tests {
     use vortex_array::expr::is_not_null;
     use vortex_array::expr::is_null;
     use vortex_array::expr::lit;
+    use vortex_array::expr::not;
     use vortex_array::expr::or;
     use vortex_array::expr::pack;
     use vortex_array::expr::root;
@@ -745,8 +773,11 @@ mod tests {
     use crate::LayoutReaderRef;
     use crate::LayoutRef;
     use crate::LayoutStrategy;
+    use crate::layouts::chunked::writer::ChunkedLayoutStrategy;
     use crate::layouts::flat::writer::FlatLayoutStrategy;
     use crate::layouts::table::TableStrategy;
+    use crate::layouts::zoned::writer::ZonedLayoutOptions;
+    use crate::layouts::zoned::writer::ZonedStrategy;
     use crate::segments::SegmentSource;
     use crate::segments::TestSegments;
     use crate::sequence::SequenceId;
@@ -1780,4 +1811,124 @@ mod tests {
         let err = result.err().unwrap().to_string();
         assert!(err.contains("Cannot compare different DTypes"), "{err}");
     }
+
+    /// Nine rows sorted by `(a, b)`; three zones of three; `c` copies `b` with nulls.
+    /// All three columns are zoned.
+    #[fixture]
+    fn sorted_zoned_layout() -> (Arc<dyn SegmentSource>, LayoutRef) {
+        let ctx = ArrayContext::empty();
+        let segments = Arc::new(TestSegments::default());
+        let (ptr, eof) = SequenceId::root().split();
+        let zoned = ZonedStrategy::new(
+            ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()),
+            FlatLayoutStrategy::default(),
+            ZonedLayoutOptions {
+                block_size: NonZeroUsize::new(3).unwrap(),
+                ..Default::default()
+            },
+        );
+        let strategy = TableStrategy::new(Arc::new(FlatLayoutStrategy::default()), Arc::new(zoned));
+        let chunk = |a: [i32; 3], b: [i32; 3]| {
+            StructArray::from_fields(
+                [
+                    ("a", buffer![a[0], a[1], a[2]].into_array()),
+                    ("b", buffer![b[0], b[1], b[2]].into_array()),
+                    (
+                        "c",
+                        PrimitiveArray::from_option_iter([Some(b[0]), Some(b[1]), None])
+                            .into_array(),
+                    ),
+                ]
+                .as_slice(),
+            )
+            .unwrap()
+            .into_array()
+        };
+        let array_stream = ChunkedArray::from_iter([
+            chunk([1, 1, 1], [1, 2, 3]),
+            chunk([1, 2, 2], [4, 1, 2]),
+            chunk([2, 2, 3], [3, 4, 1]),
+        ])
+        .into_array()
+        .to_array_stream()
+        .sequenced(ptr);
+        let segments2 = Arc::<TestSegments>::clone(&segments);
+        let layout = block_on(|handle| async move {
+            let session = new_session().with_handle(handle);
+            strategy
+                .write_stream(ctx.into(), segments2, array_stream, eof, &session)
+                .await
+        })
+        .unwrap();
+        (segments, layout)
+    }
+
+    #[rstest]
+    #[case::lexicographic_bound(
+        or(
+            gt(col("a"), lit(1)),
+            and(eq(col("a"), lit(1)), gt(col("b"), lit(4))),
+        ),
+        [false, true, true],
+    )]
+    #[case::nested_and_or(
+        and(
+            or(gt(col("a"), lit(1)), gt(col("b"), lit(4))),
+            or(eq(col("a"), lit(1)), gt(col("b"), lit(2))),
+        ),
+        [false, true, true],
+    )]
+    #[case::single_field_compound(
+        or(
+            and(gt(col("a"), lit(1)), eq(col("a"), lit(3))),
+            gt(col("b"), lit(4)),
+        ),
+        [false, false, true],
+    )]
+    #[case::unsupported_and_branch(
+        and(gt(col("a"), lit(1)), gt(col("a"), col("b"))),
+        [false, true, true],
+    )]
+    #[case::unsupported_or_branch(
+        or(gt(col("a"), lit(2)), gt(col("a"), col("b"))),
+        [true, true, true],
+    )]
+    #[case::multi_field_negation(
+        not(or(
+            gt(col("a"), lit(1)),
+            and(eq(col("a"), lit(1)), gt(col("b"), lit(4))),
+        )),
+        [true, true, true],
+    )]
+    #[case::nullable_and(
+        and(gt(col("a"), lit(1)), gt(col("c"), lit(3))),
+        [false, true, true],
+    )]
+    #[case::nullable_or(
+        or(gt(col("a"), lit(2)), gt(col("c"), lit(4))),
+        [false, false, true],
+    )]
+    fn multi_field_predicates_prune(
+        #[from(sorted_zoned_layout)] (segments, layout): (Arc<dyn SegmentSource>, LayoutRef),
+        #[case] expr: Expression,
+        #[case] expected_zones: [bool; 3],
+        #[values(
+            Mask::new_true(9),
+            Mask::new_false(9),
+            Mask::from_iter([true, false, true, false, true, false, true, false, true]),
+        )]
+        mask: Mask,
+    ) -> VortexResult<()> {
+        block_on(|handle| async {
+            let session = new_session().with_handle(handle);
+            let reader = layout.new_reader("".into(), segments, &session, &Default::default())?;
+            let expr = expr.bind(reader.dtype())?;
+            let result = reader
+                .pruning_evaluation(&(0..layout.row_count()), &expr, mask.clone())?
+                .await?;
+            let expected = Mask::from_iter(expected_zones.into_iter().flat_map(|keep| [keep; 3]));
+            assert_eq!(result, &mask & &expected);
+            Ok(())
+        })
+    }
 }