wip

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/expr/session.rs

@@ -3,6 +3,7 @@
 
 use std::sync::Arc;
 
+use vortex_error::VortexResult;
 use vortex_session::registry::Registry;
 use vortex_session::{Ref, SessionExt};
 use vortex_utils::aliases::hash_map::HashMap;
@@ -22,23 +23,142 @@ use crate::expr::exprs::root::Root;
 use crate::expr::exprs::select::Select;
 use crate::expr::transform::remove_select::RemoveSelectRule;
 use crate::expr::transform::simplify::PackGetItemRule;
-use crate::expr::transform::traits::{ChildReduceRule, ParentReduceRule, ReduceRule};
-use crate::expr::{ExprId, ExprVTable};
+use crate::expr::transform::traits::{
+    ChildReduceRule, ParentReduceRule, ReduceRule, RewriteContext,
+};
+use crate::expr::{ExprId, ExprVTable, Expression, VTable};
 
 /// Registry of expression vtables.
 pub type ExprRegistry = Registry<ExprVTable>;
 
+/// Type-erased wrapper for ReduceRule that allows dynamic dispatch.
+pub(crate) trait DynReduceRule: Send + Sync {
+    fn reduce_dyn(
+        &self,
+        expr: &Expression,
+        ctx: &dyn RewriteContext,
+    ) -> VortexResult<Option<Expression>>;
+}
+
+/// Concrete wrapper that implements DynReduceRule for a specific VTable type.
+struct ReduceRuleAdapter<V: VTable, R: ReduceRule<V>> {
+    rule: R,
+    _phantom: std::marker::PhantomData<V>,
+}
+
+impl<V: VTable, R: ReduceRule<V>> ReduceRuleAdapter<V, R> {
+    fn new(rule: R) -> Self {
+        Self {
+            rule,
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<V: VTable, R: ReduceRule<V>> DynReduceRule for ReduceRuleAdapter<V, R> {
+    fn reduce_dyn(
+        &self,
+        expr: &Expression,
+        ctx: &dyn RewriteContext,
+    ) -> VortexResult<Option<Expression>> {
+        let Some(view) = expr.as_opt::<V>() else {
+            return Ok(None);
+        };
+        self.rule.reduce(&view, ctx)
+    }
+}
+
+/// Type-erased wrapper for ChildReduceRule that allows dynamic dispatch.
+pub(crate) trait DynChildReduceRule: Send + Sync {
+    fn reduce_child_dyn(
+        &self,
+        expr: &Expression,
+        child: &Expression,
+        child_idx: usize,
+        ctx: &dyn RewriteContext,
+    ) -> VortexResult<Option<Expression>>;
+}
+
+/// Concrete wrapper that implements DynChildReduceRule for a specific VTable type.
+struct ChildReduceRuleAdapter<V: VTable, R: ChildReduceRule<V>> {
+    rule: R,
+    _phantom: std::marker::PhantomData<V>,
+}
+
+impl<V: VTable, R: ChildReduceRule<V>> ChildReduceRuleAdapter<V, R> {
+    fn new(rule: R) -> Self {
+        Self {
+            rule,
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<V: VTable, R: ChildReduceRule<V>> DynChildReduceRule for ChildReduceRuleAdapter<V, R> {
+    fn reduce_child_dyn(
+        &self,
+        expr: &Expression,
+        child: &Expression,
+        child_idx: usize,
+        ctx: &dyn RewriteContext,
+    ) -> VortexResult<Option<Expression>> {
+        let Some(view) = expr.as_opt::<V>() else {
+            return Ok(None);
+        };
+        self.rule.reduce_child(&view, child, child_idx, ctx)
+    }
+}
+
+/// Type-erased wrapper for ParentReduceRule that allows dynamic dispatch.
+pub(crate) trait DynParentReduceRule: Send + Sync {
+    fn reduce_parent_dyn(
+        &self,
+        expr: &Expression,
+        parent: &Expression,
+        ctx: &dyn RewriteContext,
+    ) -> VortexResult<Option<Expression>>;
+}
+
+/// Concrete wrapper that implements DynParentReduceRule for a specific VTable type.
+struct ParentReduceRuleAdapter<V: VTable, R: ParentReduceRule<V>> {
+    rule: R,
+    _phantom: std::marker::PhantomData<V>,
+}
+
+impl<V: VTable, R: ParentReduceRule<V>> ParentReduceRuleAdapter<V, R> {
+    fn new(rule: R) -> Self {
+        Self {
+            rule,
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<V: VTable, R: ParentReduceRule<V>> DynParentReduceRule for ParentReduceRuleAdapter<V, R> {
+    fn reduce_parent_dyn(
+        &self,
+        expr: &Expression,
+        parent: &Expression,
+        ctx: &dyn RewriteContext,
+    ) -> VortexResult<Option<Expression>> {
+        let Some(view) = expr.as_opt::<V>() else {
+            return Ok(None);
+        };
+        self.rule.reduce_parent(&view, parent, ctx)
+    }
+}
+
 /// Registry of expression rewrite rules.
 ///
 /// Stores rewrite rules indexed by the expression ID they apply to.
 #[derive(Default)]
 pub struct RewriteRuleRegistry {
     /// Generic reduce rules (no context needed), indexed by expression ID
-    reduce_rules: HashMap<ExprId, Vec<Arc<dyn ReduceRule>>>,
+    reduce_rules: HashMap<ExprId, Vec<Arc<dyn DynReduceRule>>>,
     /// Child reduce rules, indexed by expression ID
-    child_rules: HashMap<ExprId, Vec<Arc<dyn ChildReduceRule>>>,
+    child_rules: HashMap<ExprId, Vec<Arc<dyn DynChildReduceRule>>>,
     /// Parent reduce rules, indexed by expression ID
-    parent_rules: HashMap<ExprId, Vec<Arc<dyn ParentReduceRule>>>,
+    parent_rules: HashMap<ExprId, Vec<Arc<dyn DynParentReduceRule>>>,
 }
 
 impl std::fmt::Debug for RewriteRuleRegistry {
@@ -57,35 +177,59 @@ impl RewriteRuleRegistry {
     }
 
     /// Register a generic reduce rule.
-    pub fn register_reduce_rule(&mut self, rule: Arc<dyn ReduceRule>) {
-        let id = rule.id();
-        self.reduce_rules.entry(id).or_default().push(rule);
+    pub fn register_reduce_rule<V: VTable, R: ReduceRule<V> + 'static>(
+        &mut self,
+        vtable: &'static V,
+        rule: R,
+    ) {
+        let id = vtable.id();
+        let adapter = ReduceRuleAdapter::new(rule);
+        self.reduce_rules
+            .entry(id)
+            .or_default()
+            .push(Arc::new(adapter));
     }
 
     /// Register a child reduce rule.
-    pub fn register_child_rule(&mut self, rule: Arc<dyn ChildReduceRule>) {
-        let id = rule.id();
-        self.child_rules.entry(id).or_default().push(rule);
+    pub fn register_child_rule<V: VTable, R: ChildReduceRule<V> + 'static>(
+        &mut self,
+        vtable: &'static V,
+        rule: R,
+    ) {
+        let id = vtable.id();
+        let adapter = ChildReduceRuleAdapter::new(rule);
+        self.child_rules
+            .entry(id)
+            .or_default()
+            .push(Arc::new(adapter));
     }
 
     /// Register a parent reduce rule.
-    pub fn register_parent_rule(&mut self, rule: Arc<dyn ParentReduceRule>) {
-        let id = rule.id();
-        self.parent_rules.entry(id).or_default().push(rule);
+    pub fn register_parent_rule<V: VTable, R: ParentReduceRule<V> + 'static>(
+        &mut self,
+        vtable: &'static V,
+        rule: R,
+    ) {
+        let id = vtable.id();
+        let adapter = ParentReduceRuleAdapter::new(rule);
+        self.parent_rules
+            .entry(id)
+            .or_default()
+            .push(Arc::new(adapter));
     }
 
     /// Get all generic reduce rules for a given expression ID.
-    pub fn reduce_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn ReduceRule>]> {
+    pub(crate) fn reduce_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn DynReduceRule>]> {
         self.reduce_rules.get(id).map(|v| v.as_slice())
     }
 
     /// Get all child reduce rules for a given expression ID.
-    pub fn child_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn ChildReduceRule>]> {
+    pub(crate) fn child_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn DynChildReduceRule>]> {
         self.child_rules.get(id).map(|v| v.as_slice())
     }
 
     /// Get all parent reduce rules for a given expression ID.
-    pub fn parent_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn ParentReduceRule>]> {
+    pub(crate) fn parent_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn DynParentReduceRule>]> {
         self.parent_rules.get(id).map(|v| v.as_slice())
     }
 }
@@ -118,18 +262,30 @@ impl ExprSession {
     }
 
     /// Register a generic reduce rule in the session.
-    pub fn register_reduce_rule(&mut self, rule: impl ReduceRule + 'static) {
-        self.rewrite_rules.register_reduce_rule(Arc::new(rule));
+    pub fn register_reduce_rule<V: VTable>(
+        &mut self,
+        vtable: &'static V,
+        rule: impl ReduceRule<V> + 'static,
+    ) {
+        self.rewrite_rules.register_reduce_rule(vtable, rule);
     }
 
     /// Register a child reduce rule in the session.
-    pub fn register_child_rule(&mut self, rule: impl ChildReduceRule + 'static) {
-        self.rewrite_rules.register_child_rule(Arc::new(rule));
+    pub fn register_child_rule<V: VTable>(
+        &mut self,
+        vtable: &'static V,
+        rule: impl ChildReduceRule<V> + 'static,
+    ) {
+        self.rewrite_rules.register_child_rule(vtable, rule);
     }
 
     /// Register a parent reduce rule in the session.
-    pub fn register_parent_rule(&mut self, rule: impl ParentReduceRule + 'static) {
-        self.rewrite_rules.register_parent_rule(Arc::new(rule));
+    pub fn register_parent_rule<V: VTable>(
+        &mut self,
+        vtable: &'static V,
+        rule: impl ParentReduceRule<V> + 'static,
+    ) {
+        self.rewrite_rules.register_parent_rule(vtable, rule);
     }
 }
 
@@ -156,8 +312,8 @@ impl Default for ExprSession {
 
         // Register built-in rewrite rules
         let mut rewrite_rules = RewriteRuleRegistry::new();
-        rewrite_rules.register_reduce_rule(Arc::new(RemoveSelectRule));
-        rewrite_rules.register_child_rule(Arc::new(PackGetItemRule));
+        rewrite_rules.register_reduce_rule(&Select, RemoveSelectRule);
+        rewrite_rules.register_child_rule(&GetItem, PackGetItemRule);
 
         Self {
             registry: expressions,

vortex-array/src/expr/transform/remove_select.rs

@@ -9,7 +9,7 @@ use crate::expr::exprs::pack::pack;
 use crate::expr::exprs::select::Select;
 use crate::expr::transform::traits::{ReduceRule, RewriteContext};
 use crate::expr::traversal::{NodeExt, Transformed};
-use crate::expr::{ExprId, Expression};
+use crate::expr::{Expression, ExpressionView};
 
 /// Replaces [crate::SelectExpr] with combination of [crate::GetItem] and [crate::Pack] expressions.
 pub(crate) fn remove_select(e: Expression, ctx: &DType) -> VortexResult<Expression> {
@@ -60,20 +60,12 @@ fn remove_select_transformer(
 /// Transforms: `select(["a", "b"], expr)` → `pack(a: get_item("a", expr), b: get_item("b", expr))`
 pub struct RemoveSelectRule;
 
-impl ReduceRule for RemoveSelectRule {
-    fn id(&self) -> ExprId {
-        ExprId::new_ref("vortex.select")
-    }
-
+impl ReduceRule<Select> for RemoveSelectRule {
     fn reduce(
         &self,
-        expr: &Expression,
+        select: &ExpressionView<Select>,
         ctx: &dyn RewriteContext,
     ) -> VortexResult<Option<Expression>> {
-        let Some(select) = expr.as_opt::<Select>() else {
-            return Ok(None);
-        };
-
         let child = select.child();
         let child_dtype = child.return_dtype(ctx.dtype())?;
         let child_nullability = child_dtype.nullability();
@@ -114,7 +106,7 @@ mod tests {
     use super::{RemoveSelectRule, remove_select};
     use crate::expr::exprs::pack::Pack;
     use crate::expr::exprs::root::root;
-    use crate::expr::exprs::select::select;
+    use crate::expr::exprs::select::{Select, select};
     use crate::expr::session::ExprSession;
     use crate::expr::transform::simplify_typed::apply_child_rules;
     use crate::expr::transform::traits::{ReduceRule, SimpleRewriteContext};
@@ -142,7 +134,8 @@ mod tests {
 
         let rule = RemoveSelectRule;
         let ctx = SimpleRewriteContext { dtype: &dtype };
-        let result = rule.reduce(&e, &ctx).unwrap();
+        let select_view = e.as_::<Select>();
+        let result = rule.reduce(&select_view, &ctx).unwrap();
 
         assert!(result.is_some());
         let transformed = result.unwrap();