diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 700aab86ce4..c54a43fbaba 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/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, 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()); diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 805bc1848bb..dfdb421814b 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -145,13 +145,26 @@ impl ZoneMap { &self, predicate: &BoundExpression, session: &VortexSession, + ) -> VortexResult { + 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 { 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 { + /// Rewrite the stats references in `predicate` into expressions over the stats table. + pub(super) fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { 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(()) + } }