From f5f53a0291ae31a939a0b29bdb152432453612a5 Mon Sep 17 00:00:00 2001
From: Joe Isaacs <joe.isaacs@live.co.uk>
Date: Fri, 7 Aug 2026 16:32:13 +0100
Subject: [PATCH] Avoid expanding unreferenced struct plan fields

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
---
 vortex-layout/src/plan/plans/pack.rs | 23 ++++------
 vortex-layout/src/plan/tests.rs      | 68 ++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 15 deletions(-)

diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs
index 709688f6c9c..a229b73bb0e 100644
--- a/vortex-layout/src/plan/plans/pack.rs
+++ b/vortex-layout/src/plan/plans/pack.rs
@@ -258,8 +258,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))?;
 
@@ -377,14 +376,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,
                 });
@@ -401,28 +399,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,
diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs
index 6a9dd1fc6e3..8a9b5740133 100644
--- a/vortex-layout/src/plan/tests.rs
+++ b/vortex-layout/src/plan/tests.rs
@@ -29,6 +29,8 @@ use vortex_array::expr::gt;
 use vortex_array::expr::is_null;
 use vortex_array::expr::lit;
 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;
@@ -275,6 +277,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);