Prune plan scans with zoned statistics

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

Cargo.lock

@@ -358,7 +358,11 @@ impl ZonedLayout {
 }
 
 impl ZonedData {
-    fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> {
+    pub(crate) fn zone_len(&self) -> usize {
+        self.zone_len
+    }
+
+    pub(crate) fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> {
         match &self.zone_map_schema {
             ZoneMapSchema::LegacyStats(stats) => stats
                 .iter()

vortex-layout/src/layouts/zoned/mod.rs

@@ -18,6 +18,7 @@ use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND;
 use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax;
 use vortex_array::aggregate_fn::fns::sum::Sum;
 use vortex_array::aggregate_fn::fns::sum::normalize_legacy_partial_array;
+use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
@@ -92,7 +93,7 @@ impl ZoneMap {
         Ok(unsafe { Self::new_unchecked(column_dtype, array, aggregate_fns, zone_len, row_count) })
     }
 
-    pub(super) unsafe fn new_unchecked(
+    pub(crate) unsafe fn new_unchecked(
         column_dtype: DType,
         array: StructArray,
         aggregate_fns: Arc<[AggregateFnRef]>,
@@ -152,19 +153,32 @@ impl ZoneMap {
         session: &VortexSession,
     ) -> VortexResult<Mask> {
         let mut ctx = session.create_execution_ctx();
-        let num_zones = self.array.len();
-        let predicate = self.lower_stats(predicate.clone())?;
+        self.applied_predicate(predicate)?
+            .null_as_false()
+            .execute(&mut ctx)
+    }
 
-        let array = self.array.clone().into_array();
-        let applied = array.apply_bound(&predicate)?;
+    /// Evaluates a pruning predicate while preserving unknown (null) proof values.
+    pub(crate) fn evaluate(
+        &self,
+        predicate: &BoundExpression,
+        session: &VortexSession,
+    ) -> VortexResult<BoolArray> {
+        let mut ctx = session.create_execution_ctx();
+        self.applied_predicate(predicate)?
+            .execute::<BoolArray>(&mut ctx)
+    }
 
+    fn applied_predicate(&self, predicate: &BoundExpression) -> VortexResult<ArrayRef> {
+        let num_zones = self.array.len();
+        let predicate = self.lower_stats(predicate.clone())?;
+        let applied = self.array.clone().into_array().apply_bound(&predicate)?;
         if !contains_row_count(&applied) {
-            return applied.null_as_false().execute(&mut ctx);
+            return Ok(applied);
         }
 
         let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones)?;
-        let substituted = substitute_row_count(applied, &row_count_array)?;
-        substituted.null_as_false().execute(&mut ctx)
+        substitute_row_count(applied, &row_count_array)
     }
 
     fn lower_stats(&self, predicate: BoundExpression) -> VortexResult<BoundExpression> {

vortex-layout/src/layouts/zoned/zone_map.rs

@@ -148,9 +148,18 @@ fn lazy_children(layout: LayoutRef, slots: Vec<usize>) -> PlanChildren {
 
 fn lower_zoned(layout: &LayoutRef) -> VortexResult<ZonedPlan> {
     // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary zones.
+    let metadata = if let Some(layout) = layout.as_opt::<Zoned>() {
+        layout.data()
+    } else if let Some(layout) = layout.as_opt::<LegacyStats>() {
+        layout.data()
+    } else {
+        vortex_bail!("Zoned plan requires a zoned layout")
+    };
     Ok(ZonedPlan::from_children(
         layout.dtype().clone(),
         layout.row_count(),
         lazy_children(Arc::clone(layout), vec![0, 1]),
+        u64::try_from(metadata.zone_len())?,
+        metadata.aggregate_fns(),
     ))
 }

vortex-layout/src/plan/lower.rs

@@ -58,6 +58,7 @@ pub use plans::Take;
 pub use plans::TakeData;
 pub use plans::TakePlan;
 pub use plans::Zoned;
+pub use plans::ZonedData;
 pub use plans::ZonedPlan;
 pub use plans::row_idx_dtype;
 pub use typed::DynPlan;

vortex-layout/src/plan/mod.rs

@@ -16,10 +16,12 @@ use super::Pack;
 use super::PlanRef;
 use super::RowIdx;
 use super::Take;
+use super::Zoned;
 use super::plans::ExpressionConcatRule;
 use super::plans::ExpressionPackRule;
 use super::plans::ExpressionRowIdxRule;
 use super::plans::ExpressionTakeRule;
+use super::plans::ExpressionZonedRule;
 
 static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter<Concat, ExpressionConcatRule> =
     PlanParentReduceRuleAdapter::new(ExpressionConcatRule);
@@ -29,12 +31,15 @@ static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter<RowIdx, ExpressionRo
     PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule);
 static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter<Pack, ExpressionPackRule> =
     PlanParentReduceRuleAdapter::new(ExpressionPackRule);
+static EXPRESSION_ZONED_RULE: PlanParentReduceRuleAdapter<Zoned, ExpressionZonedRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionZonedRule);
 
 static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
     &EXPRESSION_CONCAT_RULE,
     &EXPRESSION_TAKE_RULE,
     &EXPRESSION_ROW_IDX_RULE,
     &EXPRESSION_PACK_RULE,
+    &EXPRESSION_ZONED_RULE,
 ]);
 
 /// Attempts a static rewrite for `parent` and its child at `child_idx`.

vortex-layout/src/plan/optimizer/mod.rs

@@ -45,5 +45,7 @@ pub(crate) use take::ExpressionTakeRule;
 pub use take::Take;
 pub use take::TakeData;
 pub use take::TakePlan;
+pub(crate) use zoned::ExpressionZonedRule;
 pub use zoned::Zoned;
+pub use zoned::ZonedData;
 pub use zoned::ZonedPlan;