Expression optimisation rules don't introduce irreducible children

vortex-array/src/expr/optimize.rs

@@ -210,6 +210,7 @@ mod tests {
     use vortex_error::vortex_err;
 
     use crate::dtype::DType;
+    use crate::dtype::FieldName;
     use crate::dtype::Nullability;
     use crate::dtype::PType;
     use crate::dtype::StructFields;
@@ -218,8 +219,13 @@ mod tests {
     use crate::expr::get_item;
     use crate::expr::lit;
     use crate::expr::lt_eq;
+    use crate::expr::merge;
     use crate::expr::or;
+    use crate::expr::pack;
     use crate::expr::root;
+    use crate::expr::select;
+    use crate::expr::transform::replace;
+    use crate::expr::transform::replace_root_fields;
     use crate::scalar::Scalar;
     use crate::scalar_fn::fns::literal::Literal;
 
@@ -272,4 +278,78 @@ mod tests {
         assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable));
         Ok(())
     }
+
+    /// `Select` over a `pack` rewrites to a `pack` of the selected child fields. It must read
+    /// those fields out of the child `pack` itself, because the optimizer visits children before
+    /// their parent and so never revisits the children a rule introduces.
+    #[test]
+    fn optimize_select_of_pack_reads_pack_fields() -> VortexResult<()> {
+        let field_names = (0..4)
+            .map(|idx| FieldName::from(format!("field_{idx}")))
+            .collect::<Vec<_>>();
+        let fields = StructFields::new(
+            field_names.clone().into(),
+            vec![DType::Primitive(PType::U64, Nullability::NonNullable); field_names.len()],
+        );
+        let scope = DType::Struct(fields.clone(), Nullability::NonNullable);
+
+        let expr = select(
+            field_names
+                .iter()
+                .cloned()
+                .chain(["input_row_idx".into(), "input_file_idx".into()])
+                .collect::<Vec<_>>(),
+            merge([
+                root(),
+                pack(
+                    [
+                        ("input_row_idx", lit(0_u64)),
+                        ("input_file_idx", lit(1_u64)),
+                    ],
+                    Nullability::NonNullable,
+                ),
+            ]),
+        );
+        let expr = replace(expr, &root(), replace_root_fields(root(), &fields));
+
+        let expected = pack(
+            field_names
+                .into_iter()
+                .map(|name| (name.clone(), get_item(name, root())))
+                .chain([
+                    ("input_row_idx".into(), lit(0_u64)),
+                    ("input_file_idx".into(), lit(1_u64)),
+                ])
+                .collect::<Vec<_>>(),
+            Nullability::NonNullable,
+        );
+
+        let optimized = expr.optimize_recursive(&scope)?;
+        assert_eq!(optimized, expected, "got {optimized}");
+        assert_eq!(optimized.optimize_recursive(&scope)?, expected);
+
+        Ok(())
+    }
+
+    /// `Merge` reduces to a `pack` that reads each field out of its children, so it must also
+    /// read straight out of a `pack` child.
+    #[test]
+    fn optimize_merge_of_packs_reads_pack_fields() -> VortexResult<()> {
+        let scope = DType::Struct(StructFields::empty(), Nullability::NonNullable);
+        let expr = merge([
+            pack([("a", lit(0_u64))], Nullability::NonNullable),
+            pack([("b", lit(1_u64))], Nullability::NonNullable),
+        ]);
+
+        let expected = pack(
+            [("a", lit(0_u64)), ("b", lit(1_u64))],
+            Nullability::NonNullable,
+        );
+
+        let optimized = expr.optimize_recursive(&scope)?;
+        assert_eq!(optimized, expected, "got {optimized}");
+        assert_eq!(optimized.optimize_recursive(&scope)?, expected);
+
+        Ok(())
+    }
 }

vortex-array/src/scalar_fn/fns/get_item.rs

@@ -3,6 +3,7 @@
 
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::slice;
 
 use prost::Message;
 use vortex_error::VortexResult;
@@ -51,6 +52,57 @@ impl GetItem {
     ) -> VortexResult<ScalarFnArray> {
         ScalarFnArray::try_new(GetItem.bind(field_name.into()), vec![input])
     }
+
+    /// Read `field_name` out of `struct_node`, without a `get_item` node where possible.
+    ///
+    /// Rules that read a field out of their own child use this instead of building the
+    /// `get_item` node themselves. The optimizer visits children before their parent, so a
+    /// `get_item` that a rule introduces is never simplified again. Removing it here keeps the
+    /// output of the rule a fixed point.
+    pub(crate) fn project_node<T: ReduceNode>(
+        struct_node: &T,
+        field_name: &FieldName,
+    ) -> VortexResult<T> {
+        if let Some(field) = Self::project_pack_node(struct_node, field_name)? {
+            return Ok(field);
+        }
+
+        struct_node.new_node(
+            GetItem.bind(field_name.clone()),
+            slice::from_ref(struct_node),
+        )
+    }
+
+    /// Read `field_name` straight out of a `pack` node, which holds each field as a child.
+    ///
+    /// Returns `None` if `struct_node` is not a `pack`, or does not hold `field_name`.
+    fn project_pack_node<T: ReduceNode>(
+        struct_node: &T,
+        field_name: &FieldName,
+    ) -> VortexResult<Option<T>> {
+        let Some(pack) = struct_node
+            .scalar_fn()
+            .and_then(|scalar_fn| scalar_fn.as_opt::<Pack>())
+        else {
+            return Ok(None);
+        };
+        let Some(idx) = pack.names.find(field_name) else {
+            return Ok(None);
+        };
+
+        let mut field = struct_node.child(idx);
+
+        // A nullable pack intersects its validity into the field, so mask the field with an
+        // all-true literal to make its dtype nullable.
+        if pack.nullability.is_nullable() {
+            field = struct_node.new_node(
+                Mask.bind(EmptyOptions),
+                &[field, struct_node.new_node(Literal.bind(true.into()), &[])?],
+            )?;
+        }
+
+        Ok(Some(field))
+    }
 }
 
 impl ScalarFnVTable for GetItem {
@@ -136,25 +188,7 @@ impl ScalarFnVTable for GetItem {
     }
 
     fn reduce<T: ReduceNode>(&self, field_name: &FieldName, node: &T) -> VortexResult<Option<T>> {
-        let child = node.child(0);
-        if let Some(child_fn) = child.scalar_fn()
-            && let Some(pack) = child_fn.as_opt::<Pack>()
-            && let Some(idx) = pack.names.find(field_name)
-        {
-            let mut field = child.child(idx);
-
-            // Possibly mask the field if the pack is nullable
-            if pack.nullability.is_nullable() {
-                field = node.new_node(
-                    Mask.bind(EmptyOptions),
-                    &[field, node.new_node(Literal.bind(true.into()), &[])?],
-                )?;
-            }
-
-            return Ok(Some(field));
-        }
-
-        Ok(None)
+        Self::project_pack_node(&node.child(0), field_name)
     }
 
     fn simplify_untyped(

vortex-array/src/scalar_fn/fns/merge.rs

@@ -175,12 +175,12 @@ impl ScalarFnVTable for Merge {
 
     fn reduce<T: ReduceNode>(&self, options: &Self::Options, node: &T) -> VortexResult<Option<T>> {
         let mut names = Vec::with_capacity(node.child_count() * 2);
-        let mut children = Vec::with_capacity(node.child_count() * 2);
+        let mut fields: Vec<T> = Vec::with_capacity(node.child_count() * 2);
         let mut duplicate_names = HashSet::<_>::new();
 
         for child in (0..node.child_count()).map(|i| node.child(i)) {
             let child_dtype = child.node_dtype()?;
-            if !child_dtype.is_struct() {
+            if !child_dtype.is_struct() || child_dtype.is_nullable() {
                 vortex_bail!(
                     "Merge child must return a non-nullable struct dtype, got {}",
                     child_dtype
@@ -192,12 +192,18 @@ impl ScalarFnVTable for Merge {
                 .vortex_expect("expected struct");
 
             for name in child_dtype.names().iter() {
+                // `project_node` reads straight out of a `pack` child, so this reduction leaves
+                // behind no `get_item` for a later pass to simplify. The child is non-nullable,
+                // so the read cannot intersect a struct validity into the field, which is the
+                // only way it could change the field dtype.
+                let field = GetItem::project_node(&child, name)?;
+
                 if let Some(idx) = names.iter().position(|n| n == name) {
                     duplicate_names.insert(name.clone());
-                    children[idx] = child.clone();
+                    fields[idx] = field;
                 } else {
                     names.push(name.clone());
-                    children.push(child.clone());
+                    fields.push(field);
                 }
             }
 
@@ -209,18 +215,12 @@ impl ScalarFnVTable for Merge {
             }
         }
 
-        let pack_children: Vec<_> = names
-            .iter()
-            .zip(children)
-            .map(|(name, child)| node.new_node(GetItem.bind(name.clone()), &[child]))
-            .try_collect()?;
-
         let pack_expr = node.new_node(
             Pack.bind(PackOptions {
                 names: FieldNames::from(names),
                 nullability: node.node_dtype()?.nullability(),
             }),
-            &pack_children,
+            &fields,
         )?;
 
         Ok(Some(pack_expr))

vortex-array/src/scalar_fn/fns/select.rs

@@ -6,6 +6,7 @@ use std::fmt::Formatter;
 
 use itertools::Itertools;
 use prost::Message;
+use vortex_error::VortexError;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -27,7 +28,6 @@ use crate::dtype::FieldNames;
 use crate::expr::display::ExprDisplay;
 use crate::expr::expression::Expression;
 use crate::expr::field::DisplayFieldNames;
-use crate::expr::get_item;
 use crate::expr::pack;
 use crate::scalar_fn::Arity;
 use crate::scalar_fn::ChildName;
@@ -205,29 +205,41 @@ impl ScalarFnVTable for Select {
             return Ok(Some(pack(empty, struct_nullability)));
         }
 
-        // We cannot always convert a `select` into a `pack(get_item(f1), get_item(f2), ...)`.
-        // This is because `get_item` does a validity intersection of the struct validity with its
-        // fields, which is not the same as just "masking" out the unwanted fields (a selection).
+        // We cannot always convert a `select` into a projection of its child fields. This is
+        // because reading a field does a validity intersection of the struct validity with the
+        // field, which is not the same as just "masking" out the unwanted fields (a selection).
         //
         // We can, however, make this simplification when the child of the `select` is already a
-        // `pack` and we know that `get_item` will do no validity intersections.
-        let child_is_pack = child_struct.is::<Pack>();
-
-        // `get_item` only performs validity intersection when the struct is nullable but the field
-        // is not. This would change the semantics of a `select`, so we can only simplify when this
-        // won't happen.
+        // `pack` and we know that reading a field will do no validity intersections.
+        //
+        // The intersection only happens when the struct is nullable but the field is not. That
+        // would change the semantics of a `select`, so we can only simplify when it won't happen.
         let would_intersect_validity =
             struct_nullability.is_nullable() && !all_included_fields_are_nullable;
 
-        if child_is_pack && !would_intersect_validity {
-            let pack_expr = pack(
-                included_fields
-                    .into_iter()
-                    .map(|name| (name.clone(), get_item(name, child_struct.clone()))),
-                struct_nullability,
-            );
-
-            return Ok(Some(pack_expr));
+        if let Some(child_pack) = child_struct.as_opt::<Pack>()
+            && !would_intersect_validity
+        {
+            // Take the field expressions straight out of the child `pack` rather than wrap the
+            // `pack` in `get_item`. The guard above proves the two are equivalent here, and the
+            // optimizer visits children before their parent, so a `get_item` introduced now
+            // would never be simplified away.
+            let fields: Vec<_> = included_fields
+                .iter()
+                .map(|name| {
+                    let idx = child_pack.names.find(name).ok_or_else(|| {
+                        vortex_err!(
+                            "Cannot find field {} in pack fields {:?}",
+                            name,
+                            child_pack.names
+                        )
+                    })?;
+
+                    Ok((name.clone(), child_struct.child(idx).clone()))
+                })
+                .try_collect::<_, _, VortexError>()?;
+
+            return Ok(Some(pack(fields, struct_nullability)));
         }
 
         Ok(None)