wip

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

Author: joseph-isaacs

vortex-array/src/expr/session.rs

@@ -1,10 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
 use vortex_session::registry::Registry;
 use vortex_session::{Ref, SessionExt};
+use vortex_utils::aliases::hash_map::HashMap;
 
-use crate::expr::ExprVTable;
 use crate::expr::exprs::between::Between;
 use crate::expr::exprs::binary::Binary;
 use crate::expr::exprs::cast::Cast;
@@ -18,21 +20,93 @@ use crate::expr::exprs::not::Not;
 use crate::expr::exprs::pack::Pack;
 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};
 
 /// Registry of expression vtables.
 pub type ExprRegistry = Registry<ExprVTable>;
 
-/// Session state for expression vtables.
+/// 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>>>,
+    /// Child reduce rules, indexed by expression ID
+    child_rules: HashMap<ExprId, Vec<Arc<dyn ChildReduceRule>>>,
+    /// Parent reduce rules, indexed by expression ID
+    parent_rules: HashMap<ExprId, Vec<Arc<dyn ParentReduceRule>>>,
+}
+
+impl std::fmt::Debug for RewriteRuleRegistry {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("RewriteRuleRegistry")
+            .field("reduce_rules_count", &self.reduce_rules.len())
+            .field("child_rules_count", &self.child_rules.len())
+            .field("parent_rules_count", &self.parent_rules.len())
+            .finish()
+    }
+}
+
+impl RewriteRuleRegistry {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// 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);
+    }
+
+    /// 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);
+    }
+
+    /// 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);
+    }
+
+    /// Get all generic reduce rules for a given expression ID.
+    pub fn reduce_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn ReduceRule>]> {
+        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>]> {
+        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>]> {
+        self.parent_rules.get(id).map(|v| v.as_slice())
+    }
+}
+
+/// Session state for expression vtables and rewrite rules.
 #[derive(Debug)]
 pub struct ExprSession {
     registry: ExprRegistry,
+    rewrite_rules: RewriteRuleRegistry,
 }
 
 impl ExprSession {
     pub fn registry(&self) -> &ExprRegistry {
         &self.registry
     }
 
+    /// Get the rewrite rule registry.
+    pub fn rewrite_rules(&self) -> &RewriteRuleRegistry {
+        &self.rewrite_rules
+    }
+
     /// Register an expression vtable in the session, replacing any existing vtable with the same ID.
     pub fn register(&self, expr: ExprVTable) {
         self.registry.register(expr)
@@ -42,6 +116,21 @@ impl ExprSession {
     pub fn register_many(&self, exprs: impl IntoIterator<Item = ExprVTable>) {
         self.registry.register_many(exprs);
     }
+
+    /// 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));
+    }
+
+    /// 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));
+    }
+
+    /// 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));
+    }
 }
 
 impl Default for ExprSession {
@@ -65,8 +154,14 @@ impl Default for ExprSession {
             ExprVTable::from_static(&Select),
         ]);
 
+        // 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));
+
         Self {
             registry: expressions,
+            rewrite_rules,
         }
     }
 }

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

@@ -7,21 +7,14 @@ pub mod immediate_access;
 pub(crate) mod match_between;
 mod partition;
 mod remove_merge;
-mod remove_select;
+pub mod remove_select;
 mod replace;
-mod simplify;
+pub mod simplify;
 mod simplify_typed;
+pub mod traits;
 
 pub use partition::*;
 pub use replace::*;
 pub use simplify::*;
 pub use simplify_typed::*;
-
-use crate::expr::traversal::Transformed;
-use crate::expr::{ExprId, Expression};
-
-trait ExpressionParentTransformer {
-    fn id(&self) -> ExprId;
-
-    fn reduce_parent(&self, expr: &Expression, parent: &Expression) -> Transformed<Expression>;
-}
+pub use traits::*;

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

@@ -4,11 +4,12 @@
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_err};
 
-use crate::expr::Expression;
 use crate::expr::exprs::get_item::get_item;
 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};
 
 /// Replaces [crate::SelectExpr] with combination of [crate::GetItem] and [crate::Pack] expressions.
 pub(crate) fn remove_select(e: Expression, ctx: &DType) -> VortexResult<Expression> {
@@ -20,38 +21,87 @@ fn remove_select_transformer(
     node: Expression,
     ctx: &DType,
 ) -> VortexResult<Transformed<Expression>> {
-    match node.as_opt::<Select>() {
-        None => Ok(Transformed::no(node)),
-        Some(select) => {
-            let child = select.child();
-            let child_dtype = child.return_dtype(ctx)?;
-            let child_nullability = child_dtype.nullability();
-
-            let child_dtype = child_dtype.as_struct_fields_opt().ok_or_else(|| {
-                vortex_err!(
-                    "Select child must return a struct dtype, however it was a {}",
-                    child_dtype
-                )
-            })?;
-
-            let expr = pack(
-                select
-                    .data()
-                    .as_include_names(child_dtype.names())
-                    .map_err(|e| {
-                        e.with_context(format!(
-                            "Select fields {:?} must be a subset of child fields {:?}",
-                            select.data(),
-                            child_dtype.names()
-                        ))
-                    })?
-                    .iter()
-                    .map(|name| (name.clone(), get_item(name.clone(), child.clone()))),
-                child_nullability,
-            );
-
-            Ok(Transformed::yes(expr))
-        }
+    if let Some(select) = node.as_opt::<Select>() {
+        let child = select.child();
+        let child_dtype = child.return_dtype(ctx)?;
+        let child_nullability = child_dtype.nullability();
+
+        let child_dtype = child_dtype.as_struct_fields_opt().ok_or_else(|| {
+            vortex_err!(
+                "Select child must return a struct dtype, however it was a {}",
+                child_dtype
+            )
+        })?;
+
+        let expr = pack(
+            select
+                .data()
+                .as_include_names(child_dtype.names())
+                .map_err(|e| {
+                    e.with_context(format!(
+                        "Select fields {:?} must be a subset of child fields {:?}",
+                        select.data(),
+                        child_dtype.names()
+                    ))
+                })?
+                .iter()
+                .map(|name| (name.clone(), get_item(name.clone(), child.clone()))),
+            child_nullability,
+        );
+
+        Ok(Transformed::yes(expr))
+    } else {
+        Ok(Transformed::no(node))
+    }
+}
+
+/// Rule that removes Select expressions by converting them to Pack + GetItem.
+///
+/// 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")
+    }
+
+    fn reduce(
+        &self,
+        expr: &Expression,
+        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();
+
+        let child_dtype = child_dtype.as_struct_fields_opt().ok_or_else(|| {
+            vortex_err!(
+                "Select child must return a struct dtype, however it was a {}",
+                child_dtype
+            )
+        })?;
+
+        let expr = pack(
+            select
+                .data()
+                .as_include_names(child_dtype.names())
+                .map_err(|e| {
+                    e.with_context(format!(
+                        "Select fields {:?} must be a subset of child fields {:?}",
+                        select.data(),
+                        child_dtype.names()
+                    ))
+                })?
+                .iter()
+                .map(|name| (name.clone(), get_item(name.clone(), child.clone()))),
+            child_nullability,
+        );
+
+        Ok(Some(expr))
     }
 }
 
@@ -61,10 +111,13 @@ mod tests {
     use vortex_dtype::PType::I32;
     use vortex_dtype::{DType, StructFields};
 
-    use super::remove_select;
+    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::session::ExprSession;
+    use crate::expr::transform::simplify_typed::apply_child_rules;
+    use crate::expr::transform::traits::{ReduceRule, SimpleRewriteContext};
 
     #[test]
     fn test_remove_select() {
@@ -78,4 +131,50 @@ mod tests {
         assert!(e.is::<Pack>());
         assert!(e.return_dtype(&dtype).unwrap().is_nullable());
     }
+
+    #[test]
+    fn test_remove_select_rule_direct() {
+        let dtype = DType::Struct(
+            StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]),
+            Nullable,
+        );
+        let e = select(["a", "b"], root());
+
+        let rule = RemoveSelectRule;
+        let ctx = SimpleRewriteContext { dtype: &dtype };
+        let result = rule.reduce(&e, &ctx).unwrap();
+
+        assert!(result.is_some());
+        let transformed = result.unwrap();
+        assert!(transformed.is::<Pack>());
+        assert!(transformed.return_dtype(&dtype).unwrap().is_nullable());
+    }
+
+    #[test]
+    fn test_remove_select_via_session() {
+        let dtype = DType::Struct(
+            StructFields::new(
+                ["a", "b", "c"].into(),
+                vec![I32.into(), I32.into(), I32.into()],
+            ),
+            Nullable,
+        );
+
+        // Create expression: select(["a", "c"], root())
+        let e = select(["a", "c"], root());
+
+        // Use session which has RemoveSelectRule registered
+        let session = ExprSession::default();
+        let result = apply_child_rules(e, &dtype, &session).unwrap();
+
+        // Should be transformed to Pack
+        assert!(result.is::<Pack>());
+
+        // Verify the dtype has only selected fields
+        let result_dtype = result.return_dtype(&dtype).unwrap();
+        let fields = result_dtype.as_struct_fields_opt().unwrap();
+        assert_eq!(fields.names().len(), 2);
+        assert_eq!(fields.names()[0].as_ref(), "a");
+        assert_eq!(fields.names()[1].as_ref(), "c");
+    }
 }