From b24a5f1f4cf3783444ff71716fd89ecf57137736 Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Mon, 3 Aug 2026 16:03:22 +0200 Subject: [PATCH] feat!: assign extension anchors per plan, not per registry `ExtensionRegistry` handed out `function_anchor` / `extension_urn_anchor` values at registration time and builders stamped those registry-global numbers into plans. But anchors are plan-local in Substrait, which caused two problems: - Plans were not reproducible. A single-`add` plan emitted `function_anchor: 284` against the default extension set and `4` against a minimal one, because the value encoded how many functions the other YAMLs defined and the order the `functions*.yaml` glob returned them (filesystem order, not sorted). - Extending a plan built elsewhere silently corrupted it. The merge helpers dedupe by identity and document "assumes that there are no collisions", with nothing enforcing it, so a foreign plan already using a given anchor produced two URNs at one anchor and two functions at another -- leaving `function_reference` ambiguous, with no error. Introduce `ExtensionCollector`, which owns those anchors for the duration of one build: function references are allocated on first use from 1, and URN anchors are derived at emit time (nothing outside `SimpleExtensionDeclaration` refers to one). It follows substrait-java's `io.substrait.extension.ExtensionCollector`, including that numbering. The collector reaches builders through a contextvar, as the builders' other per-build state already does (`_rel_anchor_counter`, `outer_schemas`, `anchor_scope`). An incoming materialized plan has its declarations read back to `(urn, name)` identities and its references re-derived rather than trusted, so independently numbered inputs cannot disagree about what a reference means. This is what the SQL translator needs, as it builds a set operation's two sides as separate plans before merging them. Identities come off the declaration rather than a catalog lookup, so a plan naming functions absent from the registry still round-trips; declarations too under-specified to re-derive (pyarrow emits a bare `extension_function { name: "add" }`) are preserved verbatim. Because the collector accumulates once per build, the per-level extension merging in the builders is gone rather than optimized: an N-verb chain scanned 230 declarations across 80 merge calls at N=40, and now does none. This is the extension half of #207; the schema re-inference half is untouched. `ExtensionRegistry` is now a pure catalog. `lookup_urn` and `FunctionEntry.anchor` are deprecated (`has_urn` / `urns()` replace the former); the urn->function mapping, signature matching and extension-relation registration are unchanged. BREAKING CHANGE: emitted extension anchors are now numbered per plan, so plans compared byte-for-byte against output from an earlier release will differ. Anchors are plan-local by spec, so plan semantics are unaffected. `ExtensionRegistry.lookup_urn` and `FunctionEntry.anchor` now warn and no longer return registry-global anchors. Closes #236 --- src/substrait/builders/extended_expression.py | 218 ++------- src/substrait/builders/plan.py | 174 +++---- src/substrait/dataframe/expr.py | 7 - src/substrait/extension_registry/__init__.py | 12 + src/substrait/extension_registry/collector.py | 221 +++++++++ .../extension_registry/function_entry.py | 20 +- src/substrait/extension_registry/registry.py | 40 +- src/substrait/utils/__init__.py | 76 ++- .../test_scalar_function.py | 11 +- tests/extension_registry/test_collector.py | 436 ++++++++++++++++++ tests/sql/test_sql_anchors.py | 100 ++++ tests/test_utils.py | 109 +++++ 12 files changed, 1150 insertions(+), 274 deletions(-) create mode 100644 src/substrait/extension_registry/collector.py create mode 100644 tests/extension_registry/test_collector.py create mode 100644 tests/sql/test_sql_anchors.py diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index f692a746..3e971e5f 100644 --- a/src/substrait/builders/extended_expression.py +++ b/src/substrait/builders/extended_expression.py @@ -9,16 +9,19 @@ import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee -import substrait.extensions.extensions_pb2 as ste import substrait.type_pb2 as stp -from substrait.extension_registry import ExtensionRegistry +from substrait.extension_registry import ( + ExtensionRegistry, + build_scoped, + current_collector, + function_reference, +) from substrait.type_inference import infer_extended_expression_schema, outer_schemas from substrait.utils import ( inline_reference_rels, - merge_extension_declarations, - merge_extension_urns, plan_subtrees, + remap_function_references, type_num_names, ) @@ -85,11 +88,20 @@ def resolve_expression( base_schema: stp.NamedStruct, registry: ExtensionRegistry, ) -> stee.ExtendedExpression: - return ( - expression - if isinstance(expression, stee.ExtendedExpression) - else expression(base_schema, registry) - ) + """Resolve ``expression``, folding its extensions into the build in progress. + + An already-bound ExtendedExpression numbered its function references against + whichever build produced it, so the collector re-derives them from the durable + ``(urn, name)`` identities and the expression is rewritten to match -- the + expression-level counterpart of ``builders.plan._bind``. Unchanged when the + numbering already agrees, as it does for anything this build resolved. + """ + if not isinstance(expression, stee.ExtendedExpression): + return expression(base_schema, registry) + collector = current_collector() + if collector is None: + return expression + return remap_function_references(expression, collector.adopt(expression)) def alias( @@ -110,7 +122,7 @@ def resolve( bound_expression.referred_expr[0].output_names[0] = name return bound_expression - return resolve + return build_scoped(resolve) _EPOCH_DATE = date(1970, 1, 1) @@ -386,7 +398,7 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) def outer_reference(field: Union[str, int], steps_out: int = 1): @@ -430,7 +442,7 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) class LateralInput: @@ -476,7 +488,7 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) def column(field: Union[str, int], alias: Union[Iterable[str], str, None] = None): @@ -526,7 +538,7 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) def scalar_function( @@ -561,36 +573,14 @@ def resolve( if not func: raise Exception(f"Unknown function {function} for {signature}") - func_extension_urns = [ - ste.SimpleExtensionURN( - extension_urn_anchor=registry.lookup_urn(urn), urn=urn - ) - ] - - func_extensions = [ - ste.SimpleExtensionDeclaration( - extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( - extension_urn_reference=registry.lookup_urn(urn), - function_anchor=func[0].anchor, - name=str(func[0]), - ) - ) - ] - - extension_urns = merge_extension_urns( - func_extension_urns, *[b.extension_urns for b in bound_expressions] - ) - - extensions = merge_extension_declarations( - func_extensions, *[b.extensions for b in bound_expressions] - ) + func_ref = function_reference(urn, str(func[0])) return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( expression=stalg.Expression( scalar_function=stalg.Expression.ScalarFunction( - function_reference=func[0].anchor, + function_reference=func_ref, arguments=[ stalg.FunctionArgument( value=e.referred_expr[0].expression @@ -609,11 +599,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def aggregate_function( @@ -658,39 +646,13 @@ def resolve( if not func: raise Exception(f"Unknown function {function} for {signature}") - func_extension_urns = [ - ste.SimpleExtensionURN( - extension_urn_anchor=registry.lookup_urn(urn), urn=urn - ) - ] - - func_extensions = [ - ste.SimpleExtensionDeclaration( - extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( - extension_urn_reference=registry.lookup_urn(urn), - function_anchor=func[0].anchor, - name=str(func[0]), - ) - ) - ] - - extension_urns = merge_extension_urns( - func_extension_urns, - *[b.extension_urns for b in bound_expressions], - *[s.extension_urns for s, _ in bound_sorts], - ) - - extensions = merge_extension_declarations( - func_extensions, - *[b.extensions for b in bound_expressions], - *[s.extensions for s, _ in bound_sorts], - ) + func_ref = function_reference(urn, str(func[0])) return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( measure=stalg.AggregateFunction( - function_reference=func[0].anchor, + function_reference=func_ref, arguments=[ stalg.FunctionArgument(value=e.referred_expr[0].expression) for e in bound_expressions @@ -715,11 +677,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) # TODO bounds, sorts @@ -756,40 +716,14 @@ def resolve( if not func: raise Exception(f"Unknown function {function} for {signature}") - func_extension_urns = [ - ste.SimpleExtensionURN( - extension_urn_anchor=registry.lookup_urn(urn), urn=urn - ) - ] - - func_extensions = [ - ste.SimpleExtensionDeclaration( - extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( - extension_urn_reference=registry.lookup_urn(urn), - function_anchor=func[0].anchor, - name=str(func[0]), - ) - ) - ] - - extension_urns = merge_extension_urns( - func_extension_urns, - *[b.extension_urns for b in bound_expressions], - *[b.extension_urns for b in bound_partitions], - ) - - extensions = merge_extension_declarations( - func_extensions, - *[b.extensions for b in bound_expressions], - *[b.extensions for b in bound_partitions], - ) + func_ref = function_reference(urn, str(func[0])) return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( expression=stalg.Expression( window_function=stalg.Expression.WindowFunction( - function_reference=func[0].anchor, + function_reference=func_ref, arguments=[ stalg.FunctionArgument( value=e.referred_expr[0].expression @@ -811,11 +745,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def if_then( @@ -838,18 +770,6 @@ def resolve( bound_else = resolve_expression(_else, base_schema, registry) - extension_urns = merge_extension_urns( - *[b[0].extension_urns for b in bound_ifs], - *[b[1].extension_urns for b in bound_ifs], - bound_else.extension_urns, - ) - - extensions = merge_extension_declarations( - *[b[0].extensions for b in bound_ifs], - *[b[1].extensions for b in bound_ifs], - bound_else.extensions, - ) - return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( @@ -889,11 +809,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def switch( @@ -916,18 +834,6 @@ def resolve( ] bound_else = resolve_expression(_else, base_schema, registry) - extension_urns = merge_extension_urns( - bound_match.extension_urns, - *[b.extension_urns for _, b in bound_ifs], - bound_else.extension_urns, - ) - - extensions = merge_extension_declarations( - bound_match.extensions, - *[b.extensions for _, b in bound_ifs], - bound_else.extensions, - ) - return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( @@ -950,11 +856,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def singular_or_list( @@ -968,14 +872,6 @@ def resolve( bound_value = resolve_expression(value, base_schema, registry) bound_options = [resolve_expression(o, base_schema, registry) for o in options] - extension_urns = merge_extension_urns( - bound_value.extension_urns, *[b.extension_urns for b in bound_options] - ) - - extensions = merge_extension_declarations( - bound_value.extensions, *[b.extensions for b in bound_options] - ) - return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( @@ -993,11 +889,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def multi_or_list( @@ -1014,16 +908,6 @@ def resolve( [resolve_expression(e, base_schema, registry) for e in o] for o in options ] - extension_urns = merge_extension_urns( - *[b.extension_urns for b in bound_value], - *[e.extension_urns for b in bound_options for e in b], - ) - - extensions = merge_extension_declarations( - *[b.extensions for b in bound_value], - *[e.extensions for b in bound_options for e in b], - ) - return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( @@ -1044,11 +928,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def cast( @@ -1079,11 +961,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=bound_input.extension_urns, - extensions=bound_input.extensions, ) - return resolve + return build_scoped(resolve) # -- subqueries ----------------------------------------------------------- @@ -1100,12 +980,6 @@ def _subquery(subquery, base_schema, output_name, *extension_sources): ) ], base_schema=base_schema, - extension_urns=merge_extension_urns( - *[s.extension_urns for s in extension_sources] - ), - extensions=merge_extension_declarations( - *[s.extensions for s in extension_sources] - ), ) @@ -1141,7 +1015,7 @@ def resolve(base_schema, registry): ) return _subquery(subquery, base_schema, alias or "subquery", plan) - return resolve + return build_scoped(resolve) def set_predicate(query, op, alias: Union[str, None] = None): @@ -1156,7 +1030,7 @@ def resolve(base_schema, registry): ) return _subquery(subquery, base_schema, alias or "exists", plan) - return resolve + return build_scoped(resolve) def in_predicate(needles, query, alias: Union[str, None] = None): @@ -1172,7 +1046,7 @@ def resolve(base_schema, registry): ) return _subquery(subquery, base_schema, alias or "in_subquery", plan, *bound) - return resolve + return build_scoped(resolve) def set_comparison(left, query, reduction_op, comparison_op, alias=None): @@ -1193,7 +1067,7 @@ def resolve(base_schema, registry): subquery, base_schema, alias or "set_comparison", plan, bound_left ) - return resolve + return build_scoped(resolve) def execution_context_variable(variable: str, type_value, alias=None): @@ -1218,7 +1092,7 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) def dynamic_parameter(parameter_reference: int, type: stp.Type, alias=None): @@ -1245,4 +1119,4 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index dff7d700..3e3cdc87 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -20,7 +20,11 @@ next_rel_anchor, resolve_expression, ) -from substrait.extension_registry import ExtensionRegistry +from substrait.extension_registry import ( + ExtensionRegistry, + build_scoped, + current_collector, +) from substrait.type_inference import ( _join_output_struct, _join_struct_from_schemas, @@ -29,10 +33,9 @@ join_output_names, ) from substrait.utils import ( - merge_extension_declarations, - merge_extension_urns, plan_subtrees, rebase_reference_ordinals, + remap_function_references, ) from substrait.version import substrait_version @@ -55,21 +58,42 @@ def _create_default_version(): _create_default_version() +def _bind(plan: PlanOrUnbound, registry: ExtensionRegistry) -> stp.Plan: + """Resolve ``plan`` and fold its extensions into the build in progress. + + A plan built elsewhere -- or by an earlier, separate build -- numbered its + function references independently, so the collector re-derives them from the + durable ``(urn, name)`` identities and the plan's relations are rewritten to + match. Returns the plan untouched when the numbering already agrees, which is + always the case for one resolved by the current build (it allocated through the + same collector, and carries no declarations of its own until the outermost + resolver writes them). + + Every builder binds its inputs through here, so this is the single point at + which a foreign plan's anchor space is reconciled with ours. + """ + bound = plan if isinstance(plan, stp.Plan) else plan(registry) + collector = current_collector() + if collector is None: + return bound + return remap_function_references(bound, collector.adopt(bound)) + + def _merge_plan_metadata(*objs): """Collect the plan-level metadata a builder carries over from its inputs. - ``objs`` is a mix of input Plans and bound ExtendedExpressions. Extension - URNs and declarations are merged from all of them; the plan-level execution - behavior is carried over from the first input Plan that declares one + ``objs`` is a mix of input Plans and bound ExtendedExpressions. The plan-level + execution behavior is carried over from the first input Plan that declares one (expressions have no such field). Because every relational builder routes its inputs through here, an execution behavior set anywhere upstream is preserved on the freshly-constructed output Plan -- so it is order independent across a pipeline rather than only surviving as the last step. + + Extension URNs and declarations are *not* merged here: they belong to the + build's ``ExtensionCollector``, which writes them onto the outermost plan once + (see ``build_scoped``), rather than being re-merged at every level. """ - metadata = { - "extension_urns": merge_extension_urns(*[b.extension_urns for b in objs if b]), - "extensions": merge_extension_declarations(*[b.extensions for b in objs if b]), - } + metadata = {} for b in objs: if isinstance(b, stp.Plan) and b.HasField("execution_behavior"): metadata["execution_behavior"] = b.execution_behavior @@ -86,8 +110,7 @@ def _merge_input_subtrees(bound_inputs): Returns ``(subtree_planrels, rebased_root_inputs)``: the deduplicated combined subtrees as leading ``PlanRel(rel=...)`` entries, and, per input, its root's - input Rel with ReferenceRel ordinals rebased into the combined list. Mirrors how - ``_merge_plan_metadata`` carries extension declarations upward. + input Rel with ReferenceRel ordinals rebased into the combined list. Structurally-identical subtrees (byte-equal serialized ``Rel``) collapse to a single ordinal, so a cached frame reused across branches that later meet at a @@ -133,10 +156,11 @@ def _plan_from( Merges the shared subtrees carried by ``bound_inputs`` (deduping and rebasing ordinals), builds the output ``Rel`` by calling ``make_rel`` with the list of rebased input rels (one per bound input, in order), and prepends the combined - subtrees as leading ``rel`` entries ahead of the query root. Metadata (extension - declarations / execution behavior) is merged from ``metadata_sources`` (input - plans and bound expressions). This is the single place the CTE subtree - propagation and Plan assembly live, so every relational builder is one call. + subtrees as leading ``rel`` entries ahead of the query root. Plan-level metadata + is carried over from ``metadata_sources`` (input plans and bound expressions); + extension declarations are not, as the build's ``ExtensionCollector`` owns those. + This is the single place the CTE subtree propagation and Plan assembly live, so + every relational builder is one call. """ subtree_planrels, input_rels = _merge_input_subtrees(bound_inputs) root = stp.PlanRel( @@ -166,14 +190,14 @@ def with_execution_behavior( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) result = stp.Plan() result.CopyFrom(bound_plan) result.execution_behavior.variable_eval_mode = variable_eval_mode return result - return resolve + return build_scoped(resolve) def read_named_table( @@ -205,7 +229,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ], ) - return resolve + return build_scoped(resolve) def _require_schema(named_struct: stt.NamedStruct) -> stt.NamedStruct: @@ -262,7 +286,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ) return _read_plan(named_struct, read_rel) - return resolve + return build_scoped(resolve) def local_files( @@ -282,7 +306,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ) return _read_plan(named_struct, read_rel) - return resolve + return build_scoped(resolve) def extension_table( @@ -302,7 +326,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ) return _read_plan(named_struct, read_rel) - return resolve + return build_scoped(resolve) def project( @@ -325,7 +349,7 @@ def project( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - _plan = plan if isinstance(plan, stp.Plan) else plan(registry) + _plan = _bind(plan, registry) ns = infer_plan_schema(_plan, registry=registry) bound_expressions: Iterable[stee.ExtendedExpression] = [ resolve_expression(e, ns, registry) for e in expressions @@ -352,7 +376,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (_plan, *bound_expressions), ) - return resolve + return build_scoped(resolve) def select( @@ -375,7 +399,7 @@ def select( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - _plan = plan if isinstance(plan, stp.Plan) else plan(registry) + _plan = _bind(plan, registry) ns = infer_plan_schema(_plan, registry=registry) bound_expressions: Iterable[stee.ExtendedExpression] = [ resolve_expression(e, ns, registry) for e in expressions @@ -409,7 +433,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (_plan, *bound_expressions), ) - return resolve + return build_scoped(resolve) def filter( @@ -418,7 +442,7 @@ def filter( extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) ns = infer_plan_schema(bound_plan, registry=registry) bound_expression: stee.ExtendedExpression = resolve_expression( expression, ns, registry @@ -437,7 +461,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan, bound_expression), ) - return resolve + return build_scoped(resolve) def sort( @@ -451,7 +475,7 @@ def sort( extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) ns = infer_plan_schema(bound_plan, registry=registry) bound_expressions = [ @@ -483,12 +507,12 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan, *[e[0] for e in bound_expressions]), ) - return resolve + return build_scoped(resolve) def set(inputs: Iterable[PlanOrUnbound], op: stalg.SetRel.SetOp) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_inputs = [i if isinstance(i, stp.Plan) else i(registry) for i in inputs] + bound_inputs = [_bind(i, registry) for i in inputs] return _plan_from( bound_inputs, lambda inp: stalg.Rel(set=stalg.SetRel(inputs=inp, op=op)), @@ -496,7 +520,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(bound_inputs), ) - return resolve + return build_scoped(resolve) def reference(plan: PlanOrUnbound) -> UnboundPlan: @@ -514,7 +538,7 @@ def reference(plan: PlanOrUnbound) -> UnboundPlan: """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound = plan if isinstance(plan, stp.Plan) else plan(registry) + bound = _bind(plan, registry) nested = [stp.PlanRel(rel=s) for s in plan_subtrees(bound)] ordinal = len(nested) # the promoted root sits after the plan's own subtrees promoted = stp.PlanRel(rel=bound.relations[-1].root.input) @@ -530,7 +554,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: **_merge_plan_metadata(bound), ) - return resolve + return build_scoped(resolve) def fetch( @@ -540,7 +564,7 @@ def fetch( extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) ns = infer_plan_schema(bound_plan, registry=registry) bound_offset = resolve_expression(offset, ns, registry) if offset else None @@ -567,7 +591,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan, bound_offset, bound_count), ) - return resolve + return build_scoped(resolve) def join( @@ -580,8 +604,8 @@ def join( extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_left = left if isinstance(left, stp.Plan) else left(registry) - bound_right = right if isinstance(right, stp.Plan) else right(registry) + bound_left = _bind(left, registry) + bound_right = _bind(right, registry) left_ns = infer_plan_schema(bound_left, registry=registry) right_ns = infer_plan_schema(bound_right, registry=registry) @@ -637,7 +661,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_left, bound_right, bound_expression, bound_post), ) - return resolve + return build_scoped(resolve) def lateral_join( @@ -664,7 +688,7 @@ def lateral_join( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_left = left if isinstance(left, stp.Plan) else left(registry) + bound_left = _bind(left, registry) left_ns = infer_plan_schema(bound_left, registry=registry) anchor = next_rel_anchor() @@ -674,11 +698,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: # the current left row during that inference. with _outer_anchor_binding(anchor, left_ns.struct): unbound_right = right(handle) - bound_right = ( - unbound_right - if isinstance(unbound_right, stp.Plan) - else unbound_right(registry) - ) + bound_right = _bind(unbound_right, registry) right_ns = infer_plan_schema(bound_right, registry=registry) # The join condition binds against the combined left+right input row. @@ -738,7 +758,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_left, bound_right, bound_expression, bound_post), ) - return resolve + return build_scoped(resolve) def cross( @@ -747,8 +767,8 @@ def cross( extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_left = left if isinstance(left, stp.Plan) else left(registry) - bound_right = right if isinstance(right, stp.Plan) else right(registry) + bound_left = _bind(left, registry) + bound_right = _bind(right, registry) left_ns = infer_plan_schema(bound_left, registry=registry) right_ns = infer_plan_schema(bound_right, registry=registry) @@ -773,7 +793,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_left, bound_right), ) - return resolve + return build_scoped(resolve) def aggregate( @@ -795,7 +815,7 @@ def aggregate( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_input = input if isinstance(input, stp.Plan) else input(registry) + bound_input = _bind(input, registry) ns = infer_plan_schema(bound_input, registry=registry) bound_grouping_expressions = [ @@ -854,7 +874,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ), ) - return resolve + return build_scoped(resolve) def write_named_table( @@ -865,7 +885,7 @@ def write_named_table( output_mode: Union[stalg.WriteRel.OutputMode.ValueType, None] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_input = input if isinstance(input, stp.Plan) else input(registry) + bound_input = _bind(input, registry) ns = infer_plan_schema(bound_input, registry=registry) _table_names = [table_names] if isinstance(table_names, str) else table_names _create_mode = create_mode or stalg.WriteRel.CREATE_MODE_ERROR_IF_EXISTS @@ -890,7 +910,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: include_version=False, ) - return resolve + return build_scoped(resolve) def ddl( @@ -913,11 +933,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_inputs = [] schema = table_schema if view_definition is not None: - view_plan = ( - view_definition - if isinstance(view_definition, stp.Plan) - else view_definition(registry) - ) + view_plan = _bind(view_definition, registry) bound_inputs = [view_plan] merge_sources.append(view_plan) if schema is None: @@ -940,7 +956,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(merge_sources), ) - return resolve + return build_scoped(resolve) def update( @@ -995,7 +1011,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: **_merge_plan_metadata(*merge_sources), ) - return resolve + return build_scoped(resolve) def consistent_partition_window( @@ -1011,7 +1027,7 @@ def consistent_partition_window( extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) ns = infer_plan_schema(bound_plan, registry=registry) bound_partitions = [ @@ -1088,7 +1104,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ), ) - return resolve + return build_scoped(resolve) def expand( @@ -1105,7 +1121,7 @@ def expand( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_input = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_input = _bind(plan, registry) ns = infer_plan_schema(bound_input, registry=registry) expand_fields = [] @@ -1142,7 +1158,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(merge_sources), ) - return resolve + return build_scoped(resolve) def nested_loop_join( @@ -1155,8 +1171,8 @@ def nested_loop_join( """A NestedLoopJoinRel: join over the Cartesian product using ``expression``.""" def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_left = left if isinstance(left, stp.Plan) else left(registry) - bound_right = right if isinstance(right, stp.Plan) else right(registry) + bound_left = _bind(left, registry) + bound_right = _bind(right, registry) left_ns = infer_plan_schema(bound_left, registry=registry) right_ns = infer_plan_schema(bound_right, registry=registry) @@ -1189,7 +1205,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_left, bound_right, bound_expression), ) - return resolve + return build_scoped(resolve) def _comparison_join_keys(left_keys, right_keys, left_ns, right_ns, registry): @@ -1235,8 +1251,8 @@ def builder( extension: Optional[AdvancedExtension] = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_left = left if isinstance(left, stp.Plan) else left(registry) - bound_right = right if isinstance(right, stp.Plan) else right(registry) + bound_left = _bind(left, registry) + bound_right = _bind(right, registry) left_ns = infer_plan_schema(bound_left, registry=registry) right_ns = infer_plan_schema(bound_right, registry=registry) keys = _comparison_join_keys( @@ -1303,7 +1319,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_left, bound_right, bound_post, bound_residual), ) - return resolve + return build_scoped(resolve) return builder @@ -1338,7 +1354,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: relations=[stp.PlanRel(root=stalg.RelRoot(input=rel, names=out_names))], ) - return resolve + return build_scoped(resolve) def extension_single(plan: PlanOrUnbound, detail) -> UnboundPlan: @@ -1350,7 +1366,7 @@ def extension_single(plan: PlanOrUnbound, detail) -> UnboundPlan: """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) if hasattr(detail, "derive_schema"): input_struct = infer_plan_schema(bound_plan, registry=registry).struct names = list(detail.derive_schema(input_struct).names) @@ -1367,14 +1383,14 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan,), ) - return resolve + return build_scoped(resolve) def extension_multi(inputs: Iterable[PlanOrUnbound], detail) -> UnboundPlan: """An ExtensionMultiRel over ``inputs`` from an ExtensionMultiDetail.""" def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_inputs = [i if isinstance(i, stp.Plan) else i(registry) for i in inputs] + bound_inputs = [_bind(i, registry) for i in inputs] input_structs = [ infer_plan_schema(b, registry=registry).struct for b in bound_inputs ] @@ -1391,7 +1407,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(bound_inputs), ) - return resolve + return build_scoped(resolve) def exchange( @@ -1406,7 +1422,7 @@ def exchange( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) kind = ( {"broadcast": stalg.ExchangeRel.Broadcast()} if broadcast @@ -1425,7 +1441,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan,), ) - return resolve + return build_scoped(resolve) def top_n( @@ -1445,7 +1461,7 @@ def top_n( """ def resolve(registry: ExtensionRegistry) -> stp.Plan: - bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry) + bound_plan = _bind(plan, registry) ns = infer_plan_schema(bound_plan, registry=registry) bound_sorts = [ (resolve_expression(e, ns, registry), direction) for e, direction in sorts @@ -1480,4 +1496,4 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan, *[s for s, _ in bound_sorts], bound_count, bound_offset), ) - return resolve + return build_scoped(resolve) diff --git a/src/substrait/dataframe/expr.py b/src/substrait/dataframe/expr.py index 54a6b527..c185a374 100644 --- a/src/substrait/dataframe/expr.py +++ b/src/substrait/dataframe/expr.py @@ -65,7 +65,6 @@ set_predicate as _set_predicate, ) from substrait.type_inference import infer_extended_expression_schema -from substrait.utils import merge_extensions_into # Standard Substrait function-extension URNs used by the operators below. FUNCTIONS_COMPARISON = "extension:io.substrait:functions_comparison" @@ -687,8 +686,6 @@ def resolve(base_schema, registry): ) ], base_schema=base_schema, - extension_urns=body.extension_urns, - extensions=body.extensions, ) return scalar_function( FUNCTIONS_LIST, function, expressions=[bound_list, lambda_ee] @@ -765,8 +762,6 @@ def resolve(base_schema, registry): expr=bound_key.referred_expr[0].expression, direction=direction ) ) - # Carry over any extensions a (function-valued) sort key introduced. - merge_extensions_into(bound, bound_key) return bound return Expr(resolve) @@ -810,7 +805,6 @@ def resolve(base_schema, registry): key = p.unbound if isinstance(p, Expr) else column(p) bound_p = resolve_expression(key, base_schema, registry) wf.partitions.append(bound_p.referred_expr[0].expression) - merge_extensions_into(bound, bound_p) for k in order_keys: key = k.unbound if isinstance(k, Expr) else column(k) bound_k = resolve_expression(key, base_schema, registry) @@ -819,7 +813,6 @@ def resolve(base_schema, registry): expr=bound_k.referred_expr[0].expression, direction=direction ) ) - merge_extensions_into(bound, bound_k) frame = rows if rows is not None else range if frame is not None: wf.bounds_type = ( diff --git a/src/substrait/extension_registry/__init__.py b/src/substrait/extension_registry/__init__.py index f278ccbb..25c416e7 100644 --- a/src/substrait/extension_registry/__init__.py +++ b/src/substrait/extension_registry/__init__.py @@ -1,5 +1,12 @@ """Extension Registry module.""" +from .collector import ( + ExtensionCollector, + build_scope, + build_scoped, + current_collector, + function_reference, +) from .exceptions import UnhandledParameterizedTypeError, UnrecognizedSubstraitTypeError from .function_entry import FunctionEntry, FunctionType from .registry import ExtensionRegistry @@ -12,8 +19,13 @@ ) __all__ = [ + "ExtensionCollector", "ExtensionRegistry", "FunctionEntry", + "build_scope", + "build_scoped", + "current_collector", + "function_reference", "FunctionType", "normalize_substrait_type_names", "_check_integer_constraint", diff --git a/src/substrait/extension_registry/collector.py b/src/substrait/extension_registry/collector.py new file mode 100644 index 00000000..312792a1 --- /dev/null +++ b/src/substrait/extension_registry/collector.py @@ -0,0 +1,221 @@ +"""Per-build collection of plan-local extension anchors. + +Extension anchors (``function_anchor``, ``extension_urn_anchor``) are *plan-local* +in Substrait: they are an artifact of serializing a plan, not durable identifiers. +The durable identity of a function is its ``(urn, name)`` pair. This module owns +the mapping between the two for the duration of a single build, which is what +keeps :class:`~substrait.extension_registry.ExtensionRegistry` free to be a pure +catalog. + +Mirrors ``ExtensionCollector`` in substrait-java ( +``io.substrait.extension.ExtensionCollector``), including its first-use reference +numbering and its deferral of URN anchors to emit time. +""" + +import contextlib +import contextvars +import functools +import itertools +from typing import Optional, Union + +import substrait.extended_expression_pb2 as stee +import substrait.extensions.extensions_pb2 as ste +import substrait.plan_pb2 as stplan + +# Identity of a function as declared in a plan: (extension URN, function name). +# The name is the compound form carried by SimpleExtensionDeclaration (e.g. +# "add:i64_i64"), which is what makes an identity resolvable without the catalog. +# The URN is None for a declaration that names no resolvable extension URN. +FunctionIdentity = tuple[Optional[str], str] + +ExtensionCarrier = Union[stplan.Plan, stee.ExtendedExpression] + + +def _is_underspecified(function, urns_by_anchor: dict) -> bool: + """Whether a declaration carries too little to be re-derived, so must be kept as is. + + Some producers emit a bare ``extension_function { name: "add" }`` -- no anchor + and no URN (pyarrow's ``serialize_expressions`` does exactly this). There is + nothing to renumber: anchor 0 is not a reference any expression can name under + the spec's 1-based numbering, so re-assigning it would desynchronize the + declaration from expressions that still say 0. Such declarations are preserved + verbatim instead. + """ + return ( + function.function_anchor == 0 + and function.extension_urn_reference not in urns_by_anchor + ) + + +class ExtensionCollector: + """Owns the plan-local extension anchors for a single build. + + Function references are allocated on first use, numbered from 1 in the order + they are first referenced. URN anchors are *not* allocated during the build: + nothing outside ``SimpleExtensionDeclaration`` refers to one, so they are + derived in :meth:`write_into` from the order the declarations were collected. + """ + + def __init__(self) -> None: + self._references: dict[FunctionIdentity, int] = {} + self._reference_generator = itertools.count(1) + # Declarations kept verbatim because they carry neither a URN nor an anchor + # to re-derive from; see _is_underspecified. Deduplicated on serialized bytes. + self._passthrough: "list[ste.SimpleExtensionDeclaration]" = [] + self._passthrough_keys: set = set() + + def function_reference(self, urn: str, name: str) -> int: + """The reference for ``(urn, name)``, allocating one on first use.""" + identity = (urn, name) + reference = self._references.get(identity) + if reference is None: + reference = next(self._reference_generator) + self._references[identity] = reference + return reference + + def adopt(self, carrier: ExtensionCarrier) -> dict[int, int]: + """Take over the extension declarations of an incoming plan or expression. + + Reads ``carrier``'s declarations back to ``(urn, name)`` identities and + allocates this build's reference for each, returning the + ``{old reference: new reference}`` remap its relations/expressions need. + Pass the result to + :func:`substrait.utils.remap_function_references`; it is empty (identity) + whenever the incoming numbering already agrees with ours, which is the + common case. + + Every incoming reference is re-derived rather than trusted, so two + independently built inputs meeting at a multi-input relation cannot + disagree about what a reference number means. Identities come from the + declaration itself, so a function absent from the catalog is carried + through unchanged rather than rejected. + """ + urns_by_anchor = {u.extension_urn_anchor: u.urn for u in carrier.extension_urns} + remap = {} + for declaration in carrier.extensions: + if declaration.WhichOneof("mapping_type") != "extension_function": + # Type / type-variation declarations are not collected yet; see + # merge_extension_declarations for the same gap. + mapping_type = declaration.WhichOneof("mapping_type") + raise NotImplementedError( + f"cannot collect extension declaration of type {mapping_type!r}; " + f"only 'extension_function' declarations are supported so far" + ) + function = declaration.extension_function + if _is_underspecified(function, urns_by_anchor): + key = declaration.SerializeToString(deterministic=True) + if key not in self._passthrough_keys: + self._passthrough_keys.add(key) + self._passthrough.append(declaration) + continue + # An unresolvable URN reference still yields a usable identity: the + # function name alone. Anchors stay collision-free either way. + urn = urns_by_anchor.get(function.extension_urn_reference) + new = self.function_reference(urn, function.name) + if new != function.function_anchor: + remap[function.function_anchor] = new + return remap + + def write_into(self, carrier: ExtensionCarrier) -> None: + """Emit the collected extensions onto ``carrier``, replacing what is there. + + URN anchors are assigned here, numbered from 1 in the order each URN was + first referenced, so a plan's URN anchors are as dense and plan-local as + its function references. + """ + urn_anchors: dict[str, int] = {} + urn_anchor_generator = itertools.count(1) + # Pass-through declarations first: they keep whatever anchor they arrived + # with, and cannot clash with the 1-based numbering below. + declarations = list(self._passthrough) + for (urn, name), reference in self._references.items(): + function = ste.SimpleExtensionDeclaration.ExtensionFunction( + function_anchor=reference, name=name + ) + if urn is not None: + urn_anchor = urn_anchors.get(urn) + if urn_anchor is None: + urn_anchor = next(urn_anchor_generator) + urn_anchors[urn] = urn_anchor + function.extension_urn_reference = urn_anchor + declarations.append( + ste.SimpleExtensionDeclaration(extension_function=function) + ) + + carrier.ClearField("extension_urns") + carrier.extension_urns.extend( + ste.SimpleExtensionURN(extension_urn_anchor=anchor, urn=urn) + for urn, anchor in urn_anchors.items() + ) + carrier.ClearField("extensions") + carrier.extensions.extend(declarations) + + +# The ExtensionCollector for the build currently in progress, or None outside a +# build. Ambient rather than threaded through resolver signatures, following the +# other per-build state the builders already carry this way (_rel_anchor_counter +# in builders.extended_expression, outer_schemas / anchor_scope in type_inference). +_collector: contextvars.ContextVar = contextvars.ContextVar("_collector", default=None) + + +def current_collector() -> Optional[ExtensionCollector]: + """The collector for the build in progress, or None outside a build.""" + return _collector.get() + + +def function_reference(urn: str, name: str) -> int: + """This build's reference for the function ``(urn, name)``. + + Convenience over ``current_collector().function_reference(...)`` for builders, + which always run inside a scope. + """ + collector = _collector.get() + if collector is None: + raise RuntimeError( + "no build in progress: extension anchors are plan-local, so a builder " + "must resolve inside a build_scope() (builders are wrapped in " + "build_scoped(), which enters one)" + ) + return collector.function_reference(urn, name) + + +@contextlib.contextmanager +def build_scope(): + """Enter the current build, creating a collector if this is the outermost one. + + Yields ``(collector, owns_scope)``. ``owns_scope`` is True only for the + outermost resolver of a build -- the one responsible for writing the collected + extensions onto its output. Nested resolvers get the same collector and write + nothing, which is what lets a build accumulate extensions once instead of + re-merging them at every level. + """ + collector = _collector.get() + if collector is not None: + yield collector, False + return + collector = ExtensionCollector() + token = _collector.set(collector) + try: + yield collector, True + finally: + _collector.reset(token) + + +def build_scoped(resolve): + """Wrap a builder's resolver so it participates in a build scope. + + The outermost resolver of a build writes the collected extension declarations + onto whatever it returns; nested ones leave those fields empty for it to fill. + Signature-agnostic, since plan resolvers take ``(registry)`` and expression + resolvers take ``(base_schema, registry)``. + """ + + @functools.wraps(resolve) + def wrapper(*args, **kwargs): + with build_scope() as (collector, owns_scope): + resolved = resolve(*args, **kwargs) + if owns_scope: + collector.write_into(resolved) + return resolved + + return wrapper diff --git a/src/substrait/extension_registry/function_entry.py b/src/substrait/extension_registry/function_entry.py index acb6fffa..d623a95c 100644 --- a/src/substrait/extension_registry/function_entry.py +++ b/src/substrait/extension_registry/function_entry.py @@ -1,5 +1,6 @@ """Function entry class for extension registry.""" +import warnings from enum import Enum from typing import Optional, Union @@ -23,14 +24,12 @@ def __init__( urn: str, name: str, impl: Union[se.Impl, se.Impl1, se.Impl2], - anchor: int, function_type: FunctionType = FunctionType.SCALAR, ) -> None: self.name = name self.impl = impl self.normalized_inputs: list = [] self.urn: str = urn - self.anchor = anchor self.function_type = function_type self.arguments = [] self.nullability = ( @@ -50,6 +49,23 @@ def __init__( def __repr__(self) -> str: return f"{self.name}:{'_'.join(self.normalized_inputs)}" + @property + def anchor(self) -> int: + """Deprecated. Function anchors are plan-local, not catalog state. + + A function's durable identity is ``(urn, str(entry))``; the numeric anchor + it gets in a plan is assigned per build by + :class:`~substrait.extension_registry.ExtensionCollector`. + """ + warnings.warn( + "FunctionEntry.anchor is deprecated: function anchors are plan-local and " + "are assigned per build by ExtensionCollector. Identify a function by " + "(entry.urn, str(entry)) instead.", + DeprecationWarning, + stacklevel=2, + ) + return 1 + def satisfies_signature(self, signature: tuple | list) -> Optional[str]: if self.impl.variadic: min_args_allowed = self.impl.variadic.min or 0 diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 41c3ff0a..c454d6a6 100644 --- a/src/substrait/extension_registry/registry.py +++ b/src/substrait/extension_registry/registry.py @@ -1,7 +1,7 @@ """Extension Registry class.""" -import itertools import re +import warnings from collections import defaultdict from importlib.resources import files as importlib_files from pathlib import Path @@ -20,11 +20,17 @@ class ExtensionRegistry: + """A catalog of extension functions, keyed by URN. + + Plan-independent: the registry knows which functions exist and what signatures + they accept, but assigns no anchors. Extension anchors are plan-local in + Substrait, so they belong to a single build rather than to the catalog -- see + :class:`~substrait.extension_registry.ExtensionCollector`. + """ + def __init__(self, load_default_extensions=True) -> None: - self._urn_mapping: dict = defaultdict(dict) # URN -> anchor ID - self._urn_id_generator = itertools.count(1) + self._urns: set = set() self._function_mapping: dict = defaultdict(lambda: defaultdict(list)) - self._id_generator = itertools.count(1) # {type_url: detail class} for user-defined extension relations, so an # extension relation's output schema can be derived during inference. self._extension_relations: dict = {} @@ -72,7 +78,7 @@ def register_extension_dict(self, definitions: dict) -> None: if not unverified_urn: raise ValueError("Extension definitions must contain a 'urn' field") urn = validate_urn_format(unverified_urn) - self._urn_mapping[urn] = next(self._urn_id_generator) + self._urns.add(urn) simple_extensions = build_simple_extensions(definitions) # Helper to register functions by type @@ -89,7 +95,6 @@ def register_functions_by_type( urn=urn, name=function.name, impl=impl, - anchor=next(self._id_generator), function_type=func_type, ) for impl in function.impls @@ -175,8 +180,29 @@ def find_function( matches = self._find_matching_functions(function_name, signature, urns) return matches[0] if matches else None + def has_urn(self, urn: str) -> bool: + """Whether ``urn`` has been registered.""" + return urn in self._urns + + def urns(self) -> "list[str]": + """The registered extension URNs, in registration order-independent form.""" + return sorted(self._urns) + def lookup_urn(self, urn: str) -> Optional[int]: - return self._urn_mapping.get(urn, None) + """Deprecated. Registry-global URN anchors no longer exist. + + URN anchors are plan-local, so they are assigned per build by + :class:`~substrait.extension_registry.ExtensionCollector` rather than by the + catalog. Use :meth:`has_urn` to test membership. + """ + warnings.warn( + "ExtensionRegistry.lookup_urn is deprecated: extension URN anchors are " + "plan-local and are assigned per build by ExtensionCollector. Use " + "has_urn() to test whether a URN is registered.", + DeprecationWarning, + stacklevel=2, + ) + return 1 if urn in self._urns else None def iter_functions(self): """Yield ``(urn, name, function_type)`` for every registered function. diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index ebffd4d8..82c5ffb6 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -8,6 +8,7 @@ import substrait.extensions.extensions_pb2 as ste import substrait.plan_pb2 as stplan import substrait.type_pb2 as stp +from google.protobuf.message import Message def type_num_names(typ: stp.Type): @@ -27,6 +28,12 @@ def merge_extension_urns(*extension_urns: Iterable[ste.SimpleExtensionURN]): """Merges multiple sets of SimpleExtensionURN objects into a single set. The order of extensions is kept intact, while duplicates are discarded. Assumes that there are no collisions (different extensions having identical anchors). + + Note that anchor collisions between independently numbered inputs are real, so + that assumption does not hold in general. The builders no longer rely on it: + they route inputs through ``ExtensionCollector.adopt``, which re-derives anchors + from ``(urn, name)`` identities instead of merging pre-numbered sets. Retained + for external callers doing their own merging. """ seen_urns = set() ret = [] @@ -46,6 +53,9 @@ def merge_extension_declarations( """Merges multiple sets of SimpleExtensionDeclaration objects into a single set. The order of extension declarations is kept intact, while duplicates are discarded. Assumes that there are no collisions (different extension declarations having identical anchors). + + See :func:`merge_extension_urns` on why the builders no longer depend on that + assumption; this is retained for external callers. """ seen_extension_functions = set() @@ -176,6 +186,64 @@ def _inline_reference_rels_in_place(rel: stalg.Rel, subtrees) -> None: _inline_reference_rels_in_place(child, subtrees) +# Every field that holds a function reference, i.e. an index into a plan's +# extension declarations. Matched by name during a descriptor walk rather than +# enumerated per message type, so a reference field added to the protos upstream is +# picked up automatically instead of being silently skipped. Note that this library +# only ever emits the first of these; the others can appear in a plan built +# elsewhere. +_FUNCTION_REFERENCE_FIELDS = frozenset( + { + "function_reference", # ScalarFunction / WindowFunction / AggregateFunction + "comparison_function_reference", # SortField + "custom_function_reference", # ComparisonJoinKey.ComparisonType + } +) + + +def remap_function_references(msg, remap: dict): + """A copy of ``msg`` with every function reference remapped (old -> new). + + Used when a plan built elsewhere is folded into the build in progress: the + incoming plan numbered its functions independently, so + :meth:`~substrait.extension_registry.ExtensionCollector.adopt` re-derives the + numbering and this applies the result to the relations and expressions that + refer to it. Joins ``rebase_reference_ordinals`` and + ``to_id_based_outer_references`` as a whole-tree rewrite. + + ``msg`` may be any message (a ``Rel``, ``Plan``, or ``Expression``); it is + returned unchanged when ``remap`` is empty, which is the common case, so callers + need not special-case the no-op. + """ + if not remap: + return msg + out = type(msg)() + out.CopyFrom(msg) + _remap_function_references_in_place(out, remap) + return out + + +def _remap_function_references_in_place(msg, remap: dict) -> None: + # Cardinality is read off the value rather than the descriptor: `label` is + # deprecated in protobuf 6 while `is_repeated` is absent from older 5.x, and + # this package supports both. + # + # ListFields() snapshots the set fields, so assigning during iteration is safe, + # and unset scalars (a reference of 0, which the spec never assigns) are skipped. + for field, value in msg.ListFields(): + if isinstance(value, Message): + _remap_function_references_in_place(value, remap) + elif field.message_type is not None: + # A repeated message field, or a map whose values are messages + # (ScalarMap/MessageMap expose .values(), repeated fields do not). + items = value.values() if hasattr(value, "values") else value + for item in items: + if isinstance(item, Message): + _remap_function_references_in_place(item, remap) + elif field.name in _FUNCTION_REFERENCE_FIELDS and isinstance(value, int): + setattr(msg, field.name, remap.get(value, value)) + + def _iter_direct_subexpressions(msg): """Yield the immediate ``Expression`` messages owned by ``msg``. @@ -507,9 +575,11 @@ def merge_extensions_into(target, *sources): Appends any extension URNs / declarations carried by ``sources`` whose identity is not already present on ``target``, deduplicating with the same keys as :func:`merge_extension_urns` / :func:`merge_extension_declarations` (URN string, - resp. ``(extension URN reference, name)``). This is the identity used by - ``builders.plan._merge_extensions``, so the DataFrame/Expr layer and the plan - builders agree on when extensions collapse. + resp. ``(extension URN reference, name)``). + + No longer used by the builders or the DataFrame layer, which let the build's + ``ExtensionCollector`` accumulate declarations once instead; retained for + external callers assembling plans by hand. ``target`` and each ``source`` are messages carrying repeated ``extension_urns`` and ``extensions`` fields (a ``Plan`` or an ``ExtendedExpression``). Unlike the diff --git a/tests/builders/extended_expression/test_scalar_function.py b/tests/builders/extended_expression/test_scalar_function.py index db7c1aa7..bf4934cb 100644 --- a/tests/builders/extended_expression/test_scalar_function.py +++ b/tests/builders/extended_expression/test_scalar_function.py @@ -143,19 +143,22 @@ def test_nested_scalar_calls(): extension_urns=[ ste.SimpleExtensionURN(extension_urn_anchor=1, urn="extension:test:urn") ], + # Declarations are emitted in anchor order, which is the order the functions + # were first referenced: the inner test_func resolves before the outer + # is_positive that wraps it. extensions=[ ste.SimpleExtensionDeclaration( extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( extension_urn_reference=1, - function_anchor=2, - name="is_positive:i8", + function_anchor=1, + name="test_func:i8", ) ), ste.SimpleExtensionDeclaration( extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( extension_urn_reference=1, - function_anchor=1, - name="test_func:i8", + function_anchor=2, + name="is_positive:i8", ) ), ], diff --git a/tests/extension_registry/test_collector.py b/tests/extension_registry/test_collector.py new file mode 100644 index 00000000..ec419084 --- /dev/null +++ b/tests/extension_registry/test_collector.py @@ -0,0 +1,436 @@ +"""Tests for plan-local extension anchor assignment. + +Extension anchors are plan-local in Substrait, so they are owned by a per-build +``ExtensionCollector`` rather than by the ``ExtensionRegistry`` catalog. These tests +pin the two properties that ownership buys: anchors that depend only on the plan +(not on which extensions happen to be registered), and correct folding of a plan +built elsewhere into a new build. +""" + +import importlib.resources as importlib_resources + +import pytest +import substrait.algebra_pb2 as stalg +import substrait.extended_expression_pb2 as stee +import substrait.extensions.extensions_pb2 as ste +import substrait.plan_pb2 as stplan +import substrait.type_pb2 as stt + +from substrait.builders.extended_expression import column, scalar_function +from substrait.builders.plan import project, read_named_table +from substrait.extension_registry import ( + ExtensionCollector, + ExtensionRegistry, + build_scope, + function_reference, +) + +ARITHMETIC = "extension:io.substrait:functions_arithmetic" +COMPARISON = "extension:io.substrait:functions_comparison" + +I64 = stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)) +NAMED_STRUCT = stt.NamedStruct( + names=["a", "b"], + struct=stt.Type.Struct(types=[I64, I64], nullability=stt.Type.NULLABILITY_REQUIRED), +) + + +@pytest.fixture(scope="module") +def full_registry(): + return ExtensionRegistry(load_default_extensions=True) + + +@pytest.fixture(scope="module") +def arithmetic_only_registry(): + """A registry holding *only* functions_arithmetic. + + Anchors must not differ between this and the full default set: that + dependence is exactly what made plans non-reproducible. + """ + registry = ExtensionRegistry(load_default_extensions=False) + registry.register_extension_yaml( + next( + iter( + importlib_resources.files("substrait_extensions.extensions").glob( + "functions_arithmetic.yaml" + ) + ) + ) + ) + return registry + + +def _add_plan(registry): + """``SELECT a + b FROM t`` -- one function, so one declaration.""" + plan = read_named_table("t", NAMED_STRUCT) + return project( + plan, + expressions=[scalar_function(ARITHMETIC, "add", [column("a"), column("b")])], + )(registry) + + +def _declarations(plan): + """``{function anchor: name}`` for a plan's extension declarations.""" + return { + declaration.extension_function.function_anchor: declaration.extension_function.name + for declaration in plan.extensions + } + + +class TestCollector: + def test_allocates_from_one_on_first_use(self): + with build_scope() as (collector, owns_scope): + assert owns_scope + assert collector.function_reference(ARITHMETIC, "add:i64_i64") == 1 + assert collector.function_reference(COMPARISON, "gt:any_any") == 2 + + def test_same_identity_reuses_its_reference(self): + collector = ExtensionCollector() + first = collector.function_reference(ARITHMETIC, "add:i64_i64") + collector.function_reference(COMPARISON, "gt:any_any") + assert collector.function_reference(ARITHMETIC, "add:i64_i64") == first + + def test_urn_anchors_assigned_at_emit(self): + collector = ExtensionCollector() + collector.function_reference(ARITHMETIC, "add:i64_i64") + collector.function_reference(COMPARISON, "gt:any_any") + collector.function_reference(ARITHMETIC, "subtract:i64_i64") + + out = stplan.Plan() + collector.write_into(out) + + # Two distinct URNs, densely numbered in order of first reference. + assert [(u.extension_urn_anchor, u.urn) for u in out.extension_urns] == [ + (1, ARITHMETIC), + (2, COMPARISON), + ] + # Three functions, referencing their URN's anchor. + assert [ + ( + d.extension_function.function_anchor, + d.extension_function.name, + d.extension_function.extension_urn_reference, + ) + for d in out.extensions + ] == [ + (1, "add:i64_i64", 1), + (2, "gt:any_any", 2), + (3, "subtract:i64_i64", 1), + ] + + def test_outside_a_build_scope_is_an_error(self): + with pytest.raises(RuntimeError, match="no build in progress"): + function_reference(ARITHMETIC, "add:i64_i64") + + def test_nested_scope_defers_to_the_owner(self): + with build_scope() as (outer, outer_owns): + with build_scope() as (inner, inner_owns): + assert outer_owns and not inner_owns + assert inner is outer + + +class TestPlanLocalAnchors: + def test_anchors_start_at_one(self, full_registry): + plan = _add_plan(full_registry) + assert [u.extension_urn_anchor for u in plan.extension_urns] == [1] + assert _declarations(plan) == {1: "add:i64_i64"} + + def test_independent_of_registry_contents( + self, full_registry, arithmetic_only_registry + ): + """The headline #236 symptom: a one-function plan used to emit anchor 284 + against the default set and 4 against a minimal one.""" + assert _add_plan(full_registry).SerializeToString( + deterministic=True + ) == _add_plan(arithmetic_only_registry).SerializeToString(deterministic=True) + + def test_repeated_builds_are_byte_identical(self, full_registry): + assert _add_plan(full_registry).SerializeToString( + deterministic=True + ) == _add_plan(full_registry).SerializeToString(deterministic=True) + + def test_declarations_are_dense_and_ordered_by_first_use(self, full_registry): + plan = read_named_table("t", NAMED_STRUCT) + plan = project( + plan, + expressions=[ + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), + scalar_function(ARITHMETIC, "subtract", [column("a"), column("b")]), + # `add` again: must reuse its reference rather than allocate a new one. + scalar_function(ARITHMETIC, "add", [column("b"), column("a")]), + ], + )(full_registry) + assert _declarations(plan) == {1: "add:i64_i64", 2: "subtract:i64_i64"} + + +def _foreign_plan(registry, *, urn, name, urn_anchor, function_anchor): + """A plan built "elsewhere": a filter whose condition references its own anchors. + + Deliberately uses anchor numbers a fresh build would hand out to *different* + functions, which is what used to corrupt the merged output. + """ + plan = read_named_table("t", NAMED_STRUCT)(registry) + plan.ClearField("extension_urns") + plan.ClearField("extensions") + plan.extension_urns.append( + ste.SimpleExtensionURN(extension_urn_anchor=urn_anchor, urn=urn) + ) + plan.extensions.append( + ste.SimpleExtensionDeclaration( + extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( + extension_urn_reference=urn_anchor, + function_anchor=function_anchor, + name=name, + ) + ) + ) + root = plan.relations[-1].root + condition = stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction( + function_reference=function_anchor, + arguments=[ + stalg.FunctionArgument( + value=stalg.Expression( + selection=stalg.Expression.FieldReference( + direct_reference=stalg.Expression.ReferenceSegment( + struct_field=stalg.Expression.ReferenceSegment.StructField( + field=0 + ) + ), + root_reference=stalg.Expression.FieldReference.RootReference(), + ) + ) + ) + ], + output_type=stt.Type( + bool=stt.Type.Boolean(nullability=stt.Type.NULLABILITY_REQUIRED) + ), + ) + ) + root.input.CopyFrom( + stalg.Rel(filter=stalg.FilterRel(input=root.input, condition=condition)) + ) + return plan + + +class TestForeignPlans: + # Anchor numbers the incoming plan used. 1 is what this build would hand out + # anyway (no rewrite needed); the larger values force the incoming relations to + # be renumbered, which is the path that used to corrupt the output. + @pytest.mark.parametrize("urn_anchor,function_anchor", [(1, 1), (9, 5), (3, 284)]) + def test_extending_a_foreign_plan_does_not_collide( + self, full_registry, urn_anchor, function_anchor + ): + """#236's correctness bug: this used to emit two URNs at one anchor and two + functions at another, making the plan's function references ambiguous.""" + foreign = _foreign_plan( + full_registry, + urn=COMPARISON, + name="gt:any_any", + urn_anchor=urn_anchor, + function_anchor=function_anchor, + ) + out = project( + foreign, + expressions=[ + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]) + ], + )(full_registry) + + function_anchors = [ + d.extension_function.function_anchor for d in out.extensions + ] + urn_anchors = [u.extension_urn_anchor for u in out.extension_urns] + assert len(function_anchors) == len(set(function_anchors)) + assert len(urn_anchors) == len(set(urn_anchors)) + + declarations = _declarations(out) + assert set(declarations.values()) == {"gt:any_any", "add:i64_i64"} + + # The foreign condition's reference must still resolve to gt, and ours to add. + project_rel = out.relations[-1].root.input.project + foreign_reference = ( + project_rel.input.filter.condition.scalar_function.function_reference + ) + our_reference = project_rel.expressions[0].scalar_function.function_reference + assert declarations[foreign_reference] == "gt:any_any" + assert declarations[our_reference] == "add:i64_i64" + + def test_function_absent_from_the_registry_is_carried_through(self, full_registry): + """Identities come from the declaration, not a catalog lookup, so a plan + referencing an unknown extension function survives being extended.""" + foreign = _foreign_plan( + full_registry, + urn="extension:acme:custom", + name="acme_thing:i64", + urn_anchor=1, + function_anchor=1, + ) + out = project( + foreign, + expressions=[ + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]) + ], + )(full_registry) + + declarations = _declarations(out) + assert set(declarations.values()) == {"acme_thing:i64", "add:i64_i64"} + assert "extension:acme:custom" in {u.urn for u in out.extension_urns} + foreign_reference = out.relations[ + -1 + ].root.input.project.input.filter.condition.scalar_function.function_reference + assert declarations[foreign_reference] == "acme_thing:i64" + + def test_underspecified_declaration_is_preserved_verbatim(self, full_registry): + """pyarrow's ``serialize_expressions`` emits a bare + ``extension_function { name: "add" }`` -- no anchor, no URN, no + ``extension_urns`` entry at all. There is nothing to re-derive (anchor 0 is + not a reference any expression can name), so it must survive untouched + rather than be renumbered or rejected. + """ + expression = stee.ExtendedExpression( + base_schema=NAMED_STRUCT, + extensions=[ + ste.SimpleExtensionDeclaration( + extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( + name="add" + ) + ) + ], + referred_expr=[ + stee.ExpressionReference( + expression=stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction(output_type=I64) + ), + output_names=["total"], + ) + ], + ) + out = project(read_named_table("t", NAMED_STRUCT), expressions=[expression])( + full_registry + ) + + assert list(out.extensions) == list(expression.extensions) + assert not out.extension_urns + + def test_underspecified_declaration_coexists_with_collected_ones( + self, full_registry + ): + """A pass-through declaration keeps anchor 0, so it cannot collide with the + 1-based references allocated for functions this build resolves.""" + expression = stee.ExtendedExpression( + base_schema=NAMED_STRUCT, + extensions=[ + ste.SimpleExtensionDeclaration( + extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( + name="opaque" + ) + ) + ], + referred_expr=[ + stee.ExpressionReference( + expression=stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction(output_type=I64) + ), + output_names=["opaque"], + ) + ], + ) + out = project( + read_named_table("t", NAMED_STRUCT), + expressions=[ + expression, + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), + ], + )(full_registry) + + anchors = [d.extension_function.function_anchor for d in out.extensions] + assert anchors == [0, 1] + assert _declarations(out) == {0: "opaque", 1: "add:i64_i64"} + + def test_rebuilding_a_materialized_plan_is_stable(self, full_registry): + """Extending a plan this library already materialized needs no renumbering, + so the result is the same as building the whole chain in one go.""" + one_shot = read_named_table("t", NAMED_STRUCT) + one_shot = project( + one_shot, + expressions=[ + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]) + ], + ) + one_shot = project(one_shot, expressions=[column("add(a,b)")])(full_registry) + + stepwise = read_named_table("t", NAMED_STRUCT)(full_registry) + stepwise = project( + stepwise, + expressions=[ + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]) + ], + )(full_registry) + stepwise = project(stepwise, expressions=[column("add(a,b)")])(full_registry) + + assert one_shot.SerializeToString( + deterministic=True + ) == stepwise.SerializeToString(deterministic=True) + + +class TestNoPerLevelMerging: + """The collector accumulates declarations once per build rather than having each + verb re-merge its children's, which is the extension half of #207. + + White-box on purpose: the point is that the builders no longer reach for the + merge helpers at all, so a refactor reintroducing per-level merging trips this. + """ + + def test_building_a_chain_never_re_merges_declarations( + self, full_registry, monkeypatch + ): + import substrait.builders.plan as builders_plan + import substrait.utils + + def fail(*args, **kwargs): + raise AssertionError( + "builders re-merged extension declarations; the collector owns them" + ) + + monkeypatch.setattr(substrait.utils, "merge_extension_declarations", fail) + monkeypatch.setattr(substrait.utils, "merge_extension_urns", fail) + for name in ("merge_extension_declarations", "merge_extension_urns"): + if hasattr(builders_plan, name): + monkeypatch.setattr(builders_plan, name, fail) + + functions = ["add", "subtract", "multiply", "divide"] + plan = read_named_table("t", NAMED_STRUCT) + for i in range(40): + plan = project( + plan, + expressions=[ + scalar_function( + ARITHMETIC, + functions[i % len(functions)], + [column("a"), column("b")], + alias=f"c{i}", + ) + ], + ) + out = plan(full_registry) + + # One declaration per distinct function, however long the chain. + assert len(out.extensions) == len(functions) + assert sorted(_declarations(out)) == [1, 2, 3, 4] + + +class TestDeprecations: + def test_lookup_urn_warns(self, full_registry): + with pytest.warns(DeprecationWarning, match="plan-local"): + full_registry.lookup_urn(ARITHMETIC) + + def test_function_entry_anchor_warns(self, full_registry): + entry, _ = full_registry.lookup_function(ARITHMETIC, "add", [I64, I64]) + with pytest.warns(DeprecationWarning, match="plan-local"): + entry.anchor + + def test_has_urn_replaces_lookup_urn(self, full_registry): + assert full_registry.has_urn(ARITHMETIC) + assert not full_registry.has_urn("extension:acme:nope") + assert ARITHMETIC in full_registry.urns() diff --git a/tests/sql/test_sql_anchors.py b/tests/sql/test_sql_anchors.py new file mode 100644 index 00000000..a28b8c6f --- /dev/null +++ b/tests/sql/test_sql_anchors.py @@ -0,0 +1,100 @@ +"""Extension anchor consistency across the SQL translator. + +The translator materializes a Plan at *every* step and builds a set operation's two +sides as independent plans before merging them (see ``sql_to_substrait.translate``). +Because extension anchors are plan-local, each side numbers its functions from 1 +independently -- so folding one plan into another has to re-derive those numbers +rather than trust them. These are plan-only assertions, deliberately not behind the +engine round-trip skip that covers the rest of this directory. +""" + +import substrait.type_pb2 as stt + +from substrait.extension_registry import ExtensionRegistry +from substrait.sql.sql_to_substrait import convert + +I64 = stt.Type(i64=stt.Type.I64(nullability=stt.Type.NULLABILITY_REQUIRED)) + + +def schema_resolver(name: str) -> stt.NamedStruct: + return stt.NamedStruct( + names=["a", "b"], + struct=stt.Type.Struct( + types=[I64, I64], nullability=stt.Type.NULLABILITY_REQUIRED + ), + ) + + +def _declarations(plan): + return { + d.extension_function.function_anchor: d.extension_function.name + for d in plan.extensions + } + + +def _function_references(plan): + """Every function reference appearing anywhere in ``plan``'s relations.""" + found = [] + + def walk(message): + for field, value in message.ListFields(): + if field.name == "function_reference" and isinstance(value, int): + found.append(value) + elif field.message_type is not None: + items = value if hasattr(value, "__len__") else [value] + for item in items: + if hasattr(item, "ListFields"): + walk(item) + + for plan_rel in plan.relations: + walk(plan_rel) + return found + + +def test_union_branches_get_distinct_anchors(): + """Each side of the union numbers its own function 1; the merge must renumber + one of them rather than let two functions share an anchor.""" + registry = ExtensionRegistry(load_default_extensions=True) + plan = convert( + "SELECT a + b FROM t UNION ALL SELECT a - b FROM t", + "generic", + schema_resolver, + registry, + ) + + declarations = _declarations(plan) + anchors = [d.extension_function.function_anchor for d in plan.extensions] + assert len(anchors) == len(set(anchors)), f"anchors collide: {anchors}" + assert set(declarations.values()) == {"add:i64_i64", "subtract:i64_i64"} + + # Every reference in the tree must resolve to a declared function. + references = _function_references(plan) + assert references, "expected the union's branches to reference functions" + assert all(reference in declarations for reference in references) + # ...and both branches' functions must actually be referenced. + assert {declarations[reference] for reference in references} == { + "add:i64_i64", + "subtract:i64_i64", + } + + +def test_urn_anchors_are_distinct_across_branches(): + """Two branches drawing on different extension URNs must not share a URN anchor.""" + registry = ExtensionRegistry(load_default_extensions=True) + plan = convert( + "SELECT a + b FROM t UNION ALL SELECT a FROM t WHERE a > b", + "generic", + schema_resolver, + registry, + ) + + urn_anchors = [u.extension_urn_anchor for u in plan.extension_urns] + assert len(urn_anchors) == len(set(urn_anchors)), ( + f"URN anchors collide: {urn_anchors}" + ) + assert len(plan.extension_urns) == 2, [u.urn for u in plan.extension_urns] + + # Each declaration must point at a URN the plan actually declares. + declared = {u.extension_urn_anchor for u in plan.extension_urns} + for declaration in plan.extensions: + assert declaration.extension_function.extension_urn_reference in declared diff --git a/tests/test_utils.py b/tests/test_utils.py index 7b221fdd..28321424 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -11,6 +11,7 @@ merge_extension_urns, merge_extensions_into, rel_anchor_of, + remap_function_references, to_id_based_outer_references, type_num_names, ) @@ -184,6 +185,114 @@ def test_merge_extension_declarations_rejects_non_function_mapping(): merge_extension_declarations([declaration]) +# --- remap_function_references ------------------------------------------------- +# +# Every proto field that holds a function reference must be rewritten, including +# the two this library never emits itself but a plan built elsewhere may carry. + + +def test_remap_function_references_rewrites_every_reference_field(): + rel = stalg.Rel( + project=stalg.ProjectRel( + expressions=[ + stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction( + function_reference=7 + ) + ), + stalg.Expression( + window_function=stalg.Expression.WindowFunction( + function_reference=7, + sorts=[stalg.SortField(comparison_function_reference=8)], + ) + ), + ] + ) + ) + out = remap_function_references(rel, {7: 1, 8: 2}) + + expressions = out.project.expressions + assert expressions[0].scalar_function.function_reference == 1 + assert expressions[1].window_function.function_reference == 1 + assert expressions[1].window_function.sorts[0].comparison_function_reference == 2 + + +def test_remap_function_references_rewrites_aggregate_and_window_rels(): + aggregate = stalg.Rel( + aggregate=stalg.AggregateRel( + measures=[ + stalg.AggregateRel.Measure( + measure=stalg.AggregateFunction(function_reference=8) + ) + ] + ) + ) + window = stalg.Rel( + window=stalg.ConsistentPartitionWindowRel( + window_functions=[ + stalg.ConsistentPartitionWindowRel.WindowRelFunction( + function_reference=7 + ) + ] + ) + ) + remap = {7: 1, 8: 2} + + assert ( + remap_function_references(aggregate, remap) + .aggregate.measures[0] + .measure.function_reference + == 2 + ) + assert ( + remap_function_references(window, remap) + .window.window_functions[0] + .function_reference + == 1 + ) + + +def test_remap_function_references_rewrites_join_key_comparison(): + """``custom_function_reference`` is never emitted by this library, but a plan + built elsewhere may use it, so the walk must still cover it.""" + key = stalg.ComparisonJoinKey( + comparison=stalg.ComparisonJoinKey.ComparisonType(custom_function_reference=8) + ) + assert ( + remap_function_references(key, {8: 2}).comparison.custom_function_reference == 2 + ) + + +def test_remap_function_references_leaves_input_alone(): + rel = stalg.Rel( + project=stalg.ProjectRel( + expressions=[ + stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction( + function_reference=7 + ) + ) + ] + ) + ) + remap_function_references(rel, {7: 1}) + assert rel.project.expressions[0].scalar_function.function_reference == 7 + + +def test_remap_function_references_empty_remap_is_the_same_object(): + """The no-op case is the common one -- callers rely on it not copying.""" + rel = stalg.Rel(read=stalg.ReadRel()) + assert remap_function_references(rel, {}) is rel + + +def test_remap_function_references_passes_through_unmapped(): + expression = stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction(function_reference=99) + ) + out = remap_function_references(expression, {7: 1}) + assert out.scalar_function.function_reference == 99 + + # --- to_id_based_outer_references ---------------------------------------------- # # Compact hand-built plans exercising the steps_out -> rel_reference conversion.