diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index bb61f021da0..4ccd5a9f325 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -5,6 +5,7 @@ mod children; mod display; +pub mod optimizer; mod plans; use std::any::Any; @@ -44,6 +45,8 @@ pub type PlanRef = Arc; /// Layout plans expose their optimizer-facing children in a stable logical order. Optional child /// slots count toward [`child_count`](Self::child_count) and are returned as `None` by /// [`child`](Self::child) when absent. Accessing a child may initialize and cache its plan. +/// Parent-child rewrites are expressed as [`optimizer::PlanParentReduceRule`]s and collected in a +/// static [`optimizer::PlanParentRuleSet`]. pub trait Plan: 'static + Send + Sync { /// Returns this plan as [`Any`] for plan-specific optimization rules. fn as_any(&self) -> &dyn Any; diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/optimizer/mod.rs new file mode 100644 index 00000000000..c9df9fb547e --- /dev/null +++ b/vortex-layout/src/plan/optimizer/mod.rs @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Static parent-child rewrite rules for physical plans. + +mod rules; + +pub use rules::DynPlanParentReduceRule; +pub use rules::PlanParentReduceRule; +pub use rules::PlanParentReduceRuleAdapter; +pub use rules::PlanParentRuleSet; diff --git a/vortex-layout/src/plan/optimizer/rules.rs b/vortex-layout/src/plan/optimizer/rules.rs new file mode 100644 index 00000000000..f87089ea1fc --- /dev/null +++ b/vortex-layout/src/plan/optimizer/rules.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Typed and type-erased interfaces for parent-child plan rewrites. + +use std::any::type_name; +use std::fmt::Debug; +use std::marker::PhantomData; + +use vortex_error::VortexResult; + +use crate::plan::Plan; +use crate::plan::PlanRef; + +/// A metadata-only rewrite where a child plan rewrites its parent plan. +pub trait PlanParentReduceRule: Debug + Send + Sync + 'static { + /// The concrete parent plan matched by this rule. + type Parent: Plan; + + /// Attempts to replace `parent` based on its child at `child_idx`. + fn reduce_parent( + &self, + child: &C, + parent: &Self::Parent, + child_idx: usize, + ) -> VortexResult>; +} + +/// Type-erased interface used by [`PlanParentRuleSet`]. +pub trait DynPlanParentReduceRule: Debug + Send + Sync + 'static { + /// Returns whether this rule supports the concrete child and parent plan types. + fn matches(&self, child: &dyn Plan, parent: &dyn Plan) -> bool; + + /// Attempts to replace `parent` based on `child` at `child_idx`. + fn reduce_parent( + &self, + child: &dyn Plan, + parent: &dyn Plan, + child_idx: usize, + ) -> VortexResult>; +} + +/// Bridges a typed [`PlanParentReduceRule`] to a type-erased static registry. +pub struct PlanParentReduceRuleAdapter { + rule: R, + _child: PhantomData C>, +} + +impl PlanParentReduceRuleAdapter { + /// Creates an adapter for a typed parent-child rule. + pub const fn new(rule: R) -> Self { + Self { + rule, + _child: PhantomData, + } + } +} + +impl Debug for PlanParentReduceRuleAdapter +where + C: Plan, + R: PlanParentReduceRule, +{ + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PlanParentReduceRuleAdapter") + .field("parent", &type_name::()) + .field("child", &type_name::()) + .field("rule", &self.rule) + .finish() + } +} + +impl DynPlanParentReduceRule for PlanParentReduceRuleAdapter +where + C: Plan, + R: PlanParentReduceRule, +{ + fn matches(&self, child: &dyn Plan, parent: &dyn Plan) -> bool { + child.as_any().is::() && parent.as_any().is::() + } + + fn reduce_parent( + &self, + child: &dyn Plan, + parent: &dyn Plan, + child_idx: usize, + ) -> VortexResult> { + let Some(child) = child.as_any().downcast_ref::() else { + return Ok(None); + }; + let Some(parent) = parent.as_any().downcast_ref::() else { + return Ok(None); + }; + self.rule.reduce_parent(child, parent, child_idx) + } +} + +/// An ordered static collection of parent-child plan rewrite rules. +pub struct PlanParentRuleSet { + rules: &'static [&'static dyn DynPlanParentReduceRule], +} + +impl PlanParentRuleSet { + /// Creates a rule set whose first successful rewrite wins. + pub const fn new(rules: &'static [&'static dyn DynPlanParentReduceRule]) -> Self { + Self { rules } + } + + /// Evaluates rules registered for the concrete `(parent, child)` pair. + pub fn evaluate( + &self, + child: &PlanRef, + parent: &PlanRef, + child_idx: usize, + ) -> VortexResult> { + for rule in self.rules { + if !rule.matches(child.as_ref(), parent.as_ref()) { + continue; + } + let Some(reduced) = rule.reduce_parent(child.as_ref(), parent.as_ref(), child_idx)? + else { + continue; + }; + + #[cfg(debug_assertions)] + { + vortex_error::vortex_ensure!( + reduced.row_count() == parent.row_count(), + "Plan rewrite from {rule:?} changed row count from {} to {}", + parent.row_count(), + reduced.row_count() + ); + vortex_error::vortex_ensure!( + reduced.dtype() == parent.dtype(), + "Plan rewrite from {rule:?} changed dtype from {} to {}", + parent.dtype(), + reduced.dtype() + ); + } + + return Ok(Some(reduced)); + } + Ok(None) + } +}