Add plan parent-reduction rule API

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<dyn Plan>;
 /// 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;

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;

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<C: Plan>: 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<Option<PlanRef>>;
+}
+
+/// 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<Option<PlanRef>>;
+}
+
+/// Bridges a typed [`PlanParentReduceRule`] to a type-erased static registry.
+pub struct PlanParentReduceRuleAdapter<C, R> {
+    rule: R,
+    _child: PhantomData<fn() -> C>,
+}
+
+impl<C, R> PlanParentReduceRuleAdapter<C, R> {
+    /// Creates an adapter for a typed parent-child rule.
+    pub const fn new(rule: R) -> Self {
+        Self {
+            rule,
+            _child: PhantomData,
+        }
+    }
+}
+
+impl<C, R> Debug for PlanParentReduceRuleAdapter<C, R>
+where
+    C: Plan,
+    R: PlanParentReduceRule<C>,
+{
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("PlanParentReduceRuleAdapter")
+            .field("parent", &type_name::<R::Parent>())
+            .field("child", &type_name::<C>())
+            .field("rule", &self.rule)
+            .finish()
+    }
+}
+
+impl<C, R> DynPlanParentReduceRule for PlanParentReduceRuleAdapter<C, R>
+where
+    C: Plan,
+    R: PlanParentReduceRule<C>,
+{
+    fn matches(&self, child: &dyn Plan, parent: &dyn Plan) -> bool {
+        child.as_any().is::<C>() && parent.as_any().is::<R::Parent>()
+    }
+
+    fn reduce_parent(
+        &self,
+        child: &dyn Plan,
+        parent: &dyn Plan,
+        child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>> {
+        let Some(child) = child.as_any().downcast_ref::<C>() else {
+            return Ok(None);
+        };
+        let Some(parent) = parent.as_any().downcast_ref::<R::Parent>() 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<Option<PlanRef>> {
+        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)
+    }
+}