diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..1dd793e2d76 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -358,6 +358,10 @@ impl ZonedLayout { } impl ZonedData { + pub(crate) fn zone_len(&self) -> usize { + self.zone_len + } + fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> { match &self.zone_map_schema { ZoneMapSchema::LegacyStats(stats) => stats diff --git a/vortex-layout/src/plan/lower.rs b/vortex-layout/src/plan/lower.rs index e8911734549..7aa8c41ea55 100644 --- a/vortex-layout/src/plan/lower.rs +++ b/vortex-layout/src/plan/lower.rs @@ -6,6 +6,8 @@ //! This module is only used to build physical-plan fixtures for tests. It is not a production //! planning API. +use std::sync::Arc; + use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -24,6 +26,8 @@ use crate::layouts::list::OFFSETS_CHILD_INDEX; use crate::layouts::list::VALIDITY_CHILD_INDEX; use crate::layouts::struct_::Struct; use crate::layouts::struct_::StructLayout; +use crate::layouts::zoned::LegacyStats; +use crate::layouts::zoned::Zoned; use crate::plan::ConcatPlan; use crate::plan::ListPackPlan; use crate::plan::PackPlan; @@ -31,6 +35,7 @@ use crate::plan::PlanChildren; use crate::plan::PlanRef; use crate::plan::SegmentScanPlan; use crate::plan::TakePlan; +use crate::plan::ZonedPlan; /// Constructs a physical-plan fixture from `layout` for tests. /// @@ -52,6 +57,9 @@ pub fn lower(layout: &LayoutRef) -> VortexResult { if let Some(layout) = layout.as_opt::() { return Ok(lower_list(layout)?.into_plan()); } + if layout.is::() || layout.is::() { + return Ok(lower_zoned(layout)?.into_plan()); + } vortex_bail!( "No physical plan implementation for layout '{}'", layout.encoding_id() @@ -151,3 +159,20 @@ fn lazy_children(layout: LayoutRef, slots: Vec) -> PlanChildren { lower(&child) }) } + +fn lower_zoned(layout: &LayoutRef) -> VortexResult { + // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary 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::from_children( + layout.dtype().clone(), + layout.row_count(), + lazy_children(Arc::clone(layout), vec![0, 1]), + u64::try_from(metadata.zone_len())?, + )) +} diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 802e210b07e..39fffd924d5 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -44,6 +44,9 @@ pub use plans::SegmentScanData; pub use plans::SegmentScanPlan; pub use plans::Take; pub use plans::TakePlan; +pub use plans::Zoned; +pub use plans::ZonedData; +pub use plans::ZonedPlan; pub use plans::plan_row_idx_expression; 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 a63006be62e..6355a97d3b1 100644 --- a/vortex-layout/src/plan/optimizer/mod.rs +++ b/vortex-layout/src/plan/optimizer/mod.rs @@ -20,10 +20,12 @@ use super::Eval; use super::Pack; use super::PlanRef; use super::Take; +use super::Zoned; use super::plans::EvalIdentityRule; use super::plans::ExpressionConcatRule; use super::plans::ExpressionPackRule; use super::plans::ExpressionTakeRule; +use super::plans::ExpressionZonedRule; static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter = PlanReduceRuleAdapter::new(EvalIdentityRule); @@ -36,11 +38,14 @@ static EXPRESSION_TAKE_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_PACK_RULE, + &EXPRESSION_ZONED_RULE, ]); /// Attempts a static rewrite for `plan`. diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index 7b6a1893d0b..f24381c8951 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -8,6 +8,7 @@ mod pack; mod row_idx; mod segment_scan; mod take; +mod zoned; pub use concat::Concat; pub use concat::ConcatData; @@ -35,3 +36,7 @@ pub use segment_scan::SegmentScanPlan; pub(crate) use take::ExpressionTakeRule; pub use take::Take; 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 new file mode 100644 index 00000000000..f458e093ce5 --- /dev/null +++ b/vortex-layout/src/plan/plans/zoned.rs @@ -0,0 +1,218 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::borrow::Cow; +use std::fmt; + +use vortex_array::EmptyMetadata; +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_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::plan::Eval; +use crate::plan::Plan; +use crate::plan::PlanChildren; +use crate::plan::PlanId; +use crate::plan::PlanParts; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; +use crate::plan::check_child_count; +use crate::plan::optimizer::PlanParentReduceRule; + +const DATA: usize = 0; +const ZONES: usize = 1; + +#[derive(Clone, Debug)] +struct ZonedPruningState { + expression: BoundExpression, +} + +/// Zoned-plan-specific data. +#[derive(Clone, Debug)] +pub struct ZonedData { + zone_len: u64, + 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. 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 or represents a zone-backed pruning proof. +pub type ZonedPlan = Plan; + +impl ZonedPlan { + pub(crate) fn from_children( + dtype: DType, + row_count: u64, + children: PlanChildren, + zone_len: u64, + ) -> Self { + PlanParts { + vtable: Zoned, + dtype, + row_count, + children, + data: ZonedData { + zone_len, + pruning: None, + }, + } + .into_typed() + } + + /// Creates a zoned plan over `data` summarised by `zones` of `zone_len` rows each. + pub fn new(data: PlanRef, zones: PlanRef, zone_len: u64) -> Self { + let dtype = data.dtype().clone(); + let row_count = data.row_count(); + Self::from_children(dtype, row_count, vec![data, zones].into(), zone_len) + } + + /// Returns the plan producing the summarised data, unless this is a pruning plan. + pub fn data_plan(&self) -> VortexResult> { + if self.is_pruning() { + return Ok(None); + } + self.child(DATA) + } + + /// Returns the plan producing the zone statistics. + pub fn zones_plan(&self) -> VortexResult { + let index = if self.is_pruning() { 0 } else { ZONES }; + self.child_required(index) + } + + /// Returns whether this plan represents 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) -> VortexResult> { + if self.data().zone_len == 0 || self.is_pruning() { + return Ok(None); + } + let mut data = self.data().clone(); + data.pruning = Some(ZonedPruningState { + expression: expression.clone(), + }); + Ok(Some( + PlanParts { + vtable: Zoned, + dtype: expression.dtype().clone(), + row_count: self.row_count(), + children: vec![self.zones_plan()?].into(), + data, + } + .into_typed(), + )) + } +} + +impl PlanVTable for Zoned { + type PlanData = ZonedData; + type Metadata = EmptyMetadata; + + fn id(&self) -> PlanId { + static ID: CachedId = CachedId::new("vortex.plan.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 { + None + } + + fn with_children( + plan: &Plan, + children: &PlanChildren, + _data: &mut Self::PlanData, + ) -> VortexResult<()> { + if plan.is_pruning() { + return check_child_count("Zoned pruning", children, 1); + } + + check_child_count("Zoned", children, 2)?; + let data = children + .get(DATA)? + .ok_or_else(|| vortex_error::vortex_err!("Zoned data child is absent"))?; + if data.dtype() != plan.dtype() || data.row_count() != plan.row_count() { + vortex_error::vortex_bail!("Zoned data child shape does not match the plan output"); + } + Ok(()) + } + + 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"), + _ => Cow::Owned(format!("child[{index}]")), + } + } +} + +/// 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 2b87f7ae3c5..b25d8fd3ca3 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -2,14 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fmt; +use std::num::NonZeroUsize; use std::sync::Arc; +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::dtype::DType; use vortex_array::dtype::Nullability; 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; @@ -23,6 +29,8 @@ use vortex_session::registry::CachedId; use vortex_session::registry::ReadContext; use super::*; +use crate::LayoutBuildContext; +use crate::LayoutEncoding; use crate::LayoutRef; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; @@ -32,6 +40,8 @@ use crate::layouts::foreign::new_foreign_layout; use crate::layouts::list::ListLayout; use crate::layouts::row_idx::row_idx; use crate::layouts::struct_::StructLayout; +use crate::layouts::zoned::LegacyStatsLayoutEncoding; +use crate::layouts::zoned::ZonedLayout; use crate::segments::SegmentId; fn primitive(ptype: PType, nullability: Nullability) -> DType { @@ -1085,3 +1095,177 @@ fn nullable_struct_keeps_expression_above_parent_validity() -> VortexResult<()> assert!(eval.child_plan()?.is::()); Ok(()) } + +#[test] +fn zoned_plan_exposes_data_and_zones() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let zones_dtype = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let zone_len = NonZeroUsize::new(3).ok_or_else(|| vortex_err!("zone length is zero"))?; + let aggregate_fns: Arc<[AggregateFnRef]> = Vec::new().into(); + let layout = ZonedLayout::try_new( + flat(5, dtype, 0), + flat(2, zones_dtype, 1), + zone_len, + aggregate_fns, + )? + .into_layout(); + + let plan = make_plan(layout)?; + assert!(plan.is::()); + insta::assert_snapshot!(plan.display_tree(), @" + root: vortex.plan.zoned(i32, rows=5) + data: vortex.plan.segment_scan(i32, rows=5) + zones: vortex.plan.segment_scan({}, rows=2) + "); + 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::try_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::try_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"))?; + let mixed_child = mixed.child_plan()?; + assert!(mixed_child.is::()); + assert!( + !mixed_child + .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 = make_plan(layout)?; + 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(plan_row_idx_expression(falsifier, source)?)?; + insta::assert_snapshot!(optimized.display_tree(), @r" + root: vortex.plan.eval(bool?, rows=5) expr=($.row_idx or $.child) + child: vortex.plan.pack({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(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); + let zones_dtype = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let children = OwnedLayoutChildren::layout_children(vec![ + flat(5, dtype.clone(), 0), + flat(2, zones_dtype, 1), + ]); + let session = vortex_array::array_session(); + let read_ctx = ReadContext::new([]); + let build_ctx = LayoutBuildContext { + session: &session, + array_read_ctx: &read_ctx, + }; + let layout = LayoutEncoding::build( + &LegacyStatsLayoutEncoding, + &dtype, + 5, + &3_u32.to_le_bytes(), + Vec::new(), + children.as_ref(), + &build_ctx, + )?; + + let plan = make_plan(layout)?; + assert!(plan.is::()); + insta::assert_snapshot!(plan.display_tree(), @" + root: vortex.plan.zoned(i32, rows=5) + data: vortex.plan.segment_scan(i32, rows=5) + zones: vortex.plan.segment_scan({}, rows=2) + "); + Ok(()) +}