diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 877126b2632..b71e719c106 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/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::().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, 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::::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, 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(()) + }) + } }