Avoid expanding unreferenced struct plan fields

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-layout/src/plan/plans/pack.rs

@@ -352,8 +352,7 @@ impl PlanParentReduceRule<Pack> for ExpressionPackRule {
                 .get(&ExactBoundExpr(expression.clone()))
                 .vortex_expect("Bound expression missing free-field annotations")
                 .clone();
-        let expanded_root = expanded_struct_root(child.dtype(), fields)?;
-        let expanded = expand_struct_root(expression.clone(), &expanded_root, fields)?;
+        let expanded = expand_struct_root(expression.clone(), fields)?;
         let partitioned =
             partition_bound(expanded.clone(), make_bound_free_field_annotator(fields))?;
 
@@ -469,14 +468,13 @@ fn expanded_struct_root(
 
 fn expand_struct_root(
     expression: BoundExpression,
-    expanded_root: &BoundExpression,
     fields: &StructFields,
 ) -> VortexResult<BoundExpression> {
     Ok(expression
         .transform_down(|node| {
             if node.is_root() {
                 return Ok(Transformed {
-                    value: expanded_root.clone(),
+                    value: expanded_struct_root(node.dtype(), fields)?,
                     changed: true,
                     order: TraversalOrder::Skip,
                 });
@@ -493,28 +491,23 @@ fn expand_struct_root(
                 return Ok(Transformed::no(node));
             }
 
-            if let Some(field_name) = scalar_fn.as_opt::<GetItem>() {
-                let index = fields.find(field_name).ok_or_else(|| {
-                    vortex_err!("Field {field_name} not found while expanding struct root")
-                })?;
+            if scalar_fn.is::<GetItem>() {
                 return Ok(Transformed {
-                    value: expanded_root.children()[index].clone(),
-                    changed: true,
+                    value: node,
+                    changed: false,
                     order: TraversalOrder::Skip,
                 });
             }
 
             if let Some(selection) = scalar_fn.as_opt::<Select>() {
                 let names = selection.normalize_to_included_fields(fields.names())?;
+                let root = node.children()[0].clone();
                 let children = names
                     .iter()
                     .map(|name| {
-                        let index = fields
-                            .find(name)
-                            .vortex_expect("normalized selection fields must exist in the root");
-                        expanded_root.children()[index].clone()
+                        BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()])
                     })
-                    .collect();
+                    .collect::<VortexResult<Vec<_>>>()?;
                 return Ok(Transformed {
                     value: bound_pack(names, children)?,
                     changed: true,

vortex-layout/src/plan/tests.rs

@@ -30,6 +30,8 @@ use vortex_array::expr::is_null;
 use vortex_array::expr::lit;
 use vortex_array::expr::pack;
 use vortex_array::expr::root;
+use vortex_array::expr::select;
+use vortex_array::expr::select_exclude;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 use vortex_io::runtime::single::block_on;
@@ -279,6 +281,72 @@ fn struct_plan_appends_validity_when_nullable() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn struct_select_preserves_struct_output() -> VortexResult<()> {
+    let field_dtype = primitive(PType::I32, Nullability::NonNullable);
+    let layout = StructLayout::new(
+        3,
+        DType::Struct(
+            StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]),
+            Nullability::NonNullable,
+        ),
+        vec![
+            flat(3, field_dtype.clone(), 0),
+            flat(3, field_dtype.clone(), 1),
+        ],
+    )
+    .into_layout();
+
+    let plan = make_eval(select(["a"], root()), make_plan(layout)?)?.into_plan();
+    let optimized = optimize(plan)?;
+
+    assert_eq!(
+        optimized.dtype(),
+        &DType::Struct(
+            StructFields::from_iter([("a", field_dtype)]),
+            Nullability::NonNullable,
+        )
+    );
+    insta::assert_snapshot!(optimized.display_tree(), @r"
+    root: vortex.plan.eval({a=i32}, rows=3) expr=pack(a: $)
+      child: vortex.plan.segment_scan(i32, rows=3)
+    ");
+    Ok(())
+}
+
+#[test]
+fn struct_select_exclude_preserves_struct_output() -> VortexResult<()> {
+    let field_dtype = primitive(PType::I32, Nullability::NonNullable);
+    let layout = StructLayout::new(
+        3,
+        DType::Struct(
+            StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]),
+            Nullability::NonNullable,
+        ),
+        vec![
+            flat(3, field_dtype.clone(), 0),
+            flat(3, field_dtype.clone(), 1),
+        ],
+    )
+    .into_layout();
+
+    let plan = make_eval(select_exclude(["b"], root()), make_plan(layout)?)?.into_plan();
+    let optimized = optimize(plan)?;
+
+    assert_eq!(
+        optimized.dtype(),
+        &DType::Struct(
+            StructFields::from_iter([("a", field_dtype)]),
+            Nullability::NonNullable,
+        )
+    );
+    insta::assert_snapshot!(optimized.display_tree(), @r"
+    root: vortex.plan.eval({a=i32}, rows=3) expr=pack(a: $)
+      child: vortex.plan.segment_scan(i32, rows=3)
+    ");
+    Ok(())
+}
+
 #[test]
 fn with_children_rejects_mismatched_arity() -> VortexResult<()> {
     let dtype = primitive(PType::I32, Nullability::NonNullable);