perf(layout): lower zoned pruning predicates once

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

@@ -102,16 +102,23 @@ impl PruningState {
                         Some(
                             async move {
                                 let zone_map = zone_map.await?;
-                                let initial_mask =
-                                    zone_map.prune(&predicate, &session).map_err(|err| {
-                                        err.with_context(format!(
+                                let lowered = zone_map.lower_stats(predicate.clone()).and_then(
+                                    |lowered_predicate| {
+                                        let mask =
+                                            zone_map.prune_lowered(&lowered_predicate, &session)?;
+                                        Ok((lowered_predicate, mask))
+                                    },
+                                );
+                                let (lowered_predicate, initial_mask) = lowered.map_err(|err| {
+                                    err.with_context(format!(
                                         "While evaluating pruning predicate {} (derived from {})",
                                         predicate, expr
                                     ))
-                                    })?;
+                                })?;
+
                                 Ok(Arc::new(PruningResult {
                                     zone_map,
-                                    predicate,
+                                    lowered_predicate,
                                     dynamic_updates,
                                     latest_result: RwLock::new((0, initial_mask)),
                                     session,
@@ -191,7 +198,7 @@ impl PruningState {
 
 pub(super) struct PruningResult {
     zone_map: ZoneMap,
-    predicate: BoundExpression,
+    lowered_predicate: BoundExpression,
     dynamic_updates: Option<DynamicExprUpdates>,
     latest_result: RwLock<(u64, Mask)>,
     session: VortexSession,
@@ -219,17 +226,17 @@ impl PruningResult {
 
         trace!(
             version,
-            predicate = %self.predicate,
+            predicate = %self.lowered_predicate,
             "recomputing pruning mask"
         );
 
         let next_mask = self
             .zone_map
-            .prune(&self.predicate, &self.session)
+            .prune_lowered(&self.lowered_predicate, &self.session)
             .map_err(|err| {
                 err.with_context(format!(
                     "While evaluating pruning predicate {}",
-                    self.predicate
+                    self.lowered_predicate
                 ))
             })?;
         *guard = (version, next_mask.clone());

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

@@ -145,13 +145,26 @@ impl ZoneMap {
         &self,
         predicate: &BoundExpression,
         session: &VortexSession,
+    ) -> VortexResult<Mask> {
+        let predicate = self.lower_stats(predicate.clone())?;
+        self.prune_lowered(&predicate, session)
+    }
+
+    /// Apply a predicate that [`Self::lower_stats`] already lowered to this zone map's
+    /// stats table.
+    ///
+    /// Lowering depends only on the column and stats-table dtypes, so callers that evaluate
+    /// the same predicate many times lower it once and call this function for each evaluation.
+    pub(super) fn prune_lowered(
+        &self,
+        predicate: &BoundExpression,
+        session: &VortexSession,
     ) -> VortexResult<Mask> {
         let mut ctx = session.create_execution_ctx();
         let num_zones = self.array.len();
-        let predicate = self.lower_stats(predicate.clone())?;
 
         let array = self.array.clone().into_array();
-        let applied = array.apply_bound(&predicate)?;
+        let applied = array.apply_bound(predicate)?;
 
         if !contains_row_count(&applied) {
             return null_as_false_mask(applied, &mut ctx);
@@ -162,7 +175,8 @@ impl ZoneMap {
         null_as_false_mask(substituted, &mut ctx)
     }
 
-    fn lower_stats(&self, predicate: BoundExpression) -> VortexResult<BoundExpression> {
+    /// Rewrite the stats references in `predicate` into expressions over the stats table.
+    pub(super) fn lower_stats(&self, predicate: BoundExpression) -> VortexResult<BoundExpression> {
         let binder = ZoneMapStatsBinder { zone_map: self };
         bind_stats(predicate, &binder)
     }
@@ -377,6 +391,8 @@ fn row_count_array(zone_len: u64, row_count: u64, num_zones: usize) -> VortexRes
 mod tests {
     use std::num::NonZeroUsize;
     use std::sync::Arc;
+    use std::sync::atomic::AtomicI32;
+    use std::sync::atomic::Ordering;
 
     use rstest::rstest;
     use vortex_array::IntoArray;
@@ -408,6 +424,7 @@ mod tests {
     use vortex_array::expr::BoundExpression;
     use vortex_array::expr::Expression;
     use vortex_array::expr::cast;
+    use vortex_array::expr::dynamic;
     use vortex_array::expr::gt;
     use vortex_array::expr::gt_eq;
     use vortex_array::expr::is_not_null;
@@ -417,6 +434,8 @@ mod tests {
     use vortex_array::expr::not_eq;
     use vortex_array::expr::root;
     use vortex_array::expr::stats::Stat;
+    use vortex_array::scalar::ScalarValue;
+    use vortex_array::scalar_fn::fns::operators::CompareOperator;
     use vortex_array::stats::all_nan;
     use vortex_array::stats::all_non_nan;
     use vortex_array::stats::all_non_null;
@@ -1076,4 +1095,60 @@ mod tests {
             &mut SESSION.create_execution_ctx()
         );
     }
+
+    #[test]
+    fn lowered_dynamic_predicate_tracks_rhs_updates() -> VortexResult<()> {
+        let max = Max.bind(NumericalAggregateOpts::skip_nans());
+        let min = Min.bind(NumericalAggregateOpts::skip_nans());
+        let zone_map = ZoneMap::try_new(
+            PType::I32.into(),
+            StructArray::from_fields(&[
+                (
+                    max.to_string(),
+                    PrimitiveArray::new(buffer![5i32, 6i32, 7i32], Validity::AllValid).into_array(),
+                ),
+                (
+                    min.to_string(),
+                    PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::AllValid).into_array(),
+                ),
+            ])?,
+            Arc::new([max, min]),
+            3,
+            10,
+        )?;
+        let ctx = &mut SESSION.create_execution_ctx();
+
+        let threshold = Arc::new(AtomicI32::new(5));
+        let rhs = Arc::clone(&threshold);
+        let expr = dynamic(
+            CompareOperator::Gt,
+            move || Some(ScalarValue::from(rhs.load(Ordering::Relaxed))),
+            PType::I32.into(),
+            true,
+            root(),
+        );
+
+        // Lower once, as the zoned reader does, and evaluate the lowered predicate after each
+        // update of the dynamic value.
+        let lowered = zone_map.lower_stats(falsify(&expr, PType::I32.into()))?;
+
+        // A > 5 => A.max <= 5
+        let mask = zone_map.prune_lowered(&lowered, &SESSION)?;
+        assert_arrays_eq!(
+            mask.into_array(),
+            BoolArray::from_iter([true, false, false]),
+            ctx
+        );
+
+        // A > 6 => A.max <= 6
+        threshold.store(6, Ordering::Relaxed);
+        let mask = zone_map.prune_lowered(&lowered, &SESSION)?;
+        assert_arrays_eq!(
+            mask.into_array(),
+            BoolArray::from_iter([true, true, false]),
+            ctx
+        );
+
+        Ok(())
+    }
 }