update

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

Author: joseph-isaacs

vortex-array/src/expr/exprs/get_item/transform.rs

@@ -5,39 +5,23 @@ use vortex_error::VortexResult;
 
 use crate::expr::exprs::get_item::GetItem;
 use crate::expr::exprs::pack::Pack;
-use crate::expr::transform::rules::{ChildReduceRule, RewriteContext};
+use crate::expr::transform::ReduceRule;
+use crate::expr::transform::rules::RewriteContext;
 use crate::expr::{Expression, ExpressionView};
 
 /// Rewrite rule: `pack(l_1: e_1, ..., l_i: e_i, ..., l_n: e_n).get_item(l_i) = e_i`
 ///
 /// Simplifies accessing a field from a pack expression by directly returning the field's
 /// expression instead of materializing the pack.
-///
-/// # Example
-/// ```
-/// # use vortex_array::expr::exprs::{get_item::get_item, literal::lit, pack::pack};
-/// # use vortex_dtype::Nullability::NonNullable;
-/// let e = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
-/// // After applying PackGetItemRule, this becomes: lit(2)
-/// ```
 pub struct PackGetItemRule;
 
-impl ChildReduceRule<GetItem, &dyn RewriteContext> for PackGetItemRule {
-    fn reduce_child(
+impl ReduceRule<GetItem, &dyn RewriteContext> for PackGetItemRule {
+    fn reduce(
         &self,
         get_item: &ExpressionView<GetItem>,
-        child: &Expression,
-        child_idx: usize,
         _ctx: &dyn RewriteContext,
     ) -> VortexResult<Option<Expression>> {
-        // Only consider the first child (child_idx == 0) of GetItem expressions
-        if child_idx != 0 {
-            return Ok(None);
-        }
-
-        // Check if child is Pack
-        if let Some(pack) = child.as_opt::<Pack>() {
-            // Extract the field from the pack
+        if let Some(pack) = get_item.child(0).as_opt::<Pack>() {
             let field_expr = pack.field(get_item.data())?;
             return Ok(Some(field_expr));
         }
@@ -52,31 +36,28 @@ mod tests {
     use vortex_dtype::{DType, PType};
 
     use super::PackGetItemRule;
-    use crate::expr::ExprId;
     use crate::expr::exprs::binary::checked_add;
     use crate::expr::exprs::get_item::{GetItem, get_item};
     use crate::expr::exprs::literal::lit;
     use crate::expr::exprs::pack::pack;
     use crate::expr::session::ExprSession;
-    use crate::expr::transform::rules::{ChildReduceRule, SimpleRewriteContext};
-    use crate::expr::transform::simplify_typed;
+    use crate::expr::transform::rules::SimpleRewriteContext;
+    use crate::expr::transform::{ReduceRule, simplify_typed};
 
     #[test]
     fn test_pack_get_item_rule() {
         let rule = PackGetItemRule;
 
         // Create: pack(a: lit(1), b: lit(2)).get_item("b")
         let pack_expr = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
-        let get_item_expr = get_item("b", pack_expr.clone());
+        let get_item_expr = get_item("b", pack_expr);
 
         // Create a dummy context
         let dtype = DType::Primitive(PType::I32, NonNullable);
         let ctx = SimpleRewriteContext { dtype: &dtype };
 
         let get_item_view = get_item_expr.as_::<GetItem>();
-        let result = rule
-            .reduce_child(&get_item_view, &pack_expr, 0, &ctx)
-            .unwrap();
+        let result = rule.reduce(&get_item_view, &ctx).unwrap();
 
         assert!(result.is_some());
         assert_eq!(&result.unwrap(), &lit(2));
@@ -94,39 +75,11 @@ mod tests {
         let ctx = SimpleRewriteContext { dtype: &dtype };
 
         let get_item_view = get_item_expr.as_::<GetItem>();
-        let result = rule
-            .reduce_child(&get_item_view, &lit_expr, 0, &ctx)
-            .unwrap();
+        let result = rule.reduce(&get_item_view, &ctx).unwrap();
 
         assert!(result.is_none());
     }
 
-    #[test]
-    fn test_pack_get_item_rule_from_session() {
-        let session = ExprSession::default();
-
-        let get_item_id = ExprId::new_ref("vortex.get_item");
-        let rules = session.rewrite_rules().child_rules_for(&get_item_id);
-
-        // Should have at least one rule registered (PackGetItemRule)
-        assert!(rules.is_some());
-        assert_eq!(rules.unwrap().len(), 1);
-
-        let pack_expr = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
-        let get_item_expr = get_item("b", pack_expr.clone());
-
-        let dtype = DType::Primitive(PType::I32, NonNullable);
-        let ctx = SimpleRewriteContext { dtype: &dtype };
-
-        let rule = &rules.unwrap()[0];
-        let result = rule
-            .reduce_child_dyn(&get_item_expr, &pack_expr, 0, &ctx)
-            .unwrap();
-
-        assert!(result.is_some());
-        assert_eq!(&result.unwrap(), &lit(2));
-    }
-
     #[test]
     fn test_multi_level_pack_get_item_simplify() {
         let inner_pack = pack([("a", lit(1)), ("b", lit(2))], NonNullable);

vortex-array/src/expr/exprs/merge/transform.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use itertools::Itertools as _;
-use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::expr::exprs::get_item::get_item;
@@ -31,10 +31,10 @@ impl ReduceRule<Merge, &dyn TypedRewriteContext> for RemoveMergeRule {
         for child in merge.children().iter() {
             let child_dtype = child.return_dtype(ctx.dtype())?;
             if !child_dtype.is_struct() {
-                return Err(vortex_err!(
+                vortex_bail!(
                     "Merge child must return a non-nullable struct dtype, got {}",
                     child_dtype
-                ));
+                )
             }
 
             let child_dtype = child_dtype

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

@@ -23,9 +23,7 @@ use crate::expr::exprs::pack::Pack;
 use crate::expr::exprs::root::Root;
 use crate::expr::exprs::select::Select;
 use crate::expr::exprs::select::transform::RemoveSelectRule;
-use crate::expr::transform::rules::{
-    ChildReduceRule, ParentReduceRule, ReduceRule, RewriteContext,
-};
+use crate::expr::transform::rules::{ParentReduceRule, ReduceRule, RewriteContext};
 use crate::expr::{ExprVTable, VTable};
 
 /// Registry of expression vtables.
@@ -69,7 +67,7 @@ impl ExprSession {
         self.rewrite_rules.register_typed_reduce_rule(vtable, rule);
     }
 
-    /// Register a generic reduce rule that only uses RewriteContext (non-typed).
+    /// Register a reduce rule that doesn't use TypedRewriteContext.
     /// Use this for rules that don't need access to dtype information.
     pub fn register_reduce_rule<V, R>(&mut self, vtable: &'static V, rule: R)
     where
@@ -80,28 +78,6 @@ impl ExprSession {
         self.rewrite_rules.register_reduce_rule(vtable, rule);
     }
 
-    /// Register a child reduce rule that uses TypedRewriteContext.
-    /// Use this for rules that need access to dtype information.
-    pub fn register_typed_child_rule<V, R>(&mut self, vtable: &'static V, rule: R)
-    where
-        V: VTable,
-        R: 'static,
-        for<'a> R: ChildReduceRule<V, &'a dyn crate::expr::transform::TypedRewriteContext>,
-    {
-        self.rewrite_rules.register_typed_child_rule(vtable, rule);
-    }
-
-    /// Register a child reduce rule that only uses RewriteContext (non-typed).
-    /// Use this for rules that don't need access to dtype information.
-    pub fn register_child_rule<V, R>(&mut self, vtable: &'static V, rule: R)
-    where
-        V: VTable,
-        R: 'static,
-        for<'a> R: ChildReduceRule<V, &'a dyn RewriteContext>,
-    {
-        self.rewrite_rules.register_child_rule(vtable, rule);
-    }
-
     /// Register a parent reduce rule in the session.
     pub fn register_parent_rule<V: VTable>(
         &mut self,
@@ -137,7 +113,7 @@ impl Default for ExprSession {
         let mut rewrite_rules = RewriteRuleRegistry::new();
         rewrite_rules.register_typed_reduce_rule(&Select, RemoveSelectRule);
         rewrite_rules.register_typed_reduce_rule(&Merge, RemoveMergeRule);
-        rewrite_rules.register_child_rule(&GetItem, PackGetItemRule);
+        rewrite_rules.register_reduce_rule(&GetItem, PackGetItemRule);
 
         Self {
             registry: expressions,