feat[vortex-expr]: add expr rewrite rules (#5348)

Breaking:

Using a vortex session in more places
 - simplify_typed/simplify
 - Layout::new_reader

---------

Signed-off-by: Joe Isaacs <[email protected]>
Co-authored-by: Nicholas Gates <[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,127 @@
+// 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::rules::{ReduceRule, RuleContext};
+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.
+pub struct PackGetItemRule;
+
+impl ReduceRule<GetItem, RuleContext> for PackGetItemRule {
+    fn reduce(
+        &self,
+        get_item: &ExpressionView<GetItem>,
+        _ctx: &RuleContext,
+    ) -> VortexResult<Option<Expression>> {
+        if let Some(pack) = get_item.child(0).as_opt::<Pack>() {
+            let field_expr = pack.field(get_item.data())?;
+            return Ok(Some(field_expr));
+        }
+
+        Ok(None)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_dtype::Nullability::NonNullable;
+    use vortex_dtype::{DType, PType};
+
+    use super::PackGetItemRule;
+    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::{ReduceRule, RuleContext};
+    use crate::expr::transform::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);
+
+        let get_item_view = get_item_expr.as_::<GetItem>();
+        let result = rule.reduce(&get_item_view, &RuleContext).unwrap();
+
+        assert!(result.is_some());
+        assert_eq!(&result.unwrap(), &lit(2));
+    }
+
+    #[test]
+    fn test_pack_get_item_rule_no_match() {
+        let rule = PackGetItemRule;
+
+        // Create: get_item("x", lit(42)) - not a pack child
+        let lit_expr = lit(42);
+        let get_item_expr = get_item("x", lit_expr);
+
+        let get_item_view = get_item_expr.as_::<GetItem>();
+        let result = rule.reduce(&get_item_view, &RuleContext).unwrap();
+
+        assert!(result.is_none());
+    }
+
+    #[test]
+    fn test_multi_level_pack_get_item_simplify() {
+        let inner_pack = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
+        let get_a = get_item("a", inner_pack);
+
+        let outer_pack = pack([("x", get_a), ("y", lit(3)), ("z", lit(4))], NonNullable);
+        let get_z = get_item("z", outer_pack);
+
+        let dtype = DType::Primitive(PType::I32, NonNullable);
+
+        let result = simplify_typed(get_z, &dtype, &ExprSession::default()).unwrap();
+
+        assert_eq!(&result, &lit(4));
+    }
+
+    #[test]
+    fn test_deeply_nested_pack_get_item() {
+        let innermost = pack([("a", lit(42))], NonNullable);
+        let get_a = get_item("a", innermost);
+
+        let level2 = pack([("b", get_a)], NonNullable);
+        let get_b = get_item("b", level2);
+
+        let level3 = pack([("c", get_b)], NonNullable);
+        let get_c = get_item("c", level3);
+
+        let outermost = pack([("final", get_c)], NonNullable);
+        let get_final = get_item("final", outermost);
+
+        let dtype = DType::Primitive(PType::I32, NonNullable);
+
+        let result = simplify_typed(get_final, &dtype, &ExprSession::default()).unwrap();
+
+        assert_eq!(&result, &lit(42));
+    }
+
+    #[test]
+    fn test_partial_pack_get_item_simplify() {
+        let inner_pack = pack([("x", lit(1)), ("y", lit(2))], NonNullable);
+        let get_x = get_item("x", inner_pack);
+        let add_expr = checked_add(get_x, lit(10));
+
+        let outer_pack = pack([("result", add_expr)], NonNullable);
+        let get_result = get_item("result", outer_pack);
+
+        let dtype = DType::Primitive(PType::I32, NonNullable);
+
+        let result = simplify_typed(get_result, &dtype, &ExprSession::default()).unwrap();
+
+        let expected = checked_add(lit(1), lit(10));
+        assert_eq!(&result, &expected);
+    }
+}

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,114 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use itertools::Itertools as _;
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
+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::rules::{ReduceRule, TypedRuleContext};
+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, TypedRuleContext> for RemoveMergeRule {
+    fn reduce(
+        &self,
+        merge: &ExpressionView<Merge>,
+        ctx: &TypedRuleContext,
+    ) -> 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() {
+                vortex_bail!(
+                    "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))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_dtype::DType;
+    use vortex_dtype::Nullability::NonNullable;
+    use vortex_dtype::PType::{I32, I64, U32, U64};
+
+    use super::RemoveMergeRule;
+    use crate::expr::exprs::get_item::get_item;
+    use crate::expr::exprs::merge::{DuplicateHandling, Merge, merge_opts};
+    use crate::expr::exprs::pack::Pack;
+    use crate::expr::exprs::root::root;
+    use crate::expr::transform::rules::{ReduceRule, TypedRuleContext};
+
+    #[test]
+    fn test_remove_merge() {
+        let dtype = DType::struct_(
+            [
+                ("0", DType::struct_([("a", I32), ("b", I64)], NonNullable)),
+                ("1", DType::struct_([("b", U32), ("c", U64)], NonNullable)),
+            ],
+            NonNullable,
+        );
+
+        let e = merge_opts(
+            [get_item("0", root()), get_item("1", root())],
+            DuplicateHandling::RightMost,
+        );
+
+        let ctx = TypedRuleContext::new(dtype.clone());
+        let rule = RemoveMergeRule;
+        let merge_view = e.as_::<Merge>();
+        let result = rule.reduce(&merge_view, &ctx).unwrap();
+
+        assert!(result.is_some());
+        let result = result.unwrap();
+        assert!(result.is::<Pack>());
+        assert_eq!(
+            result.return_dtype(&dtype).unwrap(),
+            DType::struct_([("a", I32), ("b", U32), ("c", U64)], NonNullable)
+        );
+    }
+}

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;

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

@@ -0,0 +1,113 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+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::rules::{ReduceRule, TypedRuleContext};
+use crate::expr::{Expression, ExpressionView};
+
+/// 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<Select, TypedRuleContext> for RemoveSelectRule {
+    fn reduce(
+        &self,
+        select: &ExpressionView<Select>,
+        ctx: &TypedRuleContext,
+    ) -> VortexResult<Option<Expression>> {
+        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))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_dtype::Nullability::Nullable;
+    use vortex_dtype::PType::I32;
+    use vortex_dtype::{DType, StructFields};
+
+    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::transform::rules::{ReduceRule, TypedRuleContext};
+
+    #[test]
+    fn test_remove_select_rule() {
+        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 = TypedRuleContext::new(dtype.clone());
+        let select_view = e.as_::<Select>();
+        let result = rule.reduce(&select_view, &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_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());
+
+        let rule = RemoveSelectRule;
+        let ctx = TypedRuleContext::new(dtype.clone());
+        let select_view = e.as_::<Select>();
+        let result = rule.reduce(&select_view, &ctx).unwrap();
+
+        assert!(result.is_some());
+        let transformed = result.unwrap();
+        assert!(transformed.is::<Pack>());
+
+        // 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().as_ref(), &["a", "b"]);
+    }
+}