Avoid expanding unreferenced struct plan fields

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

Author: joseph-isaacs

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

@@ -245,8 +245,7 @@ impl PlanParentReduceRule<StructPlan> for ExpressionStructRule {
                 .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))?;
         if partitioned.partition_names.is_empty() {
@@ -368,14 +367,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,
                 });
@@ -392,28 +390,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 struct 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

@@ -27,6 +27,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;
@@ -282,6 +284,71 @@ fn struct_plan_optimization_visits_all_fields() -> 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 optimized = make_expression_plan(select(["a"], root()), make_plan(layout)?)?.optimize()?;
+
+    assert_eq!(
+        optimized.dtype(),
+        &DType::Struct(
+            StructFields::from_iter([("a", field_dtype)]),
+            Nullability::NonNullable,
+        )
+    );
+    insta::assert_snapshot!(optimized.tree_display(), @r"
+    root: ExpressionPlan({a=i32}, rows=3) expr=pack(a: $)
+      child: FlatPlan(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 optimized =
+        make_expression_plan(select_exclude(["b"], root()), make_plan(layout)?)?.optimize()?;
+
+    assert_eq!(
+        optimized.dtype(),
+        &DType::Struct(
+            StructFields::from_iter([("a", field_dtype)]),
+            Nullability::NonNullable,
+        )
+    );
+    insta::assert_snapshot!(optimized.tree_display(), @r"
+    root: ExpressionPlan({a=i32}, rows=3) expr=pack(a: $)
+      child: FlatPlan(i32, rows=3)
+    ");
+    Ok(())
+}
+
 #[test]
 fn chunked_plan_optimization_visits_all_chunks() -> VortexResult<()> {
     let dtype = primitive(PType::I32, Nullability::NonNullable);