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

@@ -17,6 +17,7 @@ use vortex_array::aggregate_fn::fns::all_null::AllNull;
 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::nan_count::NanCount;
+use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
@@ -87,7 +88,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]>,
@@ -143,18 +144,34 @@ impl ZoneMap {
     /// only after the predicate has been lowered to the zone-map table.
     pub fn prune(&self, predicate: &Expression, session: &VortexSession) -> VortexResult<Mask> {
         let mut ctx = session.create_execution_ctx();
+        self.applied_predicate(predicate)?
+            .null_as_false()
+            .execute(&mut ctx)
+    }
+
+    /// Evaluate a pruning predicate while preserving unknown (null) proof values.
+    pub(crate) fn evaluate(
+        &self,
+        predicate: &Expression,
+        session: &VortexSession,
+    ) -> VortexResult<BoolArray> {
+        let mut ctx = session.create_execution_ctx();
+        self.applied_predicate(predicate)?
+            .execute::<BoolArray>(&mut ctx)
+    }
+
+    fn applied_predicate(&self, predicate: &Expression) -> VortexResult<ArrayRef> {
         let num_zones = self.array.len();
         let predicate = self.lower_stats(predicate.clone())?;
 
         let applied = self.array.clone().into_array().apply(&predicate)?;
-
-        if !contains_row_count(&applied) {
-            return applied.null_as_false().execute(&mut ctx);
-        }
-
-        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)
+        let applied = if contains_row_count(&applied) {
+            let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones)?;
+            substitute_row_count(applied, &row_count_array)?
+        } else {
+            applied
+        };
+        Ok(applied)
     }
 
     fn lower_stats(&self, predicate: Expression) -> VortexResult<Expression> {

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

@@ -5,6 +5,7 @@ use std::fmt;
 
 use super::ExpressionPlan;
 use super::Plan;
+use super::ZonedPlan;
 
 /// Context threaded through a plan tree traversal.
 pub struct PlanTreeContext {
@@ -110,7 +111,7 @@ impl PlanTreeExtractor for PlanSummaryExtractor {
     }
 }
 
-/// Adds an expression annotation to [`ExpressionPlan`] nodes.
+/// Adds expression annotations to expression and zoned-pruning plans.
 pub struct PlanExpressionExtractor;
 
 impl PlanTreeExtractor for PlanExpressionExtractor {
@@ -120,10 +121,16 @@ impl PlanTreeExtractor for PlanExpressionExtractor {
         _context: &PlanTreeContext,
         formatter: &mut fmt::Formatter<'_>,
     ) -> fmt::Result {
-        let Some(expression_plan) = plan.downcast_ref::<ExpressionPlan>() else {
-            return Ok(());
-        };
-        write!(formatter, " expr={}", expression_plan.expression())
+        if let Some(expression_plan) = plan.downcast_ref::<ExpressionPlan>() {
+            return write!(formatter, " expr={}", expression_plan.expression());
+        }
+        if let Some(expression) = plan
+            .downcast_ref::<ZonedPlan>()
+            .and_then(ZonedPlan::pruning_expression)
+        {
+            return write!(formatter, " prune={expression}");
+        }
+        Ok(())
     }
 }
 

vortex-layout/src/plan/display.rs

@@ -16,10 +16,12 @@ use super::DictPlan;
 use super::PlanRef;
 use super::RowIdxPlan;
 use super::StructPlan;
+use super::ZonedPlan;
 use super::plans::ExpressionChunkedRule;
 use super::plans::ExpressionDictRule;
 use super::plans::ExpressionRowIdxRule;
 use super::plans::ExpressionStructRule;
+use super::plans::ExpressionZonedRule;
 
 static EXPRESSION_CHUNKED_RULE: PlanParentReduceRuleAdapter<ChunkedPlan, ExpressionChunkedRule> =
     PlanParentReduceRuleAdapter::new(ExpressionChunkedRule);
@@ -29,12 +31,15 @@ static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter<RowIdxPlan, Expressi
     PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule);
 static EXPRESSION_STRUCT_RULE: PlanParentReduceRuleAdapter<StructPlan, ExpressionStructRule> =
     PlanParentReduceRuleAdapter::new(ExpressionStructRule);
+static EXPRESSION_ZONED_RULE: PlanParentReduceRuleAdapter<ZonedPlan, ExpressionZonedRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionZonedRule);
 
 static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
     &EXPRESSION_CHUNKED_RULE,
     &EXPRESSION_DICT_RULE,
     &EXPRESSION_ROW_IDX_RULE,
     &EXPRESSION_STRUCT_RULE,
+    &EXPRESSION_ZONED_RULE,
 ]);
 
 /// Attempts a static rewrite for `parent` and its child at `child_idx`.

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

@@ -23,4 +23,5 @@ pub use row_idx::RowIdxPlan;
 pub use row_idx::RowIdxValuesPlan;
 pub(crate) use struct_::ExpressionStructRule;
 pub use struct_::StructPlan;
+pub(crate) use zoned::ExpressionZonedRule;
 pub use zoned::ZonedPlan;