fix!: carry struct columns' full names through select, project and aggregate

src/substrait/builders/extended_expression.py

@@ -577,7 +577,7 @@ def column(field: str | int, alias: Iterable[str] | str | None = None):
 
     Accepts either an index or a field name of a desired field.
     """
-    alias = [alias] if alias and isinstance(alias, str) else alias
+    alias = [alias] if alias and isinstance(alias, str) else list(alias or [])
 
     def resolve(
         base_schema: stp.NamedStruct, registry: ExtensionRegistry
@@ -607,6 +607,21 @@ def resolve(
             else None
         )
 
+        output_names = list(base_schema.names)[names_start:names_end]
+        if len(alias) == 1:
+            # A single alias renames only the top-level field, keeping a struct
+            # column's nested names.
+            output_names = [alias[0], *output_names[1:]]
+        elif alias:
+            # A full list replaces every name, so it must name each field.
+            if len(alias) != len(output_names):
+                raise ValueError(
+                    f"column() alias {alias!r} has {len(alias)} names, but field "
+                    f"{field!r} needs {len(output_names)} "
+                    f"({output_names!r})."
+                )
+            output_names = list(alias)
+
         return stee.ExtendedExpression(
             referred_expr=[
                 stee.ExpressionReference(
@@ -620,9 +635,7 @@ def resolve(
                             ),
                         )
                     ),
-                    output_names=list(base_schema.names)[names_start:names_end]
-                    if not alias
-                    else alias,
+                    output_names=output_names,
                 )
             ],
             base_schema=base_schema,

src/substrait/builders/plan.py

@@ -38,6 +38,7 @@
     plan_subtrees,
     rebase_reference_ordinals,
     remap_function_references,
+    type_num_names,
 )
 from substrait.version import substrait_version
 
@@ -390,6 +391,27 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
     return build_scoped(resolve)
 
 
+def _expression_output_names(
+    ref: stee.ExpressionReference, ns: stt.NamedStruct, registry: ExtensionRegistry
+) -> list[str]:
+    """An expression's depth-first output names, checked against its output type."""
+    names = list(ref.output_names)
+    if not names:
+        raise ValueError("Expression has no output names; give it an alias.")
+    if len(names) == 1:
+        # Struct-valued expressions other than column() still emit one name (#312).
+        return names
+    needed = type_num_names(
+        infer_expression_type(ref.expression, ns.struct, registry=registry)
+    )
+    if len(names) != needed:
+        raise ValueError(
+            f"Expression named {names!r} has {len(names)} names, but its output "
+            f"type needs {needed}."
+        )
+    return names
+
+
 def project(
     plan: PlanOrUnbound,
     expressions: Iterable[ExtendedExpressionOrUnbound],
@@ -416,8 +438,13 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             resolve_expression(e, ns, registry) for e in expressions
         ]
 
+        # output_names holds an expression's names in depth-first order, so a
+        # struct-valued expression contributes its nested names too.
         names = list(_plan.relations[-1].root.names) + [
-            e.output_names[0] for ee in bound_expressions for e in ee.referred_expr
+            name
+            for ee in bound_expressions
+            for e in ee.referred_expr
+            for name in _expression_output_names(e, ns, registry)
         ]
 
         return _plan_from(
@@ -466,10 +493,16 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
             resolve_expression(e, ns, registry) for e in expressions
         ]
 
-        start_index = len(_plan.relations[-1].root.names)
+        # Expressions are appended after the input's top-level fields; root.names
+        # counts nested names too, so it is not a field count.
+        start_index = len(ns.struct.types)
+        exprs = [e.expression for ee in bound_expressions for e in ee.referred_expr]
 
         names = [
-            e.output_names[0] for ee in bound_expressions for e in ee.referred_expr
+            name
+            for ee in bound_expressions
+            for e in ee.referred_expr
+            for name in _expression_output_names(e, ns, registry)
         ]
 
         return _plan_from(
@@ -478,15 +511,11 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
                 project=stalg.ProjectRel(
                     common=stalg.RelCommon(
                         emit=stalg.RelCommon.Emit(
-                            output_mapping=[i + start_index for i in range(len(names))]
+                            output_mapping=[i + start_index for i in range(len(exprs))]
                         )
                     ),
                     input=inp[0],
-                    expressions=[
-                        e.expression
-                        for ee in bound_expressions
-                        for e in ee.referred_expr
-                    ],
+                    expressions=exprs,
                     advanced_extension=extension,
                 )
             ),
@@ -973,7 +1002,9 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         )
 
         names = [
-            e.referred_expr[0].output_names[0] for e in bound_grouping_expressions
+            name
+            for e in bound_grouping_expressions
+            for name in _expression_output_names(e.referred_expr[0], ns, registry)
         ] + [e.referred_expr[0].output_names[0] for e in bound_measures]
 
         return _plan_from(

tests/builders/extended_expression/test_column.py

@@ -127,6 +127,46 @@ def test_column_by_ordinal_over_a_nested_struct():
     ]
 
 
+def test_column_alias_over_a_nested_struct_keeps_member_names():
+    # A single alias renames the struct column itself; its members keep their names
+    # so the output still carries one name per field in depth-first order.
+    aliased = column("shop_details", alias="shop")(nested_named_struct, None)
+
+    assert list(aliased.referred_expr[0].output_names) == [
+        "shop",
+        "shop_id",
+        "shop_total",
+    ]
+
+
+def test_column_alias_list_replaces_all_names():
+    aliased = column("shop_details", alias=["shop", "id", "total"])(
+        nested_named_struct, None
+    )
+
+    assert list(aliased.referred_expr[0].output_names) == ["shop", "id", "total"]
+
+
+@pytest.mark.parametrize(
+    "field, alias",
+    [
+        ("order_id", ["x", "y"]),
+        ("shop_details", ["shop", "id"]),
+    ],
+)
+def test_column_alias_list_must_name_every_field(field, alias):
+    # A list alias replaces the depth-first names wholesale, so a wrong length would
+    # give the output schema more or fewer names than it has fields.
+    with pytest.raises(ValueError, match="needs"):
+        column(field, alias=alias)(nested_named_struct, None)
+
+
+def test_column_empty_alias_keeps_schema_name():
+    aliased = column("order_id", alias="")(named_struct, None)
+
+    assert list(aliased.referred_expr[0].output_names) == ["order_id"]
+
+
 @pytest.mark.parametrize("name", ["missing", "shop_id"])
 def test_column_unknown_name_raises(name):
     # A name that is absent, or only names a nested field, is not a top-level column.
@@ -172,3 +212,11 @@ def test_column_by_name_after_a_map_column():
 def test_column_by_ordinal_out_of_range_raises(index):
     with pytest.raises(ValueError, match=r"out of range; schema has 3 columns"):
         column(index)(named_struct, None)
+
+
+def test_column_alias_list_on_a_field_without_names_raises_value_error():
+    # A schema can carry fewer names than fields; the error must still be reported.
+    unnamed = stt.NamedStruct(names=["order_id"], struct=struct)
+
+    with pytest.raises(ValueError, match="2 needs 0"):
+        column(2, alias=["g", "h"])(unnamed, None)

tests/builders/plan/test_aggregate.py

@@ -117,3 +117,38 @@ def test_aggregate_rejects_a_non_aggregate_measure():
         aggregate(table, grouping_expressions=[column("id")], measures=[column("id")])(
             registry
         )
+
+
+def test_aggregate_group_by_struct_column():
+    # Grouping by a struct column keeps its nested names, one per field depth-first.
+    nested = stt.NamedStruct(
+        names=["s", "x", "y", "b"],
+        struct=stt.Type.Struct(
+            types=[
+                stt.Type(
+                    struct=stt.Type.Struct(
+                        types=[i64(nullable=False), i64(nullable=False)],
+                        nullability=stt.Type.NULLABILITY_REQUIRED,
+                    )
+                ),
+                boolean(),
+            ],
+            nullability=stt.Type.NULLABILITY_REQUIRED,
+        ),
+    )
+    table = read_named_table("table", nested)
+    measure_expr = aggregate_function(
+        "extension:test:urn", "count", expressions=[column("b")], alias=["count"]
+    )
+
+    actual = aggregate(
+        table, grouping_expressions=[column("s")], measures=[measure_expr]
+    )(registry)
+
+    assert list(actual.relations[-1].root.names) == ["s", "x", "y", "count"]
+    assert list(infer_plan_schema(actual, registry=registry).names) == [
+        "s",
+        "x",
+        "y",
+        "count",
+    ]

tests/builders/plan/test_project.py

@@ -1,11 +1,13 @@
+import pytest
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
 
-from substrait.builders.extended_expression import column
+from substrait.builders.extended_expression import cast, column, literal
 from substrait.builders.plan import default_version, project, read_named_table, select
 from substrait.builders.type import boolean, i64
 from substrait.extension_registry import ExtensionRegistry
+from substrait.type_inference import infer_plan_schema
 
 registry = ExtensionRegistry(load_default_extensions=False)
 
@@ -89,3 +91,97 @@ def test_select():
     )
 
     assert actual == expected
+
+
+# A struct column takes one top-level field but several depth-first names.
+nested_named_struct = stt.NamedStruct(
+    names=["a", "s", "x", "y", "b"],
+    struct=stt.Type.Struct(
+        types=[
+            i64(nullable=False),
+            stt.Type(
+                struct=stt.Type.Struct(
+                    types=[i64(nullable=False), i64(nullable=False)],
+                    nullability=stt.Type.NULLABILITY_REQUIRED,
+                )
+            ),
+            i64(nullable=False),
+        ],
+        nullability=stt.Type.NULLABILITY_REQUIRED,
+    ),
+)
+
+
+def test_select_struct_column():
+    table = read_named_table("table", nested_named_struct)
+
+    actual = select(table, [column("a"), column("s")])(registry)
+    root = actual.relations[-1].root
+
+    # The projected expressions sit after the input's 3 top-level fields, not after
+    # its 5 names, and the struct column keeps its nested names.
+    assert list(root.input.project.common.emit.output_mapping) == [3, 4]
+    assert list(root.names) == ["a", "s", "x", "y"]
+    assert list(infer_plan_schema(actual).names) == ["a", "s", "x", "y"]
+
+
+def test_project_struct_column():
+    table = read_named_table("table", nested_named_struct)
+
+    actual = project(table, [column("s")])(registry)
+
+    assert list(actual.relations[-1].root.names) == [
+        *nested_named_struct.names,
+        "s",
+        "x",
+        "y",
+    ]
+    assert len(infer_plan_schema(actual).struct.types) == 4
+
+
+@pytest.mark.parametrize("builder", [project, select])
+def test_expression_with_more_names_than_fields_is_rejected(builder):
+    # Every output name lands in the plan root, so a scalar given two names would
+    # leave the root with more names than the output type has fields.
+    table = read_named_table("table", named_struct)
+
+    with pytest.raises(ValueError, match="has 2 names"):
+        builder(table, [literal(1, i64(), alias=["one", "two"])])(registry)
+
+
+def test_select_struct_valued_expression_keeps_full_alias():
+    table = read_named_table("table", nested_named_struct)
+    pair = stt.Type(
+        struct=stt.Type.Struct(
+            types=[i64(), i64()], nullability=stt.Type.NULLABILITY_NULLABLE
+        )
+    )
+
+    actual = select(table, [cast(column("s"), pair, alias=["c", "p", "q"])])(registry)
+
+    assert list(actual.relations[-1].root.names) == ["c", "p", "q"]
+
+
+def test_select_struct_valued_expression_with_partial_alias_is_rejected():
+    # Too few names would shift every later column's name onto the wrong field.
+    table = read_named_table("table", nested_named_struct)
+    pair = stt.Type(
+        struct=stt.Type.Struct(
+            types=[i64(), i64()], nullability=stt.Type.NULLABILITY_NULLABLE
+        )
+    )
+
+    with pytest.raises(ValueError, match="has 2 names, but its output type needs 3"):
+        select(table, [cast(column("s"), pair, alias=["c", "p"]), column("a")])(
+            registry
+        )
+
+
+@pytest.mark.parametrize("builder", [project, select])
+def test_expression_without_output_names_is_rejected(builder):
+    table = read_named_table("table", named_struct)
+    unnamed = column("id")(named_struct, registry)
+    del unnamed.referred_expr[0].output_names[:]
+
+    with pytest.raises(ValueError, match="no output names"):
+        builder(table, [unnamed])(registry)