wip

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

Author: joseph-isaacs

vortex-array/src/expr/exprs/get_item.rs …tex-array/src/expr/exprs/get_item/mod.rsvortex-array/src/expr/exprs/get_item.rs renamed to vortex-array/src/expr/exprs/get_item/mod.rs vortex-array/src/expr/exprs/get_item.rs renamed to vortex-array/src/expr/exprs/get_item/mod.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+pub mod transform;
+
 use std::fmt::Formatter;
 use std::ops::Not;
 

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

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::expr::exprs::get_item::GetItem;
+use crate::expr::exprs::pack::Pack;
+use crate::expr::transform::traits::{ChildReduceRule, 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> for PackGetItemRule {
+    fn reduce_child(
+        &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
+            let field_expr = pack.field(get_item.data())?;
+            return Ok(Some(field_expr));
+        }
+
+        Ok(None)
+    }
+}

vortex-array/src/expr/exprs/merge.rs vortex-array/src/expr/exprs/merge/mod.rsvortex-array/src/expr/exprs/merge.rs renamed to vortex-array/src/expr/exprs/merge/mod.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+pub mod transform;
+
 use std::fmt::Formatter;
 use std::hash::Hash;
 use std::sync::Arc;

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

@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use itertools::Itertools as _;
+use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
+use vortex_utils::aliases::hash_set::HashSet;
+
+use crate::expr::exprs::get_item::get_item;
+use crate::expr::exprs::merge::{DuplicateHandling, Merge};
+use crate::expr::exprs::pack::pack;
+use crate::expr::transform::traits::{ReduceRule, RewriteContext};
+use crate::expr::{Expression, ExpressionView};
+
+/// Rule that removes Merge expressions by converting them to Pack + GetItem.
+///
+/// Transforms: `merge([struct1, struct2])` → `pack(field1: get_item("field1", struct1), field2: get_item("field2", struct2), ...)`
+pub struct RemoveMergeRule;
+
+impl ReduceRule<Merge> for RemoveMergeRule {
+    fn reduce(
+        &self,
+        merge: &ExpressionView<Merge>,
+        ctx: &dyn RewriteContext,
+    ) -> VortexResult<Option<Expression>> {
+        let merge_dtype = merge.return_dtype(ctx.dtype())?;
+        let mut names = Vec::with_capacity(merge.children().len() * 2);
+        let mut children = Vec::with_capacity(merge.children().len() * 2);
+        let mut duplicate_names = HashSet::<_>::new();
+
+        for child in merge.children().iter() {
+            let child_dtype = child.return_dtype(ctx.dtype())?;
+            if !child_dtype.is_struct() {
+                return Err(vortex_err!(
+                    "Merge child must return a non-nullable struct dtype, got {}",
+                    child_dtype
+                ));
+            }
+
+            let child_dtype = child_dtype
+                .as_struct_fields_opt()
+                .vortex_expect("expected struct");
+
+            for name in child_dtype.names().iter() {
+                if let Some(idx) = names.iter().position(|n| n == name) {
+                    duplicate_names.insert(name.clone());
+                    children[idx] = child.clone();
+                } else {
+                    names.push(name.clone());
+                    children.push(child.clone());
+                }
+            }
+
+            if merge.data() == &DuplicateHandling::Error && !duplicate_names.is_empty() {
+                vortex_bail!(
+                    "merge: duplicate fields in children: {}",
+                    duplicate_names.into_iter().format(", ")
+                )
+            }
+        }
+
+        let expr = pack(
+            names
+                .into_iter()
+                .zip(children)
+                .map(|(name, child)| (name.clone(), get_item(name, child))),
+            merge_dtype.nullability(),
+        );
+
+        Ok(Some(expr))
+    }
+}

vortex-array/src/expr/exprs/select.rs vortex-array/src/expr/exprs/select/mod.rsvortex-array/src/expr/exprs/select.rs renamed to vortex-array/src/expr/exprs/select/mod.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+pub mod transform;
+
 use std::fmt::{Display, Formatter};
 
 use itertools::Itertools;

…rray/src/expr/transform/remove_select.rs …array/src/expr/exprs/select/transform.rsvortex-array/src/expr/transform/remove_select.rs renamed to vortex-array/src/expr/exprs/select/transform.rs vortex-array/src/expr/transform/remove_select.rs renamed to vortex-array/src/expr/exprs/select/transform.rs

@@ -1,60 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_err};
 
 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::{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> {
-    e.transform_up(|node| remove_select_transformer(node, ctx))
-        .map(|e| e.into_inner())
-}
-
-fn remove_select_transformer(
-    node: Expression,
-    ctx: &DType,
-) -> VortexResult<Transformed<Expression>> {
-    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))`
@@ -103,29 +57,14 @@ mod tests {
     use vortex_dtype::PType::I32;
     use vortex_dtype::{DType, StructFields};
 
-    use super::{RemoveSelectRule, remove_select};
+    use super::RemoveSelectRule;
     use crate::expr::exprs::pack::Pack;
     use crate::expr::exprs::root::root;
     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};
 
     #[test]
-    fn test_remove_select() {
-        let dtype = DType::Struct(
-            StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]),
-            Nullable,
-        );
-        let e = select(["a", "b"], root());
-        let e = remove_select(e, &dtype).unwrap();
-
-        assert!(e.is::<Pack>());
-        assert!(e.return_dtype(&dtype).unwrap().is_nullable());
-    }
-
-    #[test]
-    fn test_remove_select_rule_direct() {
+    fn test_remove_select_rule() {
         let dtype = DType::Struct(
             StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]),
             Nullable,
@@ -144,30 +83,31 @@ mod tests {
     }
 
     #[test]
-    fn test_remove_select_via_session() {
+    fn test_remove_select_rule_exclude_fields() {
+        use crate::expr::exprs::select::select_exclude;
+
         let dtype = DType::Struct(
             StructFields::new(
                 ["a", "b", "c"].into(),
                 vec![I32.into(), I32.into(), I32.into()],
             ),
             Nullable,
         );
+        let e = select_exclude(["c"], root());
 
-        // 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();
+        let rule = RemoveSelectRule;
+        let ctx = SimpleRewriteContext { dtype: &dtype };
+        let select_view = e.as_::<Select>();
+        let result = rule.reduce(&select_view, &ctx).unwrap();
 
-        // Should be transformed to Pack
-        assert!(result.is::<Pack>());
+        assert!(result.is_some());
+        let transformed = result.unwrap();
+        assert!(transformed.is::<Pack>());
 
-        // Verify the dtype has only selected fields
-        let result_dtype = result.return_dtype(&dtype).unwrap();
+        // Should exclude "c" and include "a" and "b"
+        let result_dtype = transformed.return_dtype(&dtype).unwrap();
+        assert!(result_dtype.is_nullable());
         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");
+        assert_eq!(fields.names().as_ref(), &["a", "b"]);
     }
 }

vortex-array/src/expr/session.rs

@@ -12,17 +12,18 @@ use crate::expr::exprs::between::Between;
 use crate::expr::exprs::binary::Binary;
 use crate::expr::exprs::cast::Cast;
 use crate::expr::exprs::get_item::GetItem;
+use crate::expr::exprs::get_item::transform::PackGetItemRule;
 use crate::expr::exprs::is_null::IsNull;
 use crate::expr::exprs::like::Like;
 use crate::expr::exprs::list_contains::ListContains;
 use crate::expr::exprs::literal::Literal;
 use crate::expr::exprs::merge::Merge;
+use crate::expr::exprs::merge::transform::RemoveMergeRule;
 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::exprs::select::transform::RemoveSelectRule;
 use crate::expr::transform::traits::{
     ChildReduceRule, ParentReduceRule, ReduceRule, RewriteContext,
 };
@@ -313,6 +314,7 @@ impl Default for ExprSession {
         // Register built-in rewrite rules
         let mut rewrite_rules = RewriteRuleRegistry::new();
         rewrite_rules.register_reduce_rule(&Select, RemoveSelectRule);
+        rewrite_rules.register_reduce_rule(&Merge, RemoveMergeRule);
         rewrite_rules.register_child_rule(&GetItem, PackGetItemRule);
 
         Self {

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

@@ -6,8 +6,6 @@ pub mod annotations;
 pub mod immediate_access;
 pub(crate) mod match_between;
 mod partition;
-mod remove_merge;
-pub mod remove_select;
 mod replace;
 pub mod simplify;
 mod simplify_typed;

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

@@ -209,7 +209,6 @@ mod tests {
     use crate::expr::exprs::select::select;
     use crate::expr::transform::immediate_access::annotate_scope_access;
     use crate::expr::transform::replace::replace_root_fields;
-    use crate::expr::transform::simplify::simplify;
     use crate::expr::transform::simplify_typed::simplify_typed;
 
     #[fixture]
@@ -274,7 +273,7 @@ mod tests {
 
         let split_a = partitioned.find_partition(&"a".into()).unwrap();
         assert_eq!(
-            &simplify(split_a.clone()).unwrap(),
+            &simplify_typed(split_a.clone(), &dtype).unwrap(),
             &pack(
                 [
                     ("a_0", get_item("x", get_item("a", root()))),