fix: add missing output names

Author: tokoko

src/substrait/extended_expression.py

@@ -10,6 +10,14 @@ def resolve(base_schema: stp.NamedStruct) -> stee.ExtendedExpression:
         column_index = list(base_schema.names).index(name)
         lengths = [type_num_names(t) for t in base_schema.struct.types]
         flat_indices = [0] + list(itertools.accumulate(lengths))[:-1]
+        field_index = flat_indices.index(column_index)
+
+        names_start = flat_indices[field_index]
+        names_end = (
+            flat_indices[field_index + 1]
+            if len(flat_indices) > field_index + 1
+            else None
+        )
 
         return stee.ExtendedExpression(
             referred_expr=[
@@ -19,11 +27,12 @@ def resolve(base_schema: stp.NamedStruct) -> stee.ExtendedExpression:
                             root_reference=stalg.Expression.FieldReference.RootReference(),
                             direct_reference=stalg.Expression.ReferenceSegment(
                                 struct_field=stalg.Expression.ReferenceSegment.StructField(
-                                    field=flat_indices.index(column_index)
+                                    field=field_index
                                 )
                             ),
                         )
-                    )
+                    ),
+                    output_names=list(base_schema.names)[names_start:names_end],
                 )
             ],
             base_schema=base_schema,

tests/test_extended_expression.py

@@ -55,7 +55,8 @@ def test_column_no_nesting():
                             )
                         ),
                     )
-                )
+                ),
+                output_names=["description"],
             )
         ],
         base_schema=named_struct,
@@ -75,7 +76,29 @@ def test_column_nesting():
                             )
                         ),
                     )
-                )
+                ),
+                output_names=["order_total"],
+            )
+        ],
+        base_schema=nested_named_struct,
+    )
+
+
+def test_column_nested_struct():
+    assert column("shop_details")(nested_named_struct) == stee.ExtendedExpression(
+        referred_expr=[
+            stee.ExpressionReference(
+                expression=stalg.Expression(
+                    selection=stalg.Expression.FieldReference(
+                        root_reference=stalg.Expression.FieldReference.RootReference(),
+                        direct_reference=stalg.Expression.ReferenceSegment(
+                            struct_field=stalg.Expression.ReferenceSegment.StructField(
+                                field=1
+                            )
+                        ),
+                    )
+                ),
+                output_names=["shop_details", "shop_id", "shop_total"],
             )
         ],
         base_schema=nested_named_struct,