diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index fd13d7a..ce0bf9b 100644 --- a/src/substrait/builders/extended_expression.py +++ b/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, diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 60ed0e4..6525fd2 100644 --- a/src/substrait/builders/plan.py +++ b/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( diff --git a/tests/builders/extended_expression/test_column.py b/tests/builders/extended_expression/test_column.py index 2fd4de5..e5938b8 100644 --- a/tests/builders/extended_expression/test_column.py +++ b/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) diff --git a/tests/builders/plan/test_aggregate.py b/tests/builders/plan/test_aggregate.py index b6aca40..fa00dfb 100644 --- a/tests/builders/plan/test_aggregate.py +++ b/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", + ] diff --git a/tests/builders/plan/test_project.py b/tests/builders/plan/test_project.py index bca68a5..ec6314d 100644 --- a/tests/builders/plan/test_project.py +++ b/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)