diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index f692a746..590237ea 100644 --- a/src/substrait/builders/extended_expression.py +++ b/src/substrait/builders/extended_expression.py @@ -3,6 +3,7 @@ import contextvars import itertools import uuid as uuid_module +from dataclasses import dataclass from datetime import date, datetime, time, timedelta, timezone from decimal import Decimal from typing import Any, Callable, Iterable, Union @@ -55,6 +56,24 @@ def fresh_rel_anchors(): ExtendedExpressionOrUnbound = Union[stee.ExtendedExpression, UnboundExtendedExpression] +@dataclass(frozen=True) +class EnumerationArgument: + """A simple-extension function enumeration argument. + + Enumeration arguments participate in overload resolution but are encoded as + ``FunctionArgument.enum`` rather than record expressions. + """ + + value: str + + def __post_init__(self): + if not self.value: + raise ValueError("enumeration argument must be non-empty") + + +FunctionArgumentOrExpression = Union[ExtendedExpressionOrUnbound, EnumerationArgument] + + def _alias_or_inferred( alias: Union[Iterable[str], str, None], op: str, @@ -80,6 +99,42 @@ def _function_options(options): return result +def _resolve_function_arguments( + arguments: Iterable[FunctionArgumentOrExpression], + base_schema: stp.NamedStruct, + registry: ExtensionRegistry, +): + bound = [ + argument + if isinstance(argument, EnumerationArgument) + else resolve_expression(argument, base_schema, registry) + for argument in arguments + ] + expression_arguments = [ + argument for argument in bound if isinstance(argument, stee.ExtendedExpression) + ] + expression_schemas = [ + infer_extended_expression_schema(argument, registry=registry) + for argument in expression_arguments + ] + schema_iterator = iter(expression_schemas) + signature = [] + protobuf_arguments = [] + output_names = [] + for argument in bound: + if isinstance(argument, EnumerationArgument): + signature.append(argument.value) + protobuf_arguments.append(stalg.FunctionArgument(enum=argument.value)) + output_names.append(argument.value) + else: + signature.extend(next(schema_iterator).types) + protobuf_arguments.append( + stalg.FunctionArgument(value=argument.referred_expr[0].expression) + ) + output_names.append(argument.referred_expr[0].output_names[0]) + return expression_arguments, signature, protobuf_arguments, output_names + + def resolve_expression( expression: ExtendedExpressionOrUnbound, base_schema: stp.NamedStruct, @@ -389,6 +444,49 @@ def resolve( return resolve +def nested_struct( + expressions: Iterable[ExtendedExpressionOrUnbound], + nullable: bool = False, + alias: Union[Iterable[str], str, None] = None, +) -> UnboundExtendedExpression: + """Build a struct-valued nested expression from scalar expressions.""" + + def resolve( + base_schema: stp.NamedStruct, registry: ExtensionRegistry + ) -> stee.ExtendedExpression: + bound = [resolve_expression(item, base_schema, registry) for item in expressions] + return stee.ExtendedExpression( + referred_expr=[ + stee.ExpressionReference( + expression=stalg.Expression( + nested=stalg.Expression.Nested( + nullable=nullable, + struct=stalg.Expression.Nested.Struct( + fields=[ + item.referred_expr[0].expression for item in bound + ] + ), + ) + ), + output_names=_alias_or_inferred( + alias, + "struct", + [item.referred_expr[0].output_names[0] for item in bound], + ), + ) + ], + base_schema=base_schema, + extension_urns=merge_extension_urns( + *[item.extension_urns for item in bound] + ), + extensions=merge_extension_declarations( + *[item.extensions for item in bound] + ), + ) + + return resolve + + def outer_reference(field: Union[str, int], steps_out: int = 1): """A field reference to an enclosing query's column (a correlated reference). @@ -532,7 +630,7 @@ def resolve( def scalar_function( urn: str, function: str, - expressions: Iterable[ExtendedExpressionOrUnbound], + expressions: Iterable[FunctionArgumentOrExpression], alias: Union[Iterable[str], str, None] = None, options: Union[dict, None] = None, ): @@ -545,16 +643,12 @@ def scalar_function( def resolve( base_schema: stp.NamedStruct, registry: ExtensionRegistry ) -> stee.ExtendedExpression: - bound_expressions = [ - resolve_expression(e, base_schema, registry) for e in expressions - ] - - expression_schemas = [ - infer_extended_expression_schema(b, registry=registry) - for b in bound_expressions - ] - - signature = [typ for es in expression_schemas for typ in es.types] + ( + bound_expressions, + signature, + function_arguments, + argument_names, + ) = _resolve_function_arguments(expressions, base_schema, registry) func = registry.lookup_function(urn, function, signature) @@ -591,12 +685,7 @@ def resolve( expression=stalg.Expression( scalar_function=stalg.Expression.ScalarFunction( function_reference=func[0].anchor, - arguments=[ - stalg.FunctionArgument( - value=e.referred_expr[0].expression - ) - for e in bound_expressions - ], + arguments=function_arguments, options=_function_options(options), output_type=func[1], ) @@ -604,7 +693,7 @@ def resolve( output_names=_alias_or_inferred( alias, function, - [e.referred_expr[0].output_names[0] for e in bound_expressions], + argument_names, ), ) ], @@ -619,7 +708,7 @@ def resolve( def aggregate_function( urn: str, function: str, - expressions: Iterable[ExtendedExpressionOrUnbound], + expressions: Iterable[FunctionArgumentOrExpression], alias: Union[Iterable[str], str, None] = None, invocation: Union[ "stalg.AggregateFunction.AggregationInvocation.ValueType", None @@ -638,21 +727,17 @@ def aggregate_function( def resolve( base_schema: stp.NamedStruct, registry: ExtensionRegistry ) -> stee.ExtendedExpression: - bound_expressions: Iterable[stee.ExtendedExpression] = [ - resolve_expression(e, base_schema, registry) for e in expressions - ] + ( + bound_expressions, + signature, + function_arguments, + argument_names, + ) = _resolve_function_arguments(expressions, base_schema, registry) bound_sorts = [ (resolve_expression(e, base_schema, registry), direction) for e, direction in sorts ] - expression_schemas = [ - infer_extended_expression_schema(b, registry=registry) - for b in bound_expressions - ] - - signature = [typ for es in expression_schemas for typ in es.types] - func = registry.lookup_function(urn, function, signature) if not func: @@ -691,10 +776,7 @@ def resolve( stee.ExpressionReference( measure=stalg.AggregateFunction( function_reference=func[0].anchor, - arguments=[ - stalg.FunctionArgument(value=e.referred_expr[0].expression) - for e in bound_expressions - ], + arguments=function_arguments, options=_function_options(options), output_type=func[1], invocation=invocation @@ -710,7 +792,7 @@ def resolve( output_names=_alias_or_inferred( alias, "IfThen", - [e.referred_expr[0].output_names[0] for e in bound_expressions], + argument_names, ), ) ], @@ -726,7 +808,7 @@ def resolve( def window_function( urn: str, function: str, - expressions: Iterable[ExtendedExpressionOrUnbound], + expressions: Iterable[FunctionArgumentOrExpression], partitions: Iterable[ExtendedExpressionOrUnbound] = [], alias: Union[Iterable[str], str, None] = None, options: Union[dict, None] = None, @@ -736,21 +818,17 @@ def window_function( def resolve( base_schema: stp.NamedStruct, registry: ExtensionRegistry ) -> stee.ExtendedExpression: - bound_expressions: Iterable[stee.ExtendedExpression] = [ - resolve_expression(e, base_schema, registry) for e in expressions - ] + ( + bound_expressions, + signature, + function_arguments, + argument_names, + ) = _resolve_function_arguments(expressions, base_schema, registry) bound_partitions = [ resolve_expression(e, base_schema, registry) for e in partitions ] - expression_schemas = [ - infer_extended_expression_schema(b, registry=registry) - for b in bound_expressions - ] - - signature = [typ for es in expression_schemas for typ in es.types] - func = registry.lookup_function(urn, function, signature) if not func: @@ -790,12 +868,7 @@ def resolve( expression=stalg.Expression( window_function=stalg.Expression.WindowFunction( function_reference=func[0].anchor, - arguments=[ - stalg.FunctionArgument( - value=e.referred_expr[0].expression - ) - for e in bound_expressions - ], + arguments=function_arguments, options=_function_options(options), output_type=func[1], partitions=[ @@ -806,7 +879,7 @@ def resolve( output_names=_alias_or_inferred( alias, function, - [e.referred_expr[0].output_names[0] for e in bound_expressions], + argument_names, ), ) ], diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index dff7d700..6fa7787a 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -5,6 +5,7 @@ See `examples/builder_example.py` for usage. """ +from contextvars import ContextVar import re from typing import Callable, Iterable, Optional, Union @@ -40,6 +41,10 @@ PlanOrUnbound = Union[stp.Plan, UnboundPlan] +_include_relation_output_names = ContextVar( + "include_relation_output_names", default=False +) + def _create_default_version(): p = re.compile(r"(\d+)\.(\d+)\.(\d+)") @@ -77,6 +82,75 @@ def _merge_plan_metadata(*objs): return metadata +def materialize( + plan: PlanOrUnbound, + registry: ExtensionRegistry, + *, + include_relation_output_names: bool = False, +) -> stp.Plan: + """Resolve a plan with optional standard output-name hints on every relation. + + ``RelCommon.Hint.output_names`` is the protocol-defined equivalent of + ``RelRoot.names`` for intermediate relations. Builders already infer these + names while composing a plan; this option preserves that information in the + serialized Plan without changing relation semantics or forcing shared + subtrees. The option is scoped to this materialization and is safe for + concurrent callers. + """ + + if isinstance(plan, stp.Plan): + return plan + token = _include_relation_output_names.set(include_relation_output_names) + try: + return plan(registry) + finally: + _include_relation_output_names.reset(token) + + +def with_relation_alias(plan: PlanOrUnbound, alias: str) -> UnboundPlan: + """Attach the standard ``RelCommon.Hint.alias`` to a plan's root relation. + + The alias is carried with the relation when a later builder nests it. This + makes logical relation identities available for qualification and debugging + without promoting ordinary single-consumer relations to shared subtrees. + """ + + if not alias: + raise ValueError("relation alias must not be empty") + + def resolve(registry: ExtensionRegistry) -> stp.Plan: + bound = plan if isinstance(plan, stp.Plan) else plan(registry) + result = stp.Plan() + result.CopyFrom(bound) + root_input = result.relations[-1].root.input + relation_type = root_input.WhichOneof("rel_type") + if relation_type is None: + raise ValueError("plan root does not contain a relation") + relation = getattr(root_input, relation_type) + if "common" not in relation.DESCRIPTOR.fields_by_name: + raise ValueError( + f"relation {relation_type} does not support RelCommon.Hint.alias" + ) + relation.common.hint.alias = alias + return result + + return resolve + + +def _with_output_names(rel: stalg.Rel, names: Iterable[str]) -> stalg.Rel: + if not _include_relation_output_names.get(): + return rel + relation_type = rel.WhichOneof("rel_type") + if relation_type is None: + return rel + relation = getattr(rel, relation_type) + if "common" not in relation.DESCRIPTOR.fields_by_name: + return rel + del relation.common.hint.output_names[:] + relation.common.hint.output_names.extend(names) + return rel + + def _is_identity(remap: dict) -> bool: return all(old == new for old, new in remap.items()) @@ -139,8 +213,12 @@ def _plan_from( propagation and Plan assembly live, so every relational builder is one call. """ subtree_planrels, input_rels = _merge_input_subtrees(bound_inputs) + output_names = list(names) root = stp.PlanRel( - root=stalg.RelRoot(input=make_rel(input_rels), names=list(names)) + root=stalg.RelRoot( + input=_with_output_names(make_rel(input_rels), output_names), + names=output_names, + ) ) kwargs = { "relations": [*subtree_planrels, root], @@ -189,14 +267,14 @@ def read_named_table( def resolve(registry: ExtensionRegistry) -> stp.Plan: _names = [names] if isinstance(names, str) else names - rel = stalg.Rel( + rel = _with_output_names(stalg.Rel( read=stalg.ReadRel( common=stalg.RelCommon(direct=stalg.RelCommon.Direct()), base_schema=named_struct, named_table=stalg.ReadRel.NamedTable(names=_names), advanced_extension=extension, ) - ) + ), named_struct.names) return stp.Plan( version=default_version, @@ -223,7 +301,10 @@ def _read_plan(named_struct: stt.NamedStruct, read_rel: stalg.ReadRel) -> stp.Pl relations=[ stp.PlanRel( root=stalg.RelRoot( - input=stalg.Rel(read=read_rel), names=named_struct.names + input=_with_output_names( + stalg.Rel(read=read_rel), named_struct.names + ), + names=named_struct.names, ) ) ], diff --git a/tests/builders/extended_expression/test_function_arguments.py b/tests/builders/extended_expression/test_function_arguments.py new file mode 100644 index 00000000..55a76a1a --- /dev/null +++ b/tests/builders/extended_expression/test_function_arguments.py @@ -0,0 +1,57 @@ +import substrait.algebra_pb2 as stalg +import substrait.type_pb2 as stt + +from substrait.builders.extended_expression import ( + EnumerationArgument, + column, + nested_struct, + scalar_function, +) +from substrait.extension_registry import ExtensionRegistry +from substrait.type_inference import infer_extended_expression_schema + + +named_struct = stt.NamedStruct( + names=["left", "right", "occurred_at"], + struct=stt.Type.Struct( + types=[ + stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)), + stt.Type(string=stt.Type.String(nullability=stt.Type.NULLABILITY_REQUIRED)), + stt.Type( + precision_timestamp=stt.Type.PrecisionTimestamp( + precision=6, + nullability=stt.Type.NULLABILITY_REQUIRED, + ) + ), + ], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), +) + + +def test_scalar_function_accepts_enumeration_arguments(): + registry = ExtensionRegistry() + expression = scalar_function( + "extension:io.substrait:functions_datetime", + "extract", + [EnumerationArgument("UNIX_TIME"), column("occurred_at")], + )(named_struct, registry) + + function = expression.referred_expr[0].expression.scalar_function + assert function.arguments[0] == stalg.FunctionArgument(enum="UNIX_TIME") + assert function.arguments[1].HasField("value") + assert function.output_type.WhichOneof("kind") == "i64" + + +def test_nested_struct_is_a_typed_scalar_expression(): + registry = ExtensionRegistry(load_default_extensions=False) + expression = nested_struct([column("left"), column("right")])( + named_struct, registry + ) + + inferred = infer_extended_expression_schema(expression, registry=registry) + assert inferred.types[0].WhichOneof("kind") == "struct" + assert [item.WhichOneof("kind") for item in inferred.types[0].struct.types] == [ + "i64", + "string", + ] diff --git a/tests/builders/plan/test_materialize.py b/tests/builders/plan/test_materialize.py new file mode 100644 index 00000000..01530283 --- /dev/null +++ b/tests/builders/plan/test_materialize.py @@ -0,0 +1,86 @@ +import substrait.algebra_pb2 as stalg +import substrait.type_pb2 as stt + +from substrait.builders.extended_expression import column, literal +from substrait.builders.plan import ( + filter, + materialize, + read_named_table, + select, + with_relation_alias, +) +from substrait.builders.type import boolean, i64 +from substrait.extension_registry import ExtensionRegistry + + +def test_materialize_preserves_inferred_names_on_nested_relations(): + registry = ExtensionRegistry(load_default_extensions=False) + schema = stt.NamedStruct( + names=["id", "is_applicable"], + struct=stt.Type.Struct( + types=[i64(nullable=False), boolean()], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), + ) + plan = select( + filter( + read_named_table("example", schema), + literal(True, boolean(nullable=False)), + ), + [column("id")], + ) + + bound = materialize(plan, registry, include_relation_output_names=True) + project = bound.relations[-1].root.input.project + filtered = project.input.filter + read = filtered.input.read + + assert list(project.common.hint.output_names) == ["id"] + assert list(filtered.common.hint.output_names) == ["id", "is_applicable"] + assert list(read.common.hint.output_names) == ["id", "is_applicable"] + + +def test_materialize_keeps_output_name_hints_opt_in(): + registry = ExtensionRegistry(load_default_extensions=False) + schema = stt.NamedStruct( + names=["id"], + struct=stt.Type.Struct( + types=[i64(nullable=False)], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), + ) + + bound = materialize(read_named_table("example", schema), registry) + + assert not bound.relations[-1].root.input.read.common.HasField("hint") + + +def test_relation_aliases_survive_native_relation_nesting(): + registry = ExtensionRegistry(load_default_extensions=False) + schema = stt.NamedStruct( + names=["order_id"], + struct=stt.Type.Struct( + types=[i64(nullable=False)], + nullability=stt.Type.NULLABILITY_REQUIRED, + ), + ) + base = with_relation_alias( + read_named_table("orders", schema), + "orders_read", + ) + projected = with_relation_alias( + select(base, [column("order_id")]), + "orders_project", + ) + + bound = materialize( + projected, + registry, + include_relation_output_names=True, + ) + + project = bound.relations[-1].root.input.project + assert project.common.hint.alias == "orders_project" + assert list(project.common.hint.output_names) == ["order_id"] + assert project.input.read.common.hint.alias == "orders_read" + assert list(project.input.read.common.hint.output_names) == ["order_id"]