From 16738683ab7cdba1b17a1e1ace14451fdc9f3605 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 7 Aug 2026 16:30:25 +0100 Subject: [PATCH] Prune plan scans with zoned statistics Signed-off-by: Joe Isaacs --- Cargo.lock | 1 + vortex-layout/src/layouts/zoned/mod.rs | 6 +- vortex-layout/src/layouts/zoned/zone_map.rs | 30 +- vortex-layout/src/plan/lower.rs | 14 +- vortex-layout/src/plan/mod.rs | 1 + vortex-layout/src/plan/optimizer/mod.rs | 5 + vortex-layout/src/plan/plans/mod.rs | 2 + vortex-layout/src/plan/plans/zoned.rs | 374 +++++++++++++++++++- vortex-layout/src/plan/tests.rs | 120 +++++++ vortex-scan-v2/Cargo.toml | 1 + vortex-scan-v2/src/repeated_scan.rs | 4 + vortex-scan-v2/src/scan_builder.rs | 176 ++++++--- vortex-scan-v2/src/splits.rs | 9 +- vortex-scan-v2/src/tasks.rs | 88 +++-- vortex-scan-v2/src/tests.rs | 145 ++++++++ 15 files changed, 876 insertions(+), 100 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3e11af3875c..97cb458dc0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10565,6 +10565,7 @@ version = "0.1.0" dependencies = [ "futures", "itertools 0.14.0", + "parking_lot", "tracing", "tracing-subscriber", "vortex-array", diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..7a954155835 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -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() diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index c84c0b443dd..157f121d4c9 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -16,6 +16,7 @@ use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; 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::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -84,7 +85,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]>, @@ -144,19 +145,32 @@ impl ZoneMap { session: &VortexSession, ) -> VortexResult { 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 { + let mut ctx = session.create_execution_ctx(); + self.applied_predicate(predicate)? + .execute::(&mut ctx) + } + fn applied_predicate(&self, predicate: &BoundExpression) -> VortexResult { + 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 { diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs index 9f1f8748f51..05f735ba134 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -166,5 +166,17 @@ fn lower_zoned(layout: &LayoutRef) -> VortexResult { .slot(1)? .ok_or_else(|| vortex_err!("Zoned zones child is absent"))?, )?; - Ok(ZonedPlan::new(data, zones)) + let metadata = if let Some(layout) = layout.as_opt::() { + layout.data() + } else if let Some(layout) = layout.as_opt::() { + layout.data() + } else { + vortex_bail!("Zoned plan requires a zoned layout") + }; + Ok(ZonedPlan::new( + data, + zones, + u64::try_from(metadata.zone_len())?, + metadata.aggregate_fns(), + )) } diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index d6360587dc6..ab19dcbe1bb 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -55,6 +55,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; diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/optimizer/mod.rs index 32f8ccaca25..7ab202a07e6 100644 --- a/vortex-layout/src/plan/optimizer/mod.rs +++ b/vortex-layout/src/plan/optimizer/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 = PlanParentReduceRuleAdapter::new(ExpressionConcatRule); @@ -29,12 +31,15 @@ static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter = PlanParentReduceRuleAdapter::new(ExpressionPackRule); +static EXPRESSION_ZONED_RULE: PlanParentReduceRuleAdapter = + 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`. diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index ffaf3f75f82..d0af735685e 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/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; diff --git a/vortex-layout/src/plan/plans/zoned.rs b/vortex-layout/src/plan/plans/zoned.rs index 7823f5b4a02..7254739c324 100644 --- a/vortex-layout/src/plan/plans/zoned.rs +++ b/vortex-layout/src/plan/plans/zoned.rs @@ -2,14 +2,39 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::fmt; use std::ops::Range; +use std::sync::Arc; +use std::sync::OnceLock; +use futures::FutureExt; +use futures::TryFutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::Transformed; +use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::fns::stat::StatFn; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_error::SharedVortexResult; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::registry::CachedId; +use crate::layouts::zoned::zone_map::ZoneMap; +use crate::plan::Eval; use crate::plan::Plan; use crate::plan::PlanArrayFuture; use crate::plan::PlanExecutionContext; @@ -18,48 +43,285 @@ use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; use crate::plan::check_child_count; +use crate::plan::optimizer::PlanParentReduceRule; -pub(crate) const DATA: usize = 0; -pub(crate) const ZONES: usize = 1; +const DATA: usize = 0; +const ZONES: usize = 1; + +type SharedZoneMap = Shared>>; + +#[derive(Clone)] +struct ZonedPruningState { + expression: BoundExpression, + zone_map: Arc>, +} + +impl fmt::Debug for ZonedPruningState { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ZonedPruningState") + .field("expression", &self.expression) + .finish_non_exhaustive() + } +} + +impl ZonedPruningState { + fn new(expression: BoundExpression) -> Self { + Self { + expression, + zone_map: Arc::new(OnceLock::new()), + } + } + + fn zone_map( + &self, + ctx: &PlanExecutionContext, + zones: &PlanRef, + column_dtype: &DType, + aggregate_fns: &Arc<[AggregateFnRef]>, + zone_len: u64, + row_count: u64, + ) -> VortexResult { + let zone_count = zones.row_count(); + let zone_count_usize = usize::try_from(zone_count)?; + Ok(self + .zone_map + .get_or_init(|| { + let ctx = ctx.clone(); + let zones = zones.clone(); + let column_dtype = column_dtype.clone(); + let aggregate_fns = Arc::clone(aggregate_fns); + async move { + let zones = zones.execute( + &ctx, + &(0..zone_count), + MaskFuture::new_true(zone_count_usize), + )?; + let mut execution = ctx.session().create_execution_ctx(); + let zones = zones.await?.execute::(&mut execution)?; + // SAFETY: zoned layout construction validated that the auxiliary child was + // written from this column dtype and stats-table schema. + Ok(unsafe { + ZoneMap::new_unchecked( + column_dtype, + zones, + aggregate_fns, + zone_len, + row_count, + ) + }) + } + .map_err(Arc::new) + .boxed() + .shared() + }) + .clone()) + } +} + +/// Zoned-plan-specific data. +#[derive(Clone, Debug)] +pub struct ZonedData { + column_dtype: DType, + zone_len: u64, + aggregate_fns: Arc<[AggregateFnRef]>, + pruning: Option, +} /// Reads data alongside the zone statistics summarising it. /// /// This operator covers both `vortex.zoned` layouts and legacy `vortex.stats` layouts, which have -/// the same physical child shape. +/// the same physical child shape. An expression containing abstract statistic functions can +/// rewrite it into a pruning plan that retains only the zone-statistics child. #[derive(Clone, Debug)] pub struct Zoned; -/// A plan that pairs data with its zone statistics. +/// A plan that pairs data with its zone statistics or evaluates a zone-backed pruning proof. pub type ZonedPlan = Plan; impl ZonedPlan { /// Creates a zoned plan over `data` summarised by `zones`. - pub fn new(data: PlanRef, zones: PlanRef) -> Self { - let dtype: DType = data.dtype().clone(); + pub fn new( + data: PlanRef, + zones: PlanRef, + zone_len: u64, + aggregate_fns: Arc<[AggregateFnRef]>, + ) -> Self { + let dtype = data.dtype().clone(); let row_count = data.row_count(); PlanParts { vtable: Zoned, - dtype, + dtype: dtype.clone(), row_count, children: vec![data, zones], - data: (), + data: ZonedData { + column_dtype: dtype, + zone_len, + aggregate_fns, + pruning: None, + }, } .into_typed() } - /// Returns the plan producing the summarised data. - pub fn data_plan(&self) -> &PlanRef { - &self.children()[DATA] + /// Returns the plan producing the summarised data, unless this is a pruning plan. + pub fn data_plan(&self) -> Option<&PlanRef> { + (!self.is_pruning()) + .then(|| self.children().get(DATA)) + .flatten() } /// Returns the plan producing the zone statistics. pub fn zones_plan(&self) -> &PlanRef { - &self.children()[ZONES] + let index = if self.is_pruning() { 0 } else { ZONES }; + &self.children()[index] + } + + /// Returns whether this plan evaluates a zone-backed pruning proof. + pub fn is_pruning(&self) -> bool { + self.data().pruning.is_some() + } + + /// Returns the abstract pruning proof carried by this plan, when present. + pub fn pruning_expression(&self) -> Option<&BoundExpression> { + self.data().pruning.as_ref().map(|state| &state.expression) + } + + fn with_pruning(&self, expression: BoundExpression) -> Option { + if self.data().zone_len == 0 || self.is_pruning() { + return None; + } + let mut data = self.data().clone(); + data.pruning = Some(ZonedPruningState::new(expression.clone())); + Some( + PlanParts { + vtable: Zoned, + dtype: expression.dtype().clone(), + row_count: self.row_count(), + children: vec![self.zones_plan().clone()], + data, + } + .into_typed(), + ) + } + + fn execute_pruning( + &self, + ctx: &PlanExecutionContext, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + vortex_ensure!( + row_range.start <= row_range.end && row_range.end <= self.row_count(), + "Zoned pruning row range {:?} is outside 0..{}", + row_range, + self.row_count() + ); + let range_len = usize::try_from(row_range.end - row_range.start)?; + vortex_ensure!( + mask.len() == range_len, + "Zoned pruning mask length mismatch" + ); + + let state = self + .data() + .pruning + .clone() + .ok_or_else(|| vortex_error::vortex_err!("Zoned pruning state is absent"))?; + let ctx = ctx.clone(); + let zones = self.zones_plan().clone(); + let column_dtype = self.data().column_dtype.clone(); + let output_dtype = self.dtype().clone(); + let aggregate_fns = Arc::clone(&self.data().aggregate_fns); + let zone_len = self.data().zone_len; + let row_count = self.row_count(); + let row_range = row_range.clone(); + + Ok(async move { + let input_mask = mask.await?; + if input_mask.all_false() { + return Ok(BoolArray::new( + BitBuffer::new_unset(0), + Validity::from(output_dtype.nullability()), + ) + .into_array()); + } + + let zone_map = state.zone_map( + &ctx, + &zones, + &column_dtype, + &aggregate_fns, + zone_len, + row_count, + )?; + let zone_map = zone_map.await?; + let evaluated = zone_map.evaluate(&state.expression, ctx.session())?; + let mut execution = ctx.session().create_execution_ctx(); + let zone_validity = + BoolArrayExt::validity(&evaluated).execute_mask(evaluated.len(), &mut execution)?; + let zone_values = evaluated.to_bit_buffer(); + + let zone_start = row_range.start / zone_len; + let zone_end = row_range.end.div_ceil(zone_len); + let zone_start_usize = usize::try_from(zone_start)?; + let zone_end_usize = usize::try_from(zone_end)?; + vortex_ensure!( + zone_end_usize <= evaluated.len(), + "Zoned pruning requires zones {zone_start}..{zone_end}, but only {} exist", + evaluated.len() + ); + + let mut values = BitBufferMut::with_capacity(range_len); + let mut validity = BitBufferMut::with_capacity(range_len); + let relevant_values = zone_values.slice(zone_start_usize..zone_end_usize); + let relevant_validity = zone_validity.slice(zone_start_usize..zone_end_usize); + for (offset, (value, valid)) in relevant_values + .iter() + .zip(relevant_validity.iter()) + .enumerate() + { + let zone_index = zone_start + u64::try_from(offset)?; + let zone_row_start = zone_index.saturating_mul(zone_len).min(row_count); + let zone_row_end = zone_index + .saturating_add(1) + .saturating_mul(zone_len) + .min(row_count); + let start = zone_row_start.max(row_range.start); + let end = zone_row_end.min(row_range.end); + if start < end { + let len = usize::try_from(end - start)?; + values.append_n(value, len); + validity.append_n(valid, len); + } + } + vortex_ensure!( + values.len() == range_len && validity.len() == range_len, + "Expanded zone proof length does not match row range" + ); + + let validity = if output_dtype.is_nullable() { + Validity::from(validity.freeze()) + } else { + vortex_ensure!( + validity.freeze().true_count() == range_len, + "Non-nullable zoned proof produced null values" + ); + Validity::NonNullable + }; + let output = BoolArray::new(values.freeze(), validity).into_array(); + if input_mask.all_true() { + Ok(output) + } else { + output.filter(input_mask) + } + } + .boxed()) } } impl PlanVTable for Zoned { - type PlanData = (); + type PlanData = ZonedData; type Metadata = EmptyMetadata; fn id(&self) -> PlanId { @@ -67,15 +329,39 @@ impl PlanVTable for Zoned { *ID } + fn fmt(plan: &Plan, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + if let Some(expression) = plan.pruning_expression() { + write!(formatter, " prune={expression}")?; + } + Ok(()) + } + fn metadata(_plan: &Plan) -> Option { - Some(EmptyMetadata) + None } - fn with_children(_plan: &Plan, mut children: Vec) -> VortexResult> { + fn with_children(plan: &Plan, mut children: Vec) -> VortexResult> { + if plan.is_pruning() { + check_child_count("Zoned pruning", &children, 1)?; + return Ok(PlanParts { + vtable: Zoned, + dtype: plan.dtype().clone(), + row_count: plan.row_count(), + children, + data: plan.data().clone(), + } + .into_typed()); + } + check_child_count("Zoned", &children, 2)?; let zones = children.remove(ZONES); let data = children.remove(DATA); - Ok(ZonedPlan::new(data, zones)) + Ok(ZonedPlan::new( + data, + zones, + plan.data().zone_len, + Arc::clone(&plan.data().aggregate_fns), + )) } fn execute( @@ -84,10 +370,22 @@ impl PlanVTable for Zoned { row_range: &Range, mask: MaskFuture, ) -> VortexResult { - plan.data_plan().execute(ctx, row_range, mask) + if plan.is_pruning() { + return plan.execute_pruning(ctx, row_range, mask); + } + plan.data_plan() + .ok_or_else(|| vortex_error::vortex_err!("Zoned data child is absent"))? + .execute(ctx, row_range, mask) } - fn child_name(_plan: &Plan, index: usize) -> Cow<'_, str> { + fn child_name(plan: &Plan, index: usize) -> Cow<'_, str> { + if plan.is_pruning() { + return if index == 0 { + Cow::Borrowed("zones") + } else { + Cow::Owned(format!("child[{index}]")) + }; + } match index { DATA => Cow::Borrowed("data"), ZONES => Cow::Borrowed("zones"), @@ -95,3 +393,43 @@ impl PlanVTable for Zoned { } } } + +/// Rewrites an abstract statistic expression over a zoned plan into its pruning state. +#[derive(Debug)] +pub(crate) struct ExpressionZonedRule; + +impl PlanParentReduceRule for ExpressionZonedRule { + type Parent = Eval; + + fn reduce_parent( + &self, + child: &ZonedPlan, + parent: &Plan, + _child_idx: usize, + ) -> VortexResult> { + let mut contains_stat = false; + let mut contains_root = false; + parent.expression().clone().transform_down(|expression| { + if expression + .as_scalar() + .is_some_and(|scalar_fn| scalar_fn.is::()) + { + contains_stat = true; + return Ok(Transformed { + value: expression, + order: TraversalOrder::Skip, + changed: false, + }); + } + contains_root |= expression.is_root(); + Ok(Transformed::no(expression)) + })?; + if !parent.dtype().is_boolean() || !contains_stat || contains_root { + return Ok(None); + } + + Ok(child + .with_pruning(parent.expression().clone()) + .map(Plan::into_plan)) + } +} diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 6b04026f256..6a9dd1fc6e3 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -10,6 +10,9 @@ use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::max::Max; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -19,6 +22,7 @@ use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; use vortex_array::expr::Expression; use vortex_array::expr::and; +use vortex_array::expr::bound::and as bound_and; use vortex_array::expr::checked_add; use vortex_array::expr::get_item; use vortex_array::expr::gt; @@ -900,6 +904,122 @@ fn zoned_plan_exposes_data_and_zones() -> VortexResult<()> { Ok(()) } +#[test] +fn stats_expression_rewrites_to_zoned_pruning_plan() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let max_dtype = max + .state_dtype(&dtype) + .ok_or_else(|| vortex_err!("max does not support {dtype}"))?; + let zones_dtype = DType::Struct( + StructFields::from_iter([(max.to_string(), max_dtype.as_nullable())]), + Nullability::NonNullable, + ); + let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?; + let layout = ZonedLayout::try_new( + flat(5, dtype.clone(), 0), + flat(2, zones_dtype, 1), + zone_len, + vec![max].into(), + )? + .into_layout(); + let session = vortex_array::array_session(); + let filter = gt(root(), lit(5_i32)); + let falsifier = filter + .bind(&dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("filter has no falsifier"))?; + let source = make_plan(layout)?; + let plan = EvalPlan::new(falsifier.clone(), source.clone()).into_plan(); + + insta::assert_snapshot!(plan.display_tree(), @r" + root: vortex.plan.eval(bool?, rows=5) expr=(stat($, vortex.max()) <= 5i32) + child: vortex.plan.zoned(i32, rows=5) + data: vortex.plan.segment_scan(i32, rows=5) + zones: vortex.plan.segment_scan({vortex.max()=i32?}, rows=2) + "); + + let optimized = optimize(plan)?; + insta::assert_snapshot!(optimized.display_tree(), @r" + root: vortex.plan.zoned(bool?, rows=5) prune=(stat($, vortex.max()) <= 5i32) + zones: vortex.plan.segment_scan({vortex.max()=i32?}, rows=2) + "); + let zoned = optimized + .as_opt::() + .ok_or_else(|| vortex_err!("optimized pruning plan is not zoned"))?; + assert!(zoned.is_pruning()); + assert_eq!(zoned.pruning_expression(), Some(&falsifier)); + assert!(zoned.data_plan().is_none()); + assert_eq!(zoned.children().len(), 1); + + let mixed_expression = bound_and(falsifier, gt(root(), lit(0_i32)).bind(&dtype)?); + let mixed = optimize(EvalPlan::new(mixed_expression, source).into_plan())?; + let mixed = mixed + .as_opt::() + .ok_or_else(|| vortex_err!("expression with a data reference was pushed into zones"))?; + assert!(mixed.child_plan().is::()); + assert!( + !mixed + .child_plan() + .as_opt::() + .ok_or_else(|| vortex_err!("mixed expression child is not zoned"))? + .is_pruning() + ); + Ok(()) +} + +#[test] +fn pruning_expression_partitions_across_row_idx_and_zoned_struct_field() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let max_dtype = max + .state_dtype(&value_dtype) + .ok_or_else(|| vortex_err!("max does not support {value_dtype}"))?; + let zones_dtype = DType::Struct( + StructFields::from_iter([(max.to_string(), max_dtype.as_nullable())]), + Nullability::NonNullable, + ); + let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?; + let zoned = ZonedLayout::try_new( + flat(5, value_dtype.clone(), 0), + flat(2, zones_dtype, 1), + zone_len, + vec![max].into(), + )? + .into_layout(); + let struct_dtype = DType::Struct( + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), + Nullability::NonNullable, + ); + let layout = StructLayout::new( + 5, + struct_dtype.clone(), + vec![zoned, flat(5, value_dtype, 2)], + ) + .into_layout(); + let source = RowIdxPlan::new(10, make_plan(layout)?).into_plan(); + let session = vortex_array::array_session(); + let filter = and( + gt(row_idx(), lit(11_u64)), + gt(get_item("a", root()), lit(5_i32)), + ); + let falsifier = filter + .bind(&struct_dtype)? + .falsify(&session)? + .ok_or_else(|| vortex_err!("filter has no falsifier"))?; + + let optimized = optimize(EvalPlan::new(falsifier, source).into_plan())?; + insta::assert_snapshot!(optimized.display_tree(), @r" + root: vortex.plan.eval(bool?, rows=5) expr=($.row_idx or $.child) + child: vortex.plan.row_idx_partition({row_idx=bool?, child=bool?}, rows=5) + row_idx: vortex.plan.eval(bool?, rows=5) expr=(stat($, vortex.max()) <= 11u64) + child: vortex.plan.row_idx_values(u64, rows=5) + child: vortex.plan.zoned(bool?, rows=5) prune=(stat($, vortex.max()) <= 5i32) + zones: vortex.plan.segment_scan({vortex.max()=i32?}, rows=2) + "); + Ok(()) +} + #[test] fn legacy_stats_layout_uses_zoned_plan() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-scan-v2/Cargo.toml b/vortex-scan-v2/Cargo.toml index 8f4e97cad0e..cece44f9730 100644 --- a/vortex-scan-v2/Cargo.toml +++ b/vortex-scan-v2/Cargo.toml @@ -28,6 +28,7 @@ vortex-session = { workspace = true } vortex-utils = { workspace = true } [dev-dependencies] +parking_lot = { workspace = true } tracing-subscriber = { workspace = true, features = ["env-filter"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-file = { workspace = true, features = ["tokio"] } diff --git a/vortex-scan-v2/src/repeated_scan.rs b/vortex-scan-v2/src/repeated_scan.rs index 8e51097c78e..50016a1d98d 100644 --- a/vortex-scan-v2/src/repeated_scan.rs +++ b/vortex-scan-v2/src/repeated_scan.rs @@ -33,6 +33,7 @@ use crate::tasks::split_exec; pub struct RepeatedScan { execution: PlanExecutionContext, projection: PlanRef, + pruning: Option, filter: Option, ordered: bool, row_range: Option>, @@ -80,6 +81,7 @@ impl RepeatedScan { pub(crate) fn new( execution: PlanExecutionContext, projection: PlanRef, + pruning: Option, filter: Option, ordered: bool, row_range: Option>, @@ -93,6 +95,7 @@ impl RepeatedScan { Self { execution, projection, + pruning, filter, ordered, row_range, @@ -159,6 +162,7 @@ impl RepeatedScan { let ctx = Arc::new(TaskContext { execution: self.execution.clone(), + pruning: self.pruning.clone(), filter: self.filter.clone(), projection: self.projection.clone(), mapper: Arc::clone(&self.map_fn), diff --git a/vortex-scan-v2/src/scan_builder.rs b/vortex-scan-v2/src/scan_builder.rs index be893560e53..9bed993c72b 100644 --- a/vortex-scan-v2/src/scan_builder.rs +++ b/vortex-scan-v2/src/scan_builder.rs @@ -14,6 +14,7 @@ use futures::future::BoxFuture; use futures::stream::BoxStream; use vortex_array::ArrayRef; use vortex_array::dtype::DType; +use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::root; use vortex_array::iter::ArrayIterator; @@ -33,6 +34,8 @@ use vortex_layout::plan::EvalPlan; use vortex_layout::plan::PlanExecutionContext; use vortex_layout::plan::PlanRef; use vortex_layout::plan::RowIdxPlan; +use vortex_layout::plan::RowIdxValues; +use vortex_layout::plan::Zoned; use vortex_layout::plan::lower; use vortex_layout::plan::optimize; use vortex_layout::segments::SegmentSource; @@ -253,50 +256,14 @@ impl ScanBuilder { plan = %source.display_tree(), "planning expressions over the row-index-aware source" ); - tracing::debug!( - target: "vortex_scan_v2::planner", - expression = %self.projection, - "optimizing the projection expression" - ); - let projection = self - .projection - .optimize_recursive(source.dtype())? - .bind(source.dtype())?; - let projection: PlanRef = EvalPlan::new(projection, source.clone()).into_plan(); - let projection = optimize(projection)?; - tracing::debug!( - target: "vortex_scan_v2::planner", - plan = %projection.display_tree(), - "optimized the projection physical plan" - ); - - let filter = self - .filter - .map(|filter| -> VortexResult { - tracing::debug!( - target: "vortex_scan_v2::planner", - expression = %filter, - "optimizing the filter expression" - ); - let filter = filter - .optimize_recursive(source.dtype())? - .bind(source.dtype())?; - let filter: PlanRef = EvalPlan::new(filter, source.clone()).into_plan(); - let filter = optimize(filter)?; - vortex_ensure!( - filter.dtype().is_boolean(), - "Filter plan must produce booleans" - ); - Ok(filter) - }) - .transpose()?; - if let Some(filter) = &filter { - tracing::debug!( - target: "vortex_scan_v2::planner", - plan = %filter.display_tree(), - "optimized the filter physical plan" - ); - } + let projection = optimize_projection_plan(self.projection, &source)?; + let filter_expression = optimize_filter_expression(self.filter, &source)?; + let pruning = optimize_pruning_plan( + filter_expression.as_ref(), + &source, + self.execution.session(), + )?; + let filter = optimize_filter_plan(filter_expression.as_ref(), &source)?; let splits = if let Some(ranges) = attempt_split_ranges(&self.selection, self.row_range.as_ref()) { @@ -308,6 +275,7 @@ impl ScanBuilder { .unwrap_or_else(|| 0..self.base_plan.row_count()); let mut plans = vec![&projection]; plans.extend(filter.as_ref()); + plans.extend(pruning.as_ref()); Splits::Natural(self.split_by.splits(&plans, &row_range)?) }; match &splits { @@ -328,6 +296,7 @@ impl ScanBuilder { Ok(RepeatedScan::new( self.execution, projection, + pruning, filter, self.ordered, self.row_range, @@ -363,6 +332,125 @@ impl ScanBuilder { } } +fn optimize_projection_plan(expression: Expression, source: &PlanRef) -> VortexResult { + tracing::debug!( + target: "vortex_scan_v2::planner", + %expression, + "optimizing the projection expression" + ); + let expression = expression + .optimize_recursive(source.dtype())? + .bind(source.dtype())?; + let projection = optimize(EvalPlan::new(expression, source.clone()).into_plan())?; + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %projection.display_tree(), + "optimized the projection physical plan" + ); + Ok(projection) +} + +fn optimize_filter_expression( + filter: Option, + source: &PlanRef, +) -> VortexResult> { + filter + .map(|expression| { + tracing::debug!( + target: "vortex_scan_v2::planner", + %expression, + "optimizing the filter expression" + ); + expression + .optimize_recursive(source.dtype())? + .bind(source.dtype()) + }) + .transpose() +} + +fn optimize_pruning_plan( + filter: Option<&BoundExpression>, + source: &PlanRef, + session: &VortexSession, +) -> VortexResult> { + let Some(pruning) = build_pruning_plan(filter, source, session)? else { + return Ok(None); + }; + if !uses_only_pruning_sources(&pruning) { + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %pruning.display_tree(), + "discarding a pruning plan that still requires data values" + ); + return Ok(None); + } + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %pruning.display_tree(), + "optimized the pruning physical plan" + ); + Ok(Some(pruning)) +} + +fn build_pruning_plan( + filter: Option<&BoundExpression>, + source: &PlanRef, + session: &VortexSession, +) -> VortexResult> { + let Some(filter) = filter else { + return Ok(None); + }; + let Some(falsifier) = filter.falsify(session)? else { + tracing::debug!( + target: "vortex_scan_v2::planner", + expression = %filter, + "filter has no statistics falsifier" + ); + return Ok(None); + }; + tracing::debug!( + target: "vortex_scan_v2::planner", + expression = %falsifier, + "optimizing the pruning expression" + ); + let pruning = optimize(EvalPlan::new(falsifier, source.clone()).into_plan())?; + vortex_ensure!( + pruning.dtype().is_boolean(), + "Pruning plan must produce booleans" + ); + Ok(Some(pruning)) +} + +fn optimize_filter_plan( + filter: Option<&BoundExpression>, + source: &PlanRef, +) -> VortexResult> { + let Some(expression) = filter else { + return Ok(None); + }; + let filter = optimize(EvalPlan::new(expression.clone(), source.clone()).into_plan())?; + vortex_ensure!( + filter.dtype().is_boolean(), + "Filter plan must produce booleans" + ); + tracing::debug!( + target: "vortex_scan_v2::planner", + plan = %filter.display_tree(), + "optimized the filter physical plan" + ); + Ok(Some(filter)) +} + +fn uses_only_pruning_sources(plan: &PlanRef) -> bool { + if let Some(zoned) = plan.as_opt::() { + return zoned.is_pruning(); + } + if plan.is::() { + return true; + } + !plan.children().is_empty() && plan.children().iter().all(uses_only_pruning_sources) +} + enum LazyScanState { Builder(Option>>), Preparing(PreparingScan), diff --git a/vortex-scan-v2/src/splits.rs b/vortex-scan-v2/src/splits.rs index b55ab31d127..aec90500cff 100644 --- a/vortex-scan-v2/src/splits.rs +++ b/vortex-scan-v2/src/splits.rs @@ -65,7 +65,14 @@ fn collect_plan_splits( row_range: &Range, boundaries: &mut Vec, ) -> VortexResult<()> { - if plan.is::() || plan.is::() || plan.is::() { + if let Some(zoned) = plan.as_opt::() { + if let Some(data) = zoned.data_plan() { + collect_plan_splits(data, row_offset, row_range, boundaries)?; + } + return Ok(()); + } + + if plan.is::() || plan.is::() { if let Some(child) = plan.child(0) { collect_plan_splits(child, row_offset, row_range, boundaries)?; } diff --git a/vortex-scan-v2/src/tasks.rs b/vortex-scan-v2/src/tasks.rs index e3fecf928a8..3b8dd701875 100644 --- a/vortex-scan-v2/src/tasks.rs +++ b/vortex-scan-v2/src/tasks.rs @@ -27,28 +27,57 @@ pub(crate) fn split_exec( target: "vortex_scan_v2::execution", ?row_range, selected_rows = row_mask.true_count(), + has_pruning = ctx.pruning.is_some(), has_filter = ctx.filter.is_some(), "executing a plan scan split" ); - let filter_mask = match &ctx.filter { - None => { - let row_mask = match limit { - Some(limit) if *limit == 0 => Mask::new_false(row_mask.len()), - Some(limit) => { - let true_count = row_mask.true_count(); - let mask_limit = usize::try_from(*limit) - .map(|limit| limit.min(true_count)) - .unwrap_or(true_count); - let row_mask = row_mask.limit(mask_limit); - *limit -= mask_limit as u64; - row_mask - } - None => row_mask, - }; - MaskFuture::ready(row_mask) + let row_mask = match (&ctx.filter, limit) { + (None, Some(limit)) if *limit == 0 => Mask::new_false(row_mask.len()), + (None, Some(limit)) => { + let true_count = row_mask.true_count(); + let mask_limit = usize::try_from(*limit) + .map(|limit| limit.min(true_count)) + .unwrap_or(true_count); + let row_mask = row_mask.limit(mask_limit); + *limit -= mask_limit as u64; + row_mask + } + _ => row_mask, + }; + + Ok(async move { + let mut row_mask = row_mask; + if let Some(pruning) = &ctx.pruning { + let proof = pruning.execute( + &ctx.execution, + &row_range, + MaskFuture::ready(row_mask.clone()), + )?; + let proof = proof.await?; + let mut execution = ctx.execution.session().create_execution_ctx(); + let pruned: Mask = proof.null_as_false().execute(&mut execution)?; + let pruned_rows = pruned.true_count(); + row_mask = row_mask.intersect_by_rank(&!pruned); + tracing::trace!( + target: "vortex_scan_v2::execution", + ?row_range, + pruned_rows, + remaining_rows = row_mask.true_count(), + "applied the plan pruning proof" + ); } - Some(filter) => { + + if row_mask.all_false() { + tracing::trace!( + target: "vortex_scan_v2::execution", + ?row_range, + "plan pruning skipped the scan split" + ); + return Ok(None); + } + + let filter_mask = if let Some(filter) = &ctx.filter { let predicate = filter.execute( &ctx.execution, &row_range, @@ -58,18 +87,21 @@ pub(crate) fn split_exec( MaskFuture::new(row_mask.len(), async move { let predicate = predicate.await?; let mut execution = session.create_execution_ctx(); - let predicate = predicate.null_as_false().execute(&mut execution)?; + let predicate: Mask = predicate.null_as_false().execute(&mut execution)?; Ok(row_mask.intersect_by_rank(&predicate)) }) - } - }; + } else { + MaskFuture::ready(row_mask) + }; - let projection = ctx - .projection - .execute(&ctx.execution, &row_range, filter_mask.clone())?; - let mapper = Arc::clone(&ctx.mapper); - Ok(async move { - if filter_mask.await?.all_false() { + // Register projection reads before resolving the filter mask so segments used by both + // expressions can share the same in-flight request. + let projection = ctx + .projection + .execute(&ctx.execution, &row_range, filter_mask.clone())?; + let row_mask = filter_mask.await?; + + if row_mask.all_false() { tracing::trace!( target: "vortex_scan_v2::execution", ?row_range, @@ -77,6 +109,7 @@ pub(crate) fn split_exec( ); return Ok(None); } + let array = projection.await?; tracing::trace!( target: "vortex_scan_v2::execution", @@ -85,13 +118,14 @@ pub(crate) fn split_exec( dtype = %array.dtype(), "completed a plan scan split" ); - mapper(array).map(Some) + (ctx.mapper)(array).map(Some) } .boxed()) } pub(crate) struct TaskContext { pub(crate) execution: PlanExecutionContext, + pub(crate) pruning: Option, pub(crate) filter: Option, pub(crate) projection: PlanRef, pub(crate) mapper: Arc VortexResult + Send + Sync>, diff --git a/vortex-scan-v2/src/tests.rs b/vortex-scan-v2/src/tests.rs index 155675ce216..e058727c0a6 100644 --- a/vortex-scan-v2/src/tests.rs +++ b/vortex-scan-v2/src/tests.rs @@ -1,12 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::num::NonZeroUsize; use std::sync::Arc; +use parking_lot::Mutex; use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -21,13 +24,22 @@ use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex_layout::layouts::flat::Flat; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::row_idx::row_idx; use vortex_layout::layouts::table::TableStrategy; +use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; +use vortex_layout::layouts::zoned::writer::ZonedStrategy; +use vortex_layout::segments::SegmentFuture; +use vortex_layout::segments::SegmentId; +use vortex_layout::segments::SegmentSource; +use vortex_layout::segments::SharedSegmentSource; use vortex_layout::segments::TestSegments; use vortex_layout::sequence::SequenceId; use vortex_layout::sequence::SequentialArrayStreamExt; @@ -37,6 +49,32 @@ use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; use crate::ScanBuilder; use crate::SplitBy; +#[derive(Clone)] +struct TrackingSource { + inner: Arc, + requests: Arc>>, +} + +impl TrackingSource { + fn new(inner: Arc) -> Self { + Self { + inner, + requests: Arc::new(Mutex::new(Vec::new())), + } + } + + fn requests(&self) -> Vec { + self.requests.lock().clone() + } +} + +impl SegmentSource for TrackingSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.requests.lock().push(id); + self.inner.request(id) + } +} + #[test] fn scans_layout_through_optimized_plans() -> VortexResult<()> { block_on(|handle| async { @@ -71,6 +109,113 @@ fn scans_layout_through_optimized_plans() -> VortexResult<()> { }) } +#[test] +fn filter_and_projection_share_flat_segment_request() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = PrimitiveArray::from_iter(0_i32..10).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + let data_segment = layout.as_::().segment_id(); + let tracking = TrackingSource::new(segments); + let source: Arc = Arc::new(SharedSegmentSource::new(tracking.clone())); + + let actual = ScanBuilder::try_new(&layout, source, session.clone())? + .with_filter(gt(root(), lit(4_i32))) + .with_projection(checked_add(root(), lit(1_i32))) + .into_array_stream()? + .read_all() + .await?; + let expected = PrimitiveArray::from_iter(6_i32..11).into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + assert_eq!( + tracking + .requests() + .into_iter() + .filter(|&segment| segment == data_segment) + .count(), + 1 + ); + Ok(()) + }) +} + +#[test] +fn zoned_pruning_skips_a_falsified_data_chunk() -> VortexResult<()> { + block_on(|handle| async { + let session = array_session() + .with::() + .with::() + .with_handle(handle); + let segments = Arc::new(TestSegments::default()); + let (sequence, eof) = SequenceId::root().split(); + let input = ChunkedArray::from_iter([ + buffer![1_i32, 2, 3].into_array(), + buffer![4_i32, 5, 6].into_array(), + buffer![7_i32, 8].into_array(), + ]) + .into_array(); + let strategy = ZonedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + FlatLayoutStrategy::default(), + ZonedLayoutOptions { + block_size: NonZeroUsize::new(3) + .ok_or_else(|| vortex_err!("zone length is zero"))?, + ..Default::default() + }, + ); + let layout = strategy + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + input.to_array_stream().sequenced(sequence), + eof, + &session, + ) + .await?; + + let data = layout + .slot(0)? + .ok_or_else(|| vortex_err!("zoned data child is absent"))?; + let first_data = data + .slot(0)? + .ok_or_else(|| vortex_err!("first data chunk is absent"))?; + let first_data_segment = first_data.as_::().segment_id(); + let zones = layout + .slot(1)? + .ok_or_else(|| vortex_err!("zoned stats child is absent"))?; + let zones_segment = zones.as_::().segment_id(); + + let tracking = Arc::new(TrackingSource::new(segments)); + let source: Arc = Arc::clone(&tracking) as Arc; + let actual = ScanBuilder::try_new(&layout, source, session.clone())? + .with_filter(gt(root(), lit(5_i32))) + .into_array_stream()? + .read_all() + .await?; + let expected = PrimitiveArray::from_iter(6_i32..9).into_array(); + + assert_arrays_eq!(actual, expected, &mut session.create_execution_ctx()); + let requests = tracking.requests(); + assert!(requests.contains(&zones_segment)); + assert!(!requests.contains(&first_data_segment)); + Ok(()) + }) +} + #[test] fn scans_row_idx_and_struct_expression_partitions() -> VortexResult<()> { block_on(|handle| async {