From 809abcfbe29b70babe7584ca3c841ba2220df16d 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. An input declaring two different functions at one anchor is refused rather than silently resolved to one of them. Anchor 0 is re-derived like any other. The spec marks it a valid anchor/reference (substrait-io/substrait#900, spelled out in the protos since Substrait v0.83.0), and pyarrow's `serialize_expressions` numbers from 0, emitting a bare `extension_function { name: "add" }`. Rewriting such a reference needs the remap walk to read reference fields off the descriptor rather than `ListFields()`, which omits default-valued proto3 scalars; the two reference fields that are oneof members are gated on `WhichOneof`, so an absent one is never invented. Emission stays 1-based, as those same protos ask producers to prefer non-zero values. Note this does not extend to `type_variation_anchor`, where 0 remains reserved for the system-preferred variation. Every builder folds its inputs through the collector, `_inner_rel` included: only a bare `Rel` crosses into `Expression.Subquery`, so a pre-built plan's declarations would otherwise stay behind with the discarded plan and leave its references dangling. `aggregate` now refuses a measure that is not an aggregate function, which it previously emitted as a measure that is set but empty -- the one shape where a present message does not imply a real function reference, and so the one shape that reference renumbering could not treat correctly. 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 removed (`has_urn` / `urns()` replace the former); the urn->function mapping, signature matching and extension-relation registration are unchanged. The pyarrow tests this adds read real `serialize_expressions` output, so they are coupled to a release this project does not control. Third-party integration tests now live in `tests/integration/` behind per-integration markers (`pyarrow`, `duckdb`, `datafusion`), replacing the undocumented `SUBSTRAIT_ENGINE_TESTS` env var, so any one of them can be switched off on its own as those projects catch up. The default deselects `duckdb` and `datafusion` rather than integration testing as a category: handing a lagging consumer a plan built at a newer spec version can crash the interpreter natively, so a red result there is not reliably a report and must never gate a plain `pytest`. pyarrow only produces, so it cannot take the process down and runs by default, where it can catch pyarrow drifting from the output shape the anchor handling assumes. 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` are removed; use `has_urn()` / `urns()` for URN membership, and `(entry.urn, str(entry))` as a function's durable identity. `ExtensionCollector.adopt` now raises on an input declaring two different functions at one anchor, and `aggregate` raises on a measure that is not an aggregate function; both previously produced a plan with an ambiguous or dangling function reference. Closes #236 --- CONTRIBUTING.md | 40 + pyproject.toml | 18 + src/substrait/builders/extended_expression.py | 280 ++--- src/substrait/builders/plan.py | 206 ++-- src/substrait/dataframe/expr.py | 7 - src/substrait/extension_registry/__init__.py | 12 + src/substrait/extension_registry/collector.py | 250 ++++ .../extension_registry/function_entry.py | 2 - src/substrait/extension_registry/registry.py | 25 +- src/substrait/utils/__init__.py | 135 +- .../test_scalar_function.py | 11 +- tests/builders/plan/test_aggregate.py | 21 + tests/dataframe/test_frame.py | 26 + tests/extension_registry/test_collector.py | 1089 +++++++++++++++++ tests/integration/test_pyarrow_producer.py | 192 +++ .../test_sql_engine_roundtrip.py} | 102 +- tests/sql/test_sql_anchors.py | 100 ++ tests/test_utils.py | 274 +++++ 18 files changed, 2435 insertions(+), 355 deletions(-) create mode 100644 src/substrait/extension_registry/collector.py create mode 100644 tests/extension_registry/test_collector.py create mode 100644 tests/integration/test_pyarrow_producer.py rename tests/{sql/test_sql_to_substrait.py => integration/test_sql_engine_roundtrip.py} (77%) create mode 100644 tests/sql/test_sql_anchors.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9ac13151..ffd68897 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -26,3 +26,43 @@ Run tests in the project's root dir. ``` uv run pytest ``` + +## Integration tests + +`tests/integration/` holds tests that run against third-party Substrait +implementations, which release on their own schedule: pyarrow as a producer whose +output we consume, and DuckDB and DataFusion as consumers of the plans we build. + +The default (`addopts` in `pyproject.toml`) deselects `duckdb` and `datafusion`, so +the command above and CI both skip them: handing a lagging consumer a plan built at a +newer spec version can crash the interpreter natively, which no test run can report, +so a red result there is not even reliably a report. **pyarrow runs by default** -- +it produces rather than consumes, so it cannot take the process down, and it is the +only place that would notice pyarrow's output shape drifting away from what the +extension-anchor handling assumes. + +Select with `-m`, which replaces the default rather than narrowing it: + +``` +uv run pytest -m integration # every integration test +uv run pytest -m duckdb # just one integration type +uv run pytest -m "integration and not duckdb" # everything except one +``` + +Mind that `-m` **replaces** the default expression rather than narrowing it, so a `-m` +you meant as a restriction can widen the selection: `-m "not pyarrow"` re-enables +DuckDB and DataFusion, which is the one thing the default exists to prevent. To drop +pyarrow for a single run, skip `-m` and use `uv run pytest +--ignore=tests/integration/test_pyarrow_producer.py`; to drop it for good, add +`and not pyarrow` to the `addopts`. + +Naming a path does not select a deselected marker either -- `uv run pytest +tests/integration/` still reports the engine tests as `deselected` until you pass a +`-m`. + +The per-type markers are `pyarrow`, `duckdb`, and `datafusion`; each integration test +carries `integration` plus its own, so any one of them can be switched on or off +independently as those projects catch up. If a pyarrow release starts failing, add +`and not pyarrow` to the `addopts` rather than deleting the tests -- they record what +changed. New tests in `tests/integration/` need both markers, and any new marker has +to be registered in `[tool.pytest.ini_options]`. diff --git a/pyproject.toml b/pyproject.toml index cef167f2..4a61adbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,24 @@ dev = ["pytest >= 7.0.0", "substrait-antlr==0.99.0", "pyyaml", "sqloxide", "deep [tool.pytest.ini_options] pythonpath = "src" testpaths = "tests" +markers = [ + "integration: exercises a third-party Substrait producer or consumer", + "pyarrow: integration test against pyarrow's Substrait output; runs by default", + "duckdb: integration test against duckdb's Substrait consumer; deselected by default", + "datafusion: integration test against datafusion's Substrait consumer; deselected by default", +] +# The default turns off the two integrations that cannot report their own failure, not +# integration testing as a category. DuckDB and DataFusion consume plans this library +# builds, and handing a lagging consumer a plan built at a newer spec version can crash +# the interpreter natively -- so they must not gate a plain `pytest`, which is what CI +# runs. pyarrow goes the other way (it produces, this library consumes), so it cannot +# take the process down and it currently passes: it runs by default, where it can catch +# a pyarrow release drifting away from the output shape the anchor handling assumes. +# Turn it off by adding `and not pyarrow` here if that day comes. +# +# A `-m` on the command line replaces this one rather than being ANDed with it, so +# `-m integration` runs every integration test and `-m duckdb` runs just that one. +addopts = ["-m", "not duckdb and not datafusion"] [build-system] requires = ["setuptools>=61.0.0", "setuptools_scm[toml]>=6.2.0"] diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index f692a746..2738a551 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( @@ -107,10 +119,23 @@ def resolve( base_schema: stp.NamedStruct, registry: ExtensionRegistry ) -> stee.ExtendedExpression: bound_expression = resolve_expression(expression, base_schema, registry) - bound_expression.referred_expr[0].output_names[0] = name - return bound_expression - - return resolve + # The rename lands on a copy, never on ``bound_expression``: + # ``resolve_expression`` hands back the caller's own message whenever the + # remap is empty (the common case), so renaming in place would rewrite an + # output name in an ExtendedExpression the caller still holds. Copying + # wholesale means dropping the copied declarations too -- the collector, not + # this expression, owns the anchor space until the outermost resolver writes + # it (see ``builders.plan._bind``), and left in place an enclosing builder + # would adopt the stale numbering a second time and re-apply a remap the + # expression already carries. + result = stee.ExtendedExpression() + result.CopyFrom(bound_expression) + result.ClearField("extension_urns") + result.ClearField("extensions") + result.referred_expr[0].output_names[0] = name + return result + + return build_scoped(resolve) _EPOCH_DATE = date(1970, 1, 1) @@ -386,7 +411,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 +455,7 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) class LateralInput: @@ -476,7 +501,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 +551,7 @@ def resolve( base_schema=base_schema, ) - return resolve + return build_scoped(resolve) def scalar_function( @@ -561,36 +586,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 +612,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def aggregate_function( @@ -658,39 +659,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 +690,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) # TODO bounds, sorts @@ -756,40 +729,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 +758,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def if_then( @@ -838,18 +783,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 +822,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def switch( @@ -916,18 +847,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 +869,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 +885,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 +902,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 +921,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 +941,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=extension_urns, - extensions=extensions, ) - return resolve + return build_scoped(resolve) def cast( @@ -1079,11 +974,9 @@ def resolve( ) ], base_schema=base_schema, - extension_urns=bound_input.extension_urns, - extensions=bound_input.extensions, ) - return resolve + return build_scoped(resolve) # -- subqueries ----------------------------------------------------------- @@ -1091,7 +984,7 @@ def resolve( # (a ``registry -> Plan`` callable) -- e.g. a DataFrame's underlying plan. -def _subquery(subquery, base_schema, output_name, *extension_sources): +def _subquery(subquery, base_schema, output_name): return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( @@ -1100,16 +993,17 @@ 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] - ), ) -def _inner_rel(query, registry: ExtensionRegistry, base_schema): +def _inner_rel(query, registry: ExtensionRegistry, base_schema) -> stalg.Rel: + """Resolve a subquery's ``query`` and lift the self-contained Rel to embed. + + The plan itself is not returned: everything of it that survives the subquery + boundary is either folded into the build's collector or inlined into the Rel + below, so a caller holding on to it could only reintroduce the numbering this + already reconciled. + """ # Push the enclosing schema so field references inside the subquery that use # an OuterReference (i.e. correlated columns) resolve against it. stack = outer_schemas.get() @@ -1118,6 +1012,18 @@ def _inner_rel(query, registry: ExtensionRegistry, base_schema): plan = query(registry) if callable(query) else query finally: outer_schemas.reset(token) + # Fold the inner plan's extensions into this build before lifting its Rel: only + # the Rel travels into the Expression.Subquery, so the declarations its function + # references name stay behind with the plan. A plan built elsewhere therefore + # arrives numbered against a table that is about to be discarded -- its + # references dangle, or, worse, silently name one of *this* build's declarations + # where the numbering happens to overlap. The collector re-derives them from the + # durable ``(urn, name)`` identities, the expression-level counterpart of + # ``builders.plan._bind``; unchanged for a plan this build resolved, which + # allocated through the same collector. + collector = current_collector() + if collector is not None: + plan = remap_function_references(plan, collector.adopt(plan)) rel = plan.relations[-1].root.input # An Expression.Subquery embeds only a bare Rel, but a ReferenceRel is # plan-global -- it cannot resolve once lifted out of its plan. So if the @@ -1128,58 +1034,58 @@ def _inner_rel(query, registry: ExtensionRegistry, base_schema): subtrees = plan_subtrees(plan) if subtrees: rel = inline_reference_rels(rel, subtrees) - return plan, rel + return rel def scalar_subquery(query, alias: Union[str, None] = None): """A scalar (one-row, one-column) subquery expression.""" def resolve(base_schema, registry): - plan, rel = _inner_rel(query, registry, base_schema) + rel = _inner_rel(query, registry, base_schema) subquery = stalg.Expression.Subquery( scalar=stalg.Expression.Subquery.Scalar(input=rel) ) - return _subquery(subquery, base_schema, alias or "subquery", plan) + return _subquery(subquery, base_schema, alias or "subquery") - return resolve + return build_scoped(resolve) def set_predicate(query, op, alias: Union[str, None] = None): """An EXISTS / UNIQUE subquery predicate.""" def resolve(base_schema, registry): - plan, rel = _inner_rel(query, registry, base_schema) + rel = _inner_rel(query, registry, base_schema) subquery = stalg.Expression.Subquery( set_predicate=stalg.Expression.Subquery.SetPredicate( predicate_op=op, tuples=rel ) ) - return _subquery(subquery, base_schema, alias or "exists", plan) + return _subquery(subquery, base_schema, alias or "exists") - return resolve + return build_scoped(resolve) def in_predicate(needles, query, alias: Union[str, None] = None): """A ``needles IN (subquery)`` predicate.""" def resolve(base_schema, registry): - plan, rel = _inner_rel(query, registry, base_schema) + rel = _inner_rel(query, registry, base_schema) bound = [resolve_expression(n, base_schema, registry) for n in needles] subquery = stalg.Expression.Subquery( in_predicate=stalg.Expression.Subquery.InPredicate( needles=[b.referred_expr[0].expression for b in bound], haystack=rel ) ) - return _subquery(subquery, base_schema, alias or "in_subquery", plan, *bound) + return _subquery(subquery, base_schema, alias or "in_subquery") - return resolve + return build_scoped(resolve) def set_comparison(left, query, reduction_op, comparison_op, alias=None): """A ``left ANY/ALL (subquery)`` predicate.""" def resolve(base_schema, registry): - plan, rel = _inner_rel(query, registry, base_schema) + rel = _inner_rel(query, registry, base_schema) bound_left = resolve_expression(left, base_schema, registry) subquery = stalg.Expression.Subquery( set_comparison=stalg.Expression.Subquery.SetComparison( @@ -1189,11 +1095,9 @@ def resolve(base_schema, registry): right=rel, ) ) - return _subquery( - subquery, base_schema, alias or "set_comparison", plan, bound_left - ) + return _subquery(subquery, base_schema, alias or "set_comparison") - return resolve + return build_scoped(resolve) def execution_context_variable(variable: str, type_value, alias=None): @@ -1218,7 +1122,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 +1149,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..519217f2 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,29 @@ 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) + # This is the one builder that copies its input wholesale rather than + # assembling a fresh Plan, so it is the one that has to drop the copied + # declarations: `_bind` has already renumbered the relations, and the + # collector -- not this plan -- owns the anchor space until the outermost + # resolver writes it (see `_bind`). Left in place, an enclosing builder + # would adopt the stale numbering a second time and re-apply a remap the + # relations already carry, silently pointing them at other declarations. + # Dropping `extensions` wholesale is only sound because everything it can + # hold is recoverable from the collector, which today means function + # declarations alone: `ExtensionCollector.adopt` raises NotImplementedError + # on any other kind, so whoever teaches it to collect type / type-variation + # declarations must extend `write_into` in the same change or they will be + # silently dropped here. + result.ClearField("extension_urns") + result.ClearField("extensions") result.execution_behavior.variable_eval_mode = variable_eval_mode return result - return resolve + return build_scoped(resolve) def read_named_table( @@ -205,7 +244,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ], ) - return resolve + return build_scoped(resolve) def _require_schema(named_struct: stt.NamedStruct) -> stt.NamedStruct: @@ -262,7 +301,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 +321,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 +341,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ) return _read_plan(named_struct, read_rel) - return resolve + return build_scoped(resolve) def project( @@ -325,7 +364,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 +391,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (_plan, *bound_expressions), ) - return resolve + return build_scoped(resolve) def select( @@ -375,7 +414,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 +448,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (_plan, *bound_expressions), ) - return resolve + return build_scoped(resolve) def filter( @@ -418,7 +457,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 +476,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan, bound_expression), ) - return resolve + return build_scoped(resolve) def sort( @@ -451,7 +490,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 +522,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 +535,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(bound_inputs), ) - return resolve + return build_scoped(resolve) def reference(plan: PlanOrUnbound) -> UnboundPlan: @@ -514,7 +553,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 +569,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: **_merge_plan_metadata(bound), ) - return resolve + return build_scoped(resolve) def fetch( @@ -540,7 +579,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 +606,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan, bound_offset, bound_count), ) - return resolve + return build_scoped(resolve) def join( @@ -580,8 +619,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 +676,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 +703,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 +713,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 +773,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 +782,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 +808,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_left, bound_right), ) - return resolve + return build_scoped(resolve) def aggregate( @@ -795,13 +830,30 @@ 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 = [ resolve_expression(e, ns, registry) for e in grouping_expressions ] bound_measures = [resolve_expression(e, ns, registry) for e in measures] + for m in bound_measures: + # A measure must be an aggregate_function. Reading `.measure` off a + # reference holding an `expression` instead yields a default-constructed + # AggregateFunction, and *assigning* that below would emit a measure that + # is set but empty -- no function reference, no output type. Such a plan + # is already malformed, but it is also the one shape that breaks the + # premise `remap_function_references` relies on (a present message implies + # a real reference), so its unset reference would be renumbered along with + # the genuine ones. Refuse it here rather than emit it. + if m.referred_expr[0].WhichOneof("expr_type") != "measure": + raise ValueError( + "aggregate() measures must be aggregate functions; got a " + f"{m.referred_expr[0].WhichOneof('expr_type')!r} for " + f"{m.referred_expr[0].output_names[0]!r}. Use " + "extended_expression.aggregate_function(...) rather than " + "scalar_function(...)." + ) _filters = ( list(filters) if filters is not None else [None] * len(bound_measures) @@ -854,7 +906,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ), ) - return resolve + return build_scoped(resolve) def write_named_table( @@ -865,7 +917,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 +942,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: include_version=False, ) - return resolve + return build_scoped(resolve) def ddl( @@ -913,11 +965,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 +988,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(merge_sources), ) - return resolve + return build_scoped(resolve) def update( @@ -995,7 +1043,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: **_merge_plan_metadata(*merge_sources), ) - return resolve + return build_scoped(resolve) def consistent_partition_window( @@ -1011,7 +1059,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 +1136,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: ), ) - return resolve + return build_scoped(resolve) def expand( @@ -1105,7 +1153,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 +1190,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(merge_sources), ) - return resolve + return build_scoped(resolve) def nested_loop_join( @@ -1155,8 +1203,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 +1237,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 +1283,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 +1351,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 +1386,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 +1398,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 +1415,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 +1439,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: tuple(bound_inputs), ) - return resolve + return build_scoped(resolve) def exchange( @@ -1406,7 +1454,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 +1473,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: (bound_plan,), ) - return resolve + return build_scoped(resolve) def top_n( @@ -1445,7 +1493,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 +1528,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..332e3d92 --- /dev/null +++ b/src/substrait/extension_registry/collector.py @@ -0,0 +1,250 @@ +"""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. + +Anchor 0 is an ordinary anchor here. The spec marks 0 a valid anchor/reference +(spelled out in the protos since Substrait v0.83.0), and +:func:`substrait.utils.remap_function_references` rewrites a reference of 0 like +any other, so an incoming declaration at anchor 0 is renumbered along with +the rest instead of being preserved. Emission stays 1-based, because those same +protos ask producers to "prefer non-zero values for ergonomics". That wording +covers function, type and URN anchors only -- ``type_variation_anchor`` still +reserves 0 for the system-preferred variation, so nothing here should be read as a +claim about type variations. +""" + +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 _identity_text(identity: FunctionIdentity) -> str: + """A ``(urn, name)`` identity as it reads in an error message.""" + urn, name = identity + return f"{name!r} (urn {urn!r})" if urn is not None else f"{name!r} (no URN)" + + +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 -- 0 is a reference the spec accepts but asks + producers to avoid, so it is read on input (see :meth:`adopt`) and never + emitted. 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) + + 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. Anchor 0 takes part in that like any + other anchor -- which matters for the real-world producer that motivated + this: pyarrow's ``serialize_expressions`` numbers densely from 0 and writes + a bare ``extension_function { name: "add" }`` naming no URN, so it identifies + as ``(None, "add")`` and its ``function_reference: 0`` is rewritten to + whatever this build assigns. Two such expressions folded into one build + therefore get two distinct anchors rather than colliding at 0. + + Raises ``ValueError`` if ``carrier`` itself declares two different functions + at one anchor: no remap can resolve a reference to that anchor, and silently + picking one is exactly the ambiguity this collector exists to prevent. + """ + urns_by_anchor = {u.extension_urn_anchor: u.urn for u in carrier.extension_urns} + remap = {} + identity_by_anchor: dict[int, FunctionIdentity] = {} + for declaration in carrier.extensions: + mapping_type = declaration.WhichOneof("mapping_type") + if mapping_type != "extension_function": + # Type / type-variation declarations are not collected yet; see + # merge_extension_declarations for the same gap. + 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 + # 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) + identity = (urn, function.name) + # An anchor declared twice with the same identity is merely redundant, + # but two identities at one anchor make every reference to it ambiguous, + # and the remap below can only send that anchor to one of them. Anchor 0 + # is checked with the rest: an expression can legitimately reference 0, + # so two functions declared there are exactly as unresolvable as at any + # other anchor. + previous = identity_by_anchor.setdefault(function.function_anchor, identity) + if previous != identity: + raise ValueError( + "ambiguous extension declarations: function_anchor " + f"{function.function_anchor} is declared both as " + f"{_identity_text(previous)} and as {_identity_text(identity)} " + "in one input, so a reference to it names neither" + ) + 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 the emitted + declarations first name them, so a plan's URN anchors are as dense and + plan-local as its function references, and declarations agreeing on a URN + share its one anchor. A declaration whose identity named no resolvable URN + gets no ``extension_urn_reference`` at all; because URN anchors start at 1, + that unset field -- which reads back as 0, a value the spec does allow as a + reference -- still matches no anchor in the emitted table, so it cannot be + misread as naming the first URN. + + Emits function declarations only, because those are the only ones + :meth:`adopt` collects -- it refuses a type or type-variation declaration + rather than dropping it, so nothing reaches here that this cannot write back. + Teaching the collector to carry those means extending both methods together: + their anchors (``type_anchor``, ``type_variation_anchor``) are plan-local in + the same way, so re-deriving them on the way in without emitting them here + would lose them silently. See #247. + """ + urn_anchors: dict[str, int] = {} + urn_anchor_generator = itertools.count(1) + + def anchor_for_urn(urn: str) -> int: + anchor = urn_anchors.get(urn) + if anchor is None: + anchor = next(urn_anchor_generator) + urn_anchors[urn] = anchor + return anchor + + declarations = [] + for (urn, name), reference in self._references.items(): + function = ste.SimpleExtensionDeclaration.ExtensionFunction( + function_anchor=reference, name=name + ) + if urn is not None: + function.extension_urn_reference = anchor_for_urn(urn) + 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..b2bd0d50 100644 --- a/src/substrait/extension_registry/function_entry.py +++ b/src/substrait/extension_registry/function_entry.py @@ -23,14 +23,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 = ( diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 41c3ff0a..5ad082fb 100644 --- a/src/substrait/extension_registry/registry.py +++ b/src/substrait/extension_registry/registry.py @@ -1,6 +1,5 @@ """Extension Registry class.""" -import itertools import re from collections import defaultdict from importlib.resources import files as importlib_files @@ -20,11 +19,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 +77,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 +94,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 +179,13 @@ def find_function( matches = self._find_matching_functions(function_name, signature, urns) return matches[0] if matches else None - def lookup_urn(self, urn: str) -> Optional[int]: - return self._urn_mapping.get(urn, 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, sorted lexicographically.""" + return sorted(self._urns) 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..c41782a6 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,123 @@ 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( + { + # ScalarFunction / WindowFunction / AggregateFunction / WindowRelFunction + "function_reference", + "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. + + A reference of ``0`` is remapped like any other: the spec marks 0 a valid + anchor/reference (spelled out in the protos since Substrait v0.83.0), and + proto3 leaves such a field out of ``ListFields()``, so the fields to rewrite are + read off the descriptor instead -- see :func:`_remap_own_function_references` for + how presence decides which of them may be written without inventing a reference + that was never there. + + A reference packed inside a ``google.protobuf.Any`` is out of reach: the walk + descends into the wrapper, whose only fields are ``type_url`` and the opaque + ``value`` bytes, and finds no reference field there. So an + ``Extension{Single,Multi,Leaf}Rel.detail``, an ``AdvancedExtension.optimization`` + / ``.enhancement``, or a ``ReadRel.ExtensionTable.detail`` that embeds a function + reference keeps the incoming plan's numbering. That is inherent, not an + oversight: rewriting the payload means parsing it, which needs the very schema + ``Any`` withholds -- so whoever produces a packed detail owns the remapping of + the references inside it. + """ + if not remap: + return msg + out = type(msg)() + out.CopyFrom(msg) + _remap_function_references_in_place(out, remap) + return out + + +def _remap_own_function_references(msg, remap: dict) -> None: + """Remap the function-reference fields ``msg`` itself carries (not its children). + + Driven by the descriptor rather than by ``ListFields()``: proto3 omits a + default-valued scalar from ``ListFields()``, so a set-fields walk cannot see -- + let alone rewrite -- a ``function_reference: 0``, which is a reference the spec + permits (since Substrait v0.83.0). + + Field presence decides whether a reference may be written blind: + + * the four ``function_reference`` fields (``Expression.ScalarFunction``, + ``Expression.WindowFunction``, ``AggregateFunction``, + ``ConsistentPartitionWindowRel.WindowRelFunction``) have no presence, and a + containing message that is set always denotes a real function -- there is no + valid ``ScalarFunction`` with no function -- so they are remapped + unconditionally, 0 included. + * ``SortField.comparison_function_reference`` (oneof ``sort_kind``) and + ``ComparisonJoinKey.ComparisonType.custom_function_reference`` (oneof + ``inner_type``) do have presence, so they are remapped only when their oneof + selects them. Writing one blind would *invent* it: a ``SortField`` sorting by + ``direction`` would come out sorting by comparison function instead. + + ``HasField`` cannot serve as the single gate, since it raises on a no-presence + proto3 scalar; the oneof members are gated on ``WhichOneof`` instead. + """ + for name in _FUNCTION_REFERENCE_FIELDS: + field = msg.DESCRIPTOR.fields_by_name.get(name) + if field is None: + continue + oneof = field.containing_oneof + if oneof is not None and msg.WhichOneof(oneof.name) != name: + continue + current = getattr(msg, name) + # A reference field is singular today; a repeated one added upstream would + # arrive as a container, which is skipped rather than mis-assigned. See + # _remap_function_references_in_place for why cardinality is read off the + # value rather than the descriptor. + if isinstance(current, int): + setattr(msg, name, remap.get(current, current)) + + +def _remap_function_references_in_place(msg, remap: dict) -> None: + _remap_own_function_references(msg, remap) + # 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. + # + # Only *set* submessages are descended into -- an unset one holds nothing to + # rewrite, and touching it would materialize it. ListFields() snapshots the set + # fields, so the assignments above are safe to make during iteration. + 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) + + def _iter_direct_subexpressions(msg): """Yield the immediate ``Expression`` messages owned by ``msg``. @@ -507,9 +634,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/builders/plan/test_aggregate.py b/tests/builders/plan/test_aggregate.py index 350102e5..b6aca404 100644 --- a/tests/builders/plan/test_aggregate.py +++ b/tests/builders/plan/test_aggregate.py @@ -1,3 +1,4 @@ +import pytest import substrait.algebra_pb2 as stalg import substrait.extensions.extensions_pb2 as ste import substrait.plan_pb2 as stp @@ -96,3 +97,23 @@ def test_aggregate(): ) assert actual == expected + + +def test_aggregate_rejects_a_non_aggregate_measure(): + """A measure that is not an aggregate function used to be emitted as a measure + that is *set but empty*: reading ``.measure`` off a reference holding an + ``expression`` yields a default-constructed AggregateFunction, and assigning it + sets the field. + + Beyond being malformed on its own, that is the one shape breaking the premise + ``remap_function_references`` relies on -- that a present message implies a real + reference -- so folding such a plan into another build renumbered its *unset* + reference along with the genuine ones, silently pointing an empty measure at a + real function. + """ + table = read_named_table("table", named_struct) + + with pytest.raises(ValueError, match="must be aggregate functions"): + aggregate(table, grouping_expressions=[column("id")], measures=[column("id")])( + registry + ) diff --git a/tests/dataframe/test_frame.py b/tests/dataframe/test_frame.py index 76d5f9aa..a46fbb5a 100644 --- a/tests/dataframe/test_frame.py +++ b/tests/dataframe/test_frame.py @@ -549,6 +549,32 @@ def test_subquery_merges_inner_extensions(): assert "extension:io.substrait:functions_comparison" in urns +def test_subquery_over_a_prebuilt_plan_merges_its_extensions(): + """``sub.DataFrame(other.to_plan())`` as the subquery: a plan that arrives already + numbered, against a table this build is about to replace. + + A Subquery embeds a bare Rel, so the inner plan's declarations do not travel with + it. Where the two numberings happen to overlap -- as they do here, the outer + ``gt`` taking the anchor the inner plan gave ``add`` -- the reference loads fine + and names the wrong function, so this checks by name rather than by anchor. + """ + inner = _inner().filter(sub.col("v") + 1 > 5) + prebuilt = sub.DataFrame(inner.to_plan()) + plan = _outer().filter(sub.col("x") > 0).filter(sub.exists(prebuilt)).to_plan() + + declarations = { + d.extension_function.function_anchor: d.extension_function.name + for d in plan.extensions + } + assert set(declarations.values()) == {"gt:any_any", "add:i64_i64"} + condition = plan.relations[-1].root.input.filter.condition + tuples = condition.subquery.set_predicate.tuples + lifted = tuples.filter.condition.scalar_function + assert declarations[lifted.function_reference] == "gt:any_any" + add = lifted.arguments[0].value.scalar_function + assert declarations[add.function_reference] == "add:i64_i64" + + def test_subquery_requires_dataframe(): with pytest.raises(TypeError, match="expects a DataFrame"): sub.scalar_subquery(sub.col("x")) diff --git a/tests/extension_registry/test_collector.py b/tests/extension_registry/test_collector.py new file mode 100644 index 00000000..d36f10f8 --- /dev/null +++ b/tests/extension_registry/test_collector.py @@ -0,0 +1,1089 @@ +"""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 +from collections.abc import Iterable + +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 google.protobuf.message import Message + +from substrait.builders.extended_expression import ( + alias, + column, + in_predicate, + scalar_function, + scalar_subquery, + set_comparison, + set_predicate, +) +from substrait.builders.plan import ( + filter, + project, + read_named_table, + reference, + with_execution_behavior, +) +from substrait.extension_registry import ( + ExtensionCollector, + ExtensionRegistry, + build_scope, + function_reference, +) + +ARITHMETIC = "extension:io.substrait:functions_arithmetic" +COMPARISON = "extension:io.substrait:functions_comparison" +PER_RECORD = stplan.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_RECORD + +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 + } + + +_REFERENCE_FIELDS = ( + "function_reference", + "comparison_function_reference", + "custom_function_reference", +) + + +def _function_references(msg): + """Every function reference anywhere in ``msg``, a reference of 0 included. + + Deliberately not ``substrait.utils``'s own walk: a test reusing the code under + test could not notice it missing a field. Presence is read the way it has to be + -- ``function_reference`` has none, so it counts whenever its containing message + is set, which is the only way a reference of 0 is visible at all; the two oneof + members count only when their oneof selects them. + """ + found = [] + for name in _REFERENCE_FIELDS: + field = msg.DESCRIPTOR.fields_by_name.get(name) + if field is None: + continue + oneof = field.containing_oneof + if oneof is not None and msg.WhichOneof(oneof.name) != name: + continue + found.append(getattr(msg, name)) + for _, value in msg.ListFields(): + if isinstance(value, Message): + found += _function_references(value) + elif isinstance(value, Iterable) and not isinstance(value, (str, bytes)): + items = value.values() if hasattr(value, "values") else value + for item in items: + if isinstance(item, Message): + found += _function_references(item) + return found + + +def _resolved_functions(plan): + """The set of function names every reference in ``plan`` resolves to. + + Raises ``KeyError`` on a dangling reference -- one naming an anchor the plan does + not declare -- which is the failure a lost declaration produces. + """ + declarations = _declarations(plan) + return {declarations[reference] for reference in _function_references(plan)} + + +def _declaration(name, *, function_anchor=0, urn_reference=0): + """A function declaration with its plan-local anchors spelled out.""" + return ste.SimpleExtensionDeclaration( + extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( + extension_urn_reference=urn_reference, + function_anchor=function_anchor, + name=name, + ) + ) + + +def _carrier(*declarations, urns=()): + """A carrier holding only the two fields ``adopt`` reads, from ``(anchor, urn)``. + + A Plan rather than an ExtendedExpression merely because it is the shorter of the + two to spell; ``adopt`` treats them alike. + """ + return stplan.Plan( + extension_urns=[ + ste.SimpleExtensionURN(extension_urn_anchor=anchor, urn=urn) + for anchor, urn in urns + ], + extensions=declarations, + ) + + +def _anchor_zero_expression(declaration, *, urns=(), output_name="opaque"): + """An ExtendedExpression naming ``declaration`` at anchor 0, from both ends. + + The shape pyarrow's ``serialize_expressions`` emits: a declaration at anchor 0 + and a ``function_reference`` left at 0 to name it. 0 is a valid + anchor/reference (spelled out in the protos since Substrait v0.83.0), so this is + an ordinary carrier -- what makes it worth its own + helper is that proto3 omits both zero-valued fields from the wire, so a + ``ListFields()`` walk sees neither. + """ + return stee.ExtendedExpression( + base_schema=NAMED_STRUCT, + extension_urns=[ + ste.SimpleExtensionURN(extension_urn_anchor=anchor, urn=urn) + for anchor, urn in urns + ], + extensions=[declaration], + referred_expr=[ + stee.ExpressionReference( + expression=stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction(output_type=I64) + ), + output_names=[output_name], + ) + ], + ) + + +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_declarations_dedupe_on_the_urn_they_name(self): + """Declarations are deduplicated on the identity they arrived with, not on + their bytes: the same function reaching two inputs under different incoming + URN anchors is one declaration, while two that differ only in which URN they + name stay distinct. + + Spelled with anchor-0 declarations because that is where the two could most + easily be confused -- their bytes differ only in the URN reference, which is + the field the identity is read from. + """ + collector = ExtensionCollector() + collector.adopt( + _carrier(_declaration("floor", urn_reference=7), urns=[(7, COMPARISON)]) + ) + collector.adopt( + _carrier(_declaration("floor", urn_reference=3), urns=[(3, COMPARISON)]) + ) + collector.adopt( + _carrier(_declaration("floor", urn_reference=1), urns=[(1, ARITHMETIC)]) + ) + + out = stplan.Plan() + collector.write_into(out) + + urns = {u.extension_urn_anchor: u.urn for u in out.extension_urns} + assert [ + ( + d.extension_function.function_anchor, + d.extension_function.name, + urns[d.extension_function.extension_urn_reference], + ) + for d in out.extensions + ] == [(1, "floor", COMPARISON), (2, "floor", ARITHMETIC)] + + def test_a_type_declaration_is_refused_rather_than_dropped(self): + """Type and type-variation declarations have anchors of their own that nothing + re-derives yet, so collecting one would silently lose it (the same gap + ``merge_extension_declarations`` reports).""" + carrier = _carrier( + ste.SimpleExtensionDeclaration( + extension_type=ste.SimpleExtensionDeclaration.ExtensionType( + type_anchor=1, name="point" + ) + ), + urns=[(1, "extension:acme:custom")], + ) + with pytest.raises(NotImplementedError, match="extension_type"): + ExtensionCollector().adopt(carrier) + + 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, inners=()): + """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. + + ``inners`` is a sequence of further ``(urn, name, urn_anchor, function_anchor)`` + functions, nested left to right into the outer one's argument, so the plan + carries a reference at each of several distinct paths. They are declared *after* + the outer one, so numbering by first use in declaration order can be made to + permute the incoming anchors rather than merely shift them -- see + ``TestForeignPlans.test_execution_behavior_wrapper_does_not_double_remap``. + """ + plan = read_named_table("t", NAMED_STRUCT)(registry) + plan.ClearField("extension_urns") + plan.ClearField("extensions") + declared = [(urn, name, urn_anchor, function_anchor), *inners] + for d_urn, d_name, d_urn_anchor, d_function_anchor in declared: + plan.extension_urns.append( + ste.SimpleExtensionURN(extension_urn_anchor=d_urn_anchor, urn=d_urn) + ) + plan.extensions.append( + ste.SimpleExtensionDeclaration( + extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( + extension_urn_reference=d_urn_anchor, + function_anchor=d_function_anchor, + name=d_name, + ) + ) + ) + root = plan.relations[-1].root + argument = 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(), + ) + ) + for inner in reversed(inners): + argument = stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction( + function_reference=inner[3], + arguments=[stalg.FunctionArgument(value=argument)], + output_type=I64, + ) + ) + condition = stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction( + function_reference=function_anchor, + arguments=[stalg.FunctionArgument(value=argument)], + 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" + + @pytest.mark.parametrize("nested", [False, True], ids=["outermost", "nested"]) + def test_execution_behavior_wrapper_does_not_double_remap( + self, full_registry, nested + ): + """``with_execution_behavior`` copies its input plan wholesale, so it used to + forward the input's declarations -- stale, because ``_bind`` had already + renumbered the relations against the collector. An enclosing builder adopted + that numbering a second time and re-applied a remap the relations already + carried, silently pointing each reference at the other's declaration. + + The foreign layout makes re-deriving the pair *permute* the incoming anchors + (gt@2 with subtract@1 becomes the swap ``{2: 1, 1: 2}``). A layout that only + shifts them hides this: re-applying a remap whose new values are not + themselves keys is accidentally idempotent. + """ + foreign = _foreign_plan( + full_registry, + urn=COMPARISON, + name="gt:any_any", + urn_anchor=1, + function_anchor=2, + inners=[(ARITHMETIC, "subtract:i64_i64", 2, 1)], + ) + unbound = with_execution_behavior(foreign, PER_RECORD) + if nested: + unbound = project(unbound, expressions=[column("a")]) + out = unbound(full_registry) + + root_input = out.relations[-1].root.input + condition = ( + root_input.project.input if nested else root_input + ).filter.condition.scalar_function + declarations = _declarations(out) + assert declarations[condition.function_reference] == "gt:any_any" + inner_reference = condition.arguments[ + 0 + ].value.scalar_function.function_reference + assert declarations[inner_reference] == "subtract:i64_i64" + # The wrapper still does its actual job. + assert out.execution_behavior.variable_eval_mode == PER_RECORD + + 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_anchor_zero_declaration_is_renumbered_with_its_reference( + self, full_registry + ): + """pyarrow's ``serialize_expressions`` emits a bare + ``extension_function { name: "add" }`` at anchor 0, named by a + ``function_reference`` left at 0. Both are valid + (since Substrait v0.83.0), so the declaration is re-derived like any + other -- 1-based on the way out, because the same protos ask producers to + prefer non-zero -- and the reference that named it moves with it. + + What must survive is the *identity* (``add``, naming no URN) and the + reference's *meaning*, not the number 0. + """ + expression = _anchor_zero_expression(_declaration("add"), output_name="total") + before = expression.SerializeToString(deterministic=True) + out = project(read_named_table("t", NAMED_STRUCT), expressions=[expression])( + full_registry + ) + + assert _declarations(out) == {1: "add"} + # Naming no URN, it must still name none: an unset extension_urn_reference + # reads back as 0, which is why URN anchors are emitted 1-based. + assert not out.extension_urns + assert out.extensions[0].extension_function.extension_urn_reference == 0 + project_rel = out.relations[-1].root.input.project + assert project_rel.expressions[0].scalar_function.function_reference == 1 + # The caller's expression is rewritten on a copy, not in place. + assert expression.SerializeToString(deterministic=True) == before + + def test_anchor_zero_declaration_coexists_with_collected_ones(self, full_registry): + """An incoming anchor-0 declaration is numbered into the same dense 1-based + space as the functions this build resolves, so the two cannot collide -- and + neither reference is left naming the other's declaration.""" + expression = _anchor_zero_expression(_declaration("opaque")) + out = project( + read_named_table("t", NAMED_STRUCT), + expressions=[ + expression, + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), + ], + )(full_registry) + + declarations = _declarations(out) + assert declarations == {1: "opaque", 2: "add:i64_i64"} + references = [ + e.scalar_function.function_reference + for e in out.relations[-1].root.input.project.expressions + ] + assert [declarations[r] for r in references] == ["opaque", "add:i64_i64"] + + def test_anchor_zero_declaration_naming_a_urn_is_renumbered_too( + self, full_registry + ): + """Resolving a URN changes nothing about the anchor: 0 is renumbered either + way, and the URN reference is re-derived alongside it -- 7 indexes the + incoming carrier's URN table, which this build replaces, so it must end up + naming a URN the output actually declares rather than dangling. + """ + expression = _anchor_zero_expression( + _declaration("gt", urn_reference=7), + urns=[(7, COMPARISON)], + output_name="flag", + ) + out = project( + read_named_table("t", NAMED_STRUCT), + expressions=[ + expression, + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), + ], + )(full_registry) + + declarations = _declarations(out) + assert declarations == {1: "gt", 2: "add:i64_i64"} + references = [ + e.scalar_function.function_reference + for e in out.relations[-1].root.input.project.expressions + ] + assert [declarations[r] for r in references] == ["gt", "add:i64_i64"] + + urns = {u.extension_urn_anchor: u.urn for u in out.extension_urns} + assert urns == {1: COMPARISON, 2: ARITHMETIC} + assert { + declarations[d.extension_function.function_anchor]: urns[ + d.extension_function.extension_urn_reference + ] + for d in out.extensions + } == {"gt": COMPARISON, "add:i64_i64": ARITHMETIC} + + def test_zero_based_foreign_plan_folds_in(self, full_registry): + """A producer numbering from 0 rather than 1 is not a special case: 0 is a + valid anchor/reference (since Substrait v0.83.0), so a plan whose URN + anchors, function anchors and references all run 0,1,2 is renumbered whole. + + Every one of its three references sits at a different depth, so a walk that + cannot see a reference of 0 leaves the outermost one naming this build's + first declaration instead -- a silently different function, not a load error. + """ + foreign = _foreign_plan( + full_registry, + urn=COMPARISON, + name="gt:any_any", + urn_anchor=0, + function_anchor=0, + inners=[ + (ARITHMETIC, "add:i64_i64", 1, 1), + (ARITHMETIC, "subtract:i64_i64", 2, 2), + ], + ) + out = project( + foreign, + expressions=[ + scalar_function(ARITHMETIC, "multiply", [column("a"), column("b")]) + ], + )(full_registry) + + declarations = _declarations(out) + assert declarations == { + 1: "gt:any_any", + 2: "add:i64_i64", + 3: "subtract:i64_i64", + 4: "multiply:i64_i64", + } + assert _resolved_functions(out) == set(declarations.values()) + + # Each reference by path, so the assertion above cannot be satisfied by three + # references that all happen to resolve to *some* declared function. + condition = out.relations[ + -1 + ].root.input.project.input.filter.condition.scalar_function + add = condition.arguments[0].value.scalar_function + subtract = add.arguments[0].value.scalar_function + assert declarations[condition.function_reference] == "gt:any_any" + assert declarations[add.function_reference] == "add:i64_i64" + assert declarations[subtract.function_reference] == "subtract:i64_i64" + + # The incoming plan declared ARITHMETIC at two anchors; emission collapses + # them, and no URN anchor is emitted at 0 for an unset reference to match. + urns = {u.extension_urn_anchor: u.urn for u in out.extension_urns} + assert urns == {1: COMPARISON, 2: ARITHMETIC} + + 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) + + +def _inner_query(): + """A subquery's inner query: ``SELECT * FROM u WHERE a + b > a``. + + Its two functions are declared innermost-first (``add`` at 1, ``gt`` at 2), so a + build that resolved ``gt`` before folding this in maps them ``{1: 2, 2: 1}``: a + *permutation*, which a remap applied twice would swap back and a dropped + declaration cannot fake. A layout that only shifts the anchors hides both. + """ + return filter( + read_named_table("u", NAMED_STRUCT), + expression=scalar_function( + COMPARISON, + "gt", + [ + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), + column("a"), + ], + ), + ) + + +# ``builder(query) -> UnboundExtendedExpression``, and the bound expression's path to +# the Rel the inner query was lifted into -- one entry per subquery flavour, since each +# embeds its Rel in a different field of ``Expression.Subquery``. +_SUBQUERY_BUILDERS = { + "scalar_subquery": ( + scalar_subquery, + lambda expression: expression.subquery.scalar.input, + ), + "set_predicate": ( + lambda query: set_predicate( + query, stalg.Expression.Subquery.SetPredicate.PREDICATE_OP_EXISTS + ), + lambda expression: expression.subquery.set_predicate.tuples, + ), + "in_predicate": ( + lambda query: in_predicate([column("a")], query), + lambda expression: expression.subquery.in_predicate.haystack, + ), + "set_comparison": ( + lambda query: set_comparison( + column("a"), + query, + stalg.Expression.Subquery.SetComparison.REDUCTION_OP_ANY, + stalg.Expression.Subquery.SetComparison.COMPARISON_OP_EQ, + ), + lambda expression: expression.subquery.set_comparison.right, + ), +} + + +class TestSubqueries: + """An ``Expression.Subquery`` embeds a bare Rel, so the inner query's declarations + do not travel with it -- they have to be folded into the enclosing build instead. + + ``query`` may be an UnboundPlan, which resolves inside this build's scope and so + numbers itself against the same collector, or an already-built Plan (a DataFrame's + ``to_plan()`` handed back as a subquery), which arrives numbered against a table + that is about to be discarded. The second used to lose them: its references were + left naming anchors this plan does not declare, or -- where the numbering happened + to overlap -- one of *this* build's declarations, which loads fine and computes + something else. + """ + + @pytest.mark.parametrize("prebuilt", [False, True], ids=["unbound", "pre-built"]) + @pytest.mark.parametrize( + "builder,inner_rel", + list(_SUBQUERY_BUILDERS.values()), + ids=list(_SUBQUERY_BUILDERS), + ) + def test_inner_query_functions_are_declared( + self, full_registry, builder, inner_rel, prebuilt + ): + inner = _inner_query() + query = inner(full_registry) if prebuilt else inner + before = query.SerializeToString(deterministic=True) if prebuilt else None + out = project( + read_named_table("t", NAMED_STRUCT), + expressions=[ + # Resolved first, so ``gt`` is this build's anchor 1 and folding the + # inner query in has to permute its pair rather than shift it. + scalar_function(COMPARISON, "gt", [column("a"), column("b")]), + builder(query), + ], + )(full_registry) + + declarations = _declarations(out) + assert declarations == {1: "gt:any_any", 2: "add:i64_i64"} + # Nothing anywhere in the plan names an anchor it does not declare. + assert _resolved_functions(out) == {"gt:any_any", "add:i64_i64"} + # By path, so the set above cannot be satisfied by references that resolve to + # *some* declared function: the lifted condition still means gt(add(a, b), a). + expressions = out.relations[-1].root.input.project.expressions + condition = inner_rel(expressions[1]).filter.condition.scalar_function + assert declarations[condition.function_reference] == "gt:any_any" + add = condition.arguments[0].value.scalar_function + assert declarations[add.function_reference] == "add:i64_i64" + assert declarations[expressions[0].scalar_function.function_reference] == ( + "gt:any_any" + ) + if prebuilt: + # Folded on a copy: the caller still holds its plan, numbered its own way. + assert query.SerializeToString(deterministic=True) == before + + def test_shared_subtree_lifted_into_a_subquery_is_renumbered_too( + self, full_registry + ): + """A pre-built inner query carrying a shared subtree (what ``DataFrame.cache`` + produces) is inlined into the subquery Rel, because a ReferenceRel is + plan-global and cannot resolve once lifted out of its plan. The references + inside that subtree need re-deriving like any others, so the fold has to + happen before the subtree is read -- inlining first would splice the inner + plan's own numbering into a plan that no longer has that table. + """ + inner = filter( + reference( + project( + read_named_table("u", NAMED_STRUCT), + expressions=[ + column("a"), + column("b"), + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), + ], + ) + ), + expression=scalar_function(COMPARISON, "gt", [column("a"), column("b")]), + )(full_registry) + assert _declarations(inner) == {1: "add:i64_i64", 2: "gt:any_any"} + + out = project( + read_named_table("t", NAMED_STRUCT), + expressions=[ + scalar_function(COMPARISON, "gt", [column("a"), column("b")]), + set_predicate( + inner, stalg.Expression.Subquery.SetPredicate.PREDICATE_OP_EXISTS + ), + ], + )(full_registry) + + declarations = _declarations(out) + assert declarations == {1: "gt:any_any", 2: "add:i64_i64"} + assert _resolved_functions(out) == {"gt:any_any", "add:i64_i64"} + tuples = ( + out.relations[-1] + .root.input.project.expressions[1] + .subquery.set_predicate.tuples + ) + # The subtree is inlined where the ReferenceRel stood, carrying the `add` whose + # reference the inner plan numbered 1 -- this build's `gt`. + inlined = tuples.filter.input.project.expressions[-1].scalar_function + assert declarations[inlined.function_reference] == "add:i64_i64" + assert ( + declarations[tuples.filter.condition.scalar_function.function_reference] + == "gt:any_any" + ) + + +class TestAlias: + """``alias`` binds an expression and hands back a message with one output name + changed, so it is the one expression builder whose result can be its *input*. + + Renaming in place would rewrite a column name in an ExtendedExpression the caller + still holds, and returning the input's declarations would leak an anchor space the + collector owns until the outermost resolver writes it. + """ + + @pytest.mark.parametrize( + "permuting", [False, True], ids=["identity-remap", "permuting-remap"] + ) + def test_nested_alias_declares_nothing_and_leaves_its_input_alone( + self, full_registry, permuting + ): + expression = scalar_function( + ARITHMETIC, + "add", + [ + column("a"), + scalar_function(ARITHMETIC, "subtract", [column("a"), column("b")]), + ], + )(NAMED_STRUCT, full_registry) + assert _declarations(expression) == {1: "subtract:i64_i64", 2: "add:i64_i64"} + before = expression.SerializeToString(deterministic=True) + + with build_scope() as (collector, owns_scope): + # This scope owns the build, so ``alias``'s own wrapper is nested. + assert owns_scope + if permuting: + # Allocated in the opposite order, so the fold maps {1: 2, 2: 1}: a + # remap the expression must not be seen carrying already. + collector.function_reference(ARITHMETIC, "add:i64_i64") + collector.function_reference(ARITHMETIC, "subtract:i64_i64") + out = alias(expression, "renamed")(NAMED_STRUCT, full_registry) + written = stplan.Plan() + collector.write_into(written) + + assert out.referred_expr[0].output_names[0] == "renamed" + assert not out.extensions and not out.extension_urns + # The returned copy's references resolve through the collector's numbering + # rather than the input's. + declarations = _declarations(written) + add = out.referred_expr[0].expression.scalar_function + subtract = add.arguments[1].value.scalar_function + assert declarations[add.function_reference] == "add:i64_i64" + assert declarations[subtract.function_reference] == "subtract:i64_i64" + # Untouched, whichever way the fold went: with an identity remap the fold + # hands back the caller's own message, so the rename has to land on a copy. + assert expression.SerializeToString(deterministic=True) == before + + def test_outermost_alias_still_declares_what_it_renames(self, full_registry): + """Dropping the copied declarations is sound only because the collector holds + the same information: reached directly, ``alias``'s scope is the outermost one + and writes them itself.""" + out = alias( + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), "renamed" + )(NAMED_STRUCT, full_registry) + + assert _declarations(out) == {1: "add:i64_i64"} + assert _resolved_functions(out) == {"add:i64_i64"} + assert [(u.extension_urn_anchor, u.urn) for u in out.extension_urns] == [ + (1, ARITHMETIC) + ] + assert out.referred_expr[0].output_names[0] == "renamed" + + +class TestAmbiguousInput: + """An input that declares two different functions at one anchor is rejected. + + Every reference to that anchor names both, and ``adopt``'s remap can send it to + only one of them, so folding such a plan in would silently pick a function -- the + ambiguity the collector exists to rule out. + + Parametrized over anchor 0 as well as an ordinary one: 0 is a valid + anchor/reference (since Substrait v0.83.0), so it is policed like the rest + rather than exempted from the check. + """ + + @pytest.mark.parametrize("anchor", [0, 5]) + def test_two_identities_at_one_anchor_are_rejected(self, anchor): + carrier = _carrier( + _declaration("gt:any_any", function_anchor=anchor, urn_reference=1), + _declaration("lt:any_any", function_anchor=anchor, urn_reference=1), + urns=[(1, COMPARISON)], + ) + with pytest.raises(ValueError, match=f"function_anchor {anchor}"): + ExtensionCollector().adopt(carrier) + + def test_two_bare_declarations_at_anchor_zero_are_rejected(self): + """The shape a producer numbering from 0 would have to emit to be ambiguous: + two bare ``extension_function { name: ... }`` in one carrier, both at anchor + 0, so its ``function_reference: 0`` names neither. + + pyarrow does not do this -- it numbers densely, so one carrier holds at most + one declaration at 0 (checked against real pyarrow output in + ``tests/integration/test_pyarrow_producer.py``) -- but this used to be + accepted, with both declarations emitted at anchor 0 and every reference to + them resolving to whichever the consumer found first. + """ + carrier = _carrier(_declaration("add"), _declaration("subtract")) + with pytest.raises( + ValueError, + match="function_anchor 0 is declared both as 'add' \\(no URN\\) " + "and as 'subtract' \\(no URN\\)", + ): + ExtensionCollector().adopt(carrier) + + @pytest.mark.parametrize("anchor", [0, 5]) + def test_one_identity_declared_twice_is_accepted(self, anchor): + """Redundant rather than ambiguous: both declarations mean the same function, + so a reference to the anchor still resolves to exactly one thing.""" + carrier = _carrier( + _declaration("gt:any_any", function_anchor=anchor, urn_reference=1), + _declaration("gt:any_any", function_anchor=anchor, urn_reference=1), + urns=[(1, COMPARISON)], + ) + assert ExtensionCollector().adopt(carrier) == {anchor: 1} + + @pytest.mark.parametrize("anchor", [0, 5]) + def test_one_anchor_meaning_two_things_in_separate_inputs_is_accepted(self, anchor): + """Anchors are plan-local, so two independently built inputs meeting at a + multi-input relation number theirs separately and are remapped separately. + + At anchor 0 this is the pyarrow case: each serialized expression is its own + anchor space, so two of them declaring different functions at 0 is ordinary + input that must come out as two declarations. + """ + collector = ExtensionCollector() + gt = _carrier( + _declaration("gt:any_any", function_anchor=anchor, urn_reference=1), + urns=[(1, COMPARISON)], + ) + lt = _carrier( + _declaration("lt:any_any", function_anchor=anchor, urn_reference=1), + urns=[(1, COMPARISON)], + ) + assert collector.adopt(gt) == {anchor: 1} + assert collector.adopt(lt) == {anchor: 2} + + def test_two_bare_declarations_at_anchor_zero_in_separate_inputs_are_renumbered( + self, + ): + """The same as above with no URN to resolve, which is the shape pyarrow emits. + + Kept separate because the URN is what made the case above uninteresting to the + numbering this replaced: it treated a declaration as unrewritable only when + anchor 0 came *with* an unresolvable URN, so a carrier naming a real URN was + renumbered normally either way. A bare declaration was the one that got frozen + at 0 -- so two of them, arriving from two separately serialized inputs, both + stayed at 0 and every reference to it resolved to whichever the consumer found + first. That is the pyarrow collision, and this is its pyarrow-free guard: + ``tests/integration`` covers it through the real bytes, but that suite can be + switched off. + """ + collector = ExtensionCollector() + assert collector.adopt(_carrier(_declaration("add"))) == {0: 1} + assert collector.adopt(_carrier(_declaration("multiply"))) == {0: 2} + + out = stplan.Plan() + collector.write_into(out) + assert _declarations(out) == {1: "add", 2: "multiply"} + + +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. + """ + + @pytest.fixture + def merge_helpers_are_fatal(self, monkeypatch): + """Make every legacy merge helper explode, wherever a builder could reach it. + + Patched both in ``substrait.utils`` and in each builder module's own + namespace: ``from substrait.utils import merge_...`` binds a module-local name + at import time, which a patch of ``substrait.utils`` alone would never reach. + Both builder modules are covered -- either could reintroduce the import. + """ + import substrait.builders.extended_expression as builders_expression + 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" + ) + + names = ( + "merge_extension_declarations", + "merge_extension_urns", + "merge_extensions_into", + ) + for module in (substrait.utils, builders_plan, builders_expression): + for name in names: + if hasattr(module, name): + monkeypatch.setattr(module, name, fail) + + def test_building_a_chain_never_re_merges_declarations( + self, full_registry, merge_helpers_are_fatal + ): + 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] + + def test_building_an_expression_never_re_merges_declarations( + self, full_registry, merge_helpers_are_fatal + ): + """An ExtendedExpression built on its own goes through the same collector, so + the expression builders have no more use for the helpers than the plan ones. + """ + expression = scalar_function( + ARITHMETIC, + "add", + [ + column("a"), + scalar_function(ARITHMETIC, "subtract", [column("a"), column("b")]), + ], + )(NAMED_STRUCT, full_registry) + + assert _declarations(expression) == {1: "subtract:i64_i64", 2: "add:i64_i64"} + + +class TestAnchorsAreNotARegistryConcern: + """The registry hands out no anchors, so the two members that implied it does are + gone rather than deprecated: ``lookup_urn`` returned a URN anchor and + ``FunctionEntry.anchor`` a function anchor, and neither is knowable from a catalog + (#236). Pinned as absent so a refactor cannot quietly bring them back. + """ + + def test_the_registry_has_no_lookup_urn(self, full_registry): + assert not hasattr(ExtensionRegistry, "lookup_urn") + assert not hasattr(full_registry, "lookup_urn") + + def test_a_function_entry_has_no_anchor(self, full_registry): + entry, _ = full_registry.lookup_function(ARITHMETIC, "add", [I64, I64]) + assert not hasattr(type(entry), "anchor") + assert not hasattr(entry, "anchor") + # What replaces it: the identity a declaration can be re-derived from. + assert entry.urn == ARITHMETIC + assert str(entry) == "add:i64_i64" + + 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/integration/test_pyarrow_producer.py b/tests/integration/test_pyarrow_producer.py new file mode 100644 index 00000000..097fa6d6 --- /dev/null +++ b/tests/integration/test_pyarrow_producer.py @@ -0,0 +1,192 @@ +"""Real ``pyarrow.substrait.serialize_expressions`` output folded into a build. + +The direction here is the opposite of the engine round-trips next door: pyarrow is the +*producer* and this library the consumer, so nothing hands a newer-spec plan to an +older consumer and there is no native-crash risk. What makes these integration tests +is the coupling to a release we do not control: the numbering and URN shape pyarrow +emits is an observation about pyarrow, not a contract of this package. + +They run by default anyway, unlike the engine round-trips. Nothing here can take the +process down, and while pyarrow does emit this shape these are the only tests that +would notice it changing -- the anchor handling is built on that observation, so +finding out from a red test beats finding out from a user. If a pyarrow release does +drift, add ``and not pyarrow`` to the ``addopts`` in pyproject.toml and the suite goes +green again without losing the tests that recorded what changed. + +Nothing here is the only guard on the collector behaviour it exercises. Held by +pyarrow-free tests in ``tests/extension_registry/test_collector.py``, which run by +default: + +- ``TestAmbiguousInput::test_two_bare_declarations_at_anchor_zero_in_separate_inputs_are_renumbered`` + is the collision below in pyarrow-free form: two inputs each carrying one bare + anchor-0 declaration must come out as two distinct declarations. Note its sibling + ``test_one_anchor_meaning_two_things_in_separate_inputs_is_accepted`` is *not* this + guard despite the similar name -- its carriers name a resolvable URN, which the + numbering this replaced already renumbered normally, so it passes either way. +- ``TestForeignPlans::test_zero_based_foreign_plan_folds_in`` covers a carrier + numbering densely from 0, the way one ``serialize_expressions`` call does. +- ``TestForeignPlans::test_anchor_zero_declaration_*`` cover renumbering an anchor-0 + declaration together with the reference naming it, with and without a URN. + +So what these add is only that pyarrow really does emit that shape: they go through +the actual bytes instead of an imitation of them, over the same shape +``examples/pyarrow_example.py`` exercises. +""" + +import pytest +import substrait.extended_expression_pb2 as stee + +from substrait.builders.extended_expression import column, scalar_function +from substrait.builders.plan import project, read_named_table +from substrait.extension_registry import ExtensionRegistry + +pytestmark = [pytest.mark.integration, pytest.mark.pyarrow] + +# A second guard behind the marker, so an explicit ``-m pyarrow`` on a machine without +# pyarrow skips rather than errors on collection. +pytest.importorskip("pyarrow") + +ARITHMETIC = "extension:io.substrait:functions_arithmetic" + + +@pytest.fixture(scope="module") +def full_registry(): + return ExtensionRegistry(load_default_extensions=True) + + +@pytest.fixture(scope="module") +def serialize(): + """``*(pyarrow.compute -> expression) -> ExtendedExpression``, via pyarrow. + + One call serializes all of the given expressions into a single carrier, so the + number of arguments chooses between pyarrow's one-declaration and its + densely-numbered multi-declaration output. ``importorskip`` for each piece because + pyarrow is not a dependency of this package, only the producer whose output this + has to accept -- the marker is what keeps these out of a default run, and this is + what keeps a marker-selected run honest where pyarrow is absent. + """ + pa = pytest.importorskip("pyarrow") + pc = pytest.importorskip("pyarrow.compute") + pa_substrait = pytest.importorskip("pyarrow.substrait") + schema = pa.schema([pa.field("a", pa.int64()), pa.field("b", pa.int64())]) + + def _serialize(*expressions): + buffer = pa_substrait.serialize_expressions( + exprs=[build(pc) for build in expressions], + names=[f"e{i}" for i in range(len(expressions))], + schema=schema, + ) + return stee.ExtendedExpression.FromString(bytes(buffer)) + + return _serialize + + +def _declarations(carrier): + """``{function anchor: name}`` for a carrier's extension declarations.""" + return { + declaration.extension_function.function_anchor: declaration.extension_function.name + for declaration in carrier.extensions + } + + +def _project_functions(plan): + """The function each expression of ``plan``'s ProjectRel resolves to, in order. + + Resolved through the plan's own declarations, so a reference naming an anchor the + plan does not declare raises ``KeyError`` -- the failure a dropped or misnumbered + declaration produces. Read at the one path these plans put their references rather + than by the generic walk ``test_collector`` needs: there the point is that no + reference field anywhere is missed, here the plan is built two lines above and its + shape is known. + """ + declarations = _declarations(plan) + return [ + declarations[expression.scalar_function.function_reference] + for expression in plan.relations[-1].root.input.project.expressions + ] + + +class TestPyarrowExpressions: + """pyarrow numbers densely from 0 and emits a bare + ``extension_function { name: "add" }`` naming no URN -- the producer that motivated + all of the anchor-0 handling. + """ + + def test_one_expression_folds_in(self, full_registry, serialize): + expression = serialize(lambda pc: pc.field("a") + pc.field("b")) + # The premise of everything below: anchor 0, no URN, reference 0. + assert _declarations(expression) == {0: "add"} + assert not expression.extension_urns + assert [ + e.expression.scalar_function.function_reference + for e in expression.referred_expr + ] == [0] + + out = project( + read_named_table("t", expression.base_schema), expressions=[expression] + )(full_registry) + + assert _declarations(out) == {1: "add"} + assert _project_functions(out) == ["add"] + + def test_two_expressions_get_their_own_declarations(self, full_registry, serialize): + """Each serialized expression is its own anchor space, so two of them both + declare their function at 0 and both reference 0. Folded into one plan they + used to stay there: two declarations at anchor 0, with every reference to it + resolving to whichever one the consumer found first -- ``a * b`` silently + readable as ``a + b``. + """ + add = serialize(lambda pc: pc.field("a") + pc.field("b")) + multiply = serialize(lambda pc: pc.field("a") * pc.field("b")) + assert _declarations(add) == {0: "add"} + assert _declarations(multiply) == {0: "multiply"} + + out = project( + read_named_table("t", add.base_schema), expressions=[add, multiply] + )(full_registry) + + assert _declarations(out) == {1: "add", 2: "multiply"} + assert _project_functions(out) == ["add", "multiply"] + + def test_several_expressions_in_one_carrier_fold_in(self, full_registry, serialize): + """Serialized together, pyarrow numbers them densely from 0 in one carrier -- + so anchor 0 arrives alongside anchors it must not be renumbered on top of. + """ + carrier = serialize( + lambda pc: pc.field("a") + pc.field("b"), + lambda pc: pc.field("a") * pc.field("b"), + lambda pc: pc.field("a") - pc.field("b"), + ) + assert sorted(_declarations(carrier)) == [0, 1, 2] + + out = project( + read_named_table("t", carrier.base_schema), expressions=[carrier] + )(full_registry) + + declarations = _declarations(out) + assert sorted(declarations) == [1, 2, 3] + assert set(declarations.values()) == {"add", "multiply", "subtract"} + assert _project_functions(out) == ["add", "multiply", "subtract"] + + def test_mixed_with_functions_this_build_resolves(self, full_registry, serialize): + """The two numbering spaces meet: pyarrow's ``add`` (no URN) and this + library's ``add:i64_i64`` (naming ARITHMETIC) are different identities that + must not be conflated, and no reference may end up naming the other's. + """ + expression = serialize(lambda pc: pc.field("a") + pc.field("b")) + out = project( + read_named_table("t", expression.base_schema), + expressions=[ + scalar_function(ARITHMETIC, "add", [column("a"), column("b")]), + expression, + ], + )(full_registry) + + assert _declarations(out) == {1: "add:i64_i64", 2: "add"} + assert _project_functions(out) == ["add:i64_i64", "add"] + # Only the resolved one names a URN; the URN table stays 1-based so the + # other's unset reference matches nothing. + assert [(u.extension_urn_anchor, u.urn) for u in out.extension_urns] == [ + (1, ARITHMETIC) + ] + assert out.extensions[1].extension_function.extension_urn_reference == 0 diff --git a/tests/sql/test_sql_to_substrait.py b/tests/integration/test_sql_engine_roundtrip.py similarity index 77% rename from tests/sql/test_sql_to_substrait.py rename to tests/integration/test_sql_engine_roundtrip.py index 99917c06..8edfc32e 100644 --- a/tests/sql/test_sql_to_substrait.py +++ b/tests/integration/test_sql_engine_roundtrip.py @@ -1,4 +1,17 @@ -import os +"""SQL -> Substrait -> engine round-trips, checked against the engine's own SQL. + +Behavioral tests through external Substrait consumers (DuckDB / DataFusion, both +reading plans this library builds). Those consumers lag the spec, and feeding one a +plan built at a newer spec version can crash the interpreter natively -- not a +catchable failure, so a red run here is not even reliably a report. They are +best-effort and never a gate, which is why the ``duckdb`` and ``datafusion`` markers +are deselected by the ``addopts`` in pyproject.toml -- the two integrations that cannot +report their own failure, rather than integration testing as a category (the pyarrow +tests next door run by default). Run these deliberately, once the engines catch up, +with ``pytest -m integration`` (or one engine at a time via ``-m duckdb`` / ``-m +datafusion``). +""" + import sys import tempfile @@ -9,16 +22,7 @@ from substrait.extension_registry import ExtensionRegistry from substrait.sql.sql_to_substrait import convert -# These are behavioral round-trips through external Substrait consumers -# (pyarrow / DuckDB / DataFusion). Those consumers lag the spec, and feeding -# them a plan built at a newer spec version can crash the interpreter natively -# (not a catchable failure). They are best-effort and never a gate: skipped by -# default, opt in with SUBSTRAIT_ENGINE_TESTS=1 once the engines catch up. -pytestmark = pytest.mark.skipif( - not os.environ.get("SUBSTRAIT_ENGINE_TESTS"), - reason="engine Substrait consumers lag the pinned spec; " - "set SUBSTRAIT_ENGINE_TESTS=1 to run", -) +pytestmark = pytest.mark.integration data: pyarrow.Table = pyarrow.Table.from_batches( [ @@ -128,15 +132,30 @@ def assert_query(query: str, engine: str, ignore_order=True): assert_query_datafusion(query, ignore_order) +# Not a marker: the duckdb substrait extension is unavailable on windows at all, so +# even an explicit ``-m duckdb`` has nothing to run there. +NO_DUCKDB_SUBSTRAIT_ON_WINDOWS = pytest.mark.skipif( + sys.platform.startswith("win"), + reason="duckdb substrait extension not found on windows", +) + +# One marker per engine so each can be selected or excluded on its own -- the engines +# catch up to the spec independently, so "run everything but duckdb" has to be sayable +# (``-m "integration and not duckdb"``). engines = [ + pytest.param("duckdb", marks=[pytest.mark.duckdb, NO_DUCKDB_SUBSTRAIT_ON_WINDOWS]), + pytest.param("datafusion", marks=pytest.mark.datafusion), +] + +# The same two engines where duckdb does not yet consume what we emit. Kept a separate +# list rather than an inline parametrize per test so the per-engine markers cannot +# drift out of sync with ``engines``. +engines_duckdb_xfail = [ pytest.param( "duckdb", - marks=pytest.mark.skipif( - sys.platform.startswith("win"), - reason="duckdb substrait extension not found on windows", - ), + marks=[pytest.mark.duckdb, NO_DUCKDB_SUBSTRAIT_ON_WINDOWS, pytest.mark.xfail], ), - "datafusion", + pytest.param("datafusion", marks=pytest.mark.datafusion), ] @@ -264,64 +283,19 @@ def test_order_by(engine: str): ) -@pytest.mark.parametrize( - "engine", - [ - pytest.param( - "duckdb", - marks=[ - pytest.mark.skipif( - sys.platform.startswith("win"), - reason="duckdb substrait extension not found on windows", - ), - pytest.mark.xfail, - ], - ), - "datafusion", - ], -) +@pytest.mark.parametrize("engine", engines_duckdb_xfail) def test_select_limit(engine: str): assert_query("""SELECT store_id FROM stores ORDER BY store_id LIMIT 2""", engine) -@pytest.mark.parametrize( - "engine", - [ - pytest.param( - "duckdb", - marks=[ - pytest.mark.skipif( - sys.platform.startswith("win"), - reason="duckdb substrait extension not found on windows", - ), - pytest.mark.xfail, - ], - ), - "datafusion", - ], -) +@pytest.mark.parametrize("engine", engines_duckdb_xfail) def test_select_limit_offset(engine: str): assert_query( """SELECT store_id FROM stores ORDER BY store_id LIMIT 2 OFFSET 2""", engine ) -@pytest.mark.parametrize( - "engine", - [ - pytest.param( - "duckdb", - marks=[ - pytest.mark.skipif( - sys.platform.startswith("win"), - reason="duckdb substrait extension not found on windows", - ), - pytest.mark.xfail, - ], - ), - "datafusion", - ], -) +@pytest.mark.parametrize("engine", engines_duckdb_xfail) def test_row_number(engine: str): assert_query( """SELECT sale_id, fk_store_id, row_number() over (partition by fk_store_id order by sale_id) as rn 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..f56e21bf 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,279 @@ 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_rewrites_a_reference_of_zero(): + """0 is a valid anchor/reference -- the protos have said so on ``function_anchor`` + since Substrait v0.83.0 ("0 is a valid anchor/reference, but prefer non-zero + values for ergonomics") -- so an incoming plan may name its + first function with ``function_reference: 0``. + + proto3 leaves such a field out of ``ListFields()``, so a set-fields walk cannot + see it, let alone rewrite it: the reference would silently survive into a plan + whose numbering puts something else at 0. All four fields that hold a reference + without presence must be rewritten. + """ + remap = {0: 5} + rel = stalg.Rel( + project=stalg.ProjectRel( + expressions=[ + stalg.Expression( + scalar_function=stalg.Expression.ScalarFunction( + function_reference=0 + ) + ), + stalg.Expression( + window_function=stalg.Expression.WindowFunction( + function_reference=0 + ) + ), + ] + ) + ) + aggregate = stalg.Rel( + aggregate=stalg.AggregateRel( + measures=[ + stalg.AggregateRel.Measure( + measure=stalg.AggregateFunction(function_reference=0) + ) + ] + ) + ) + window = stalg.Rel( + window=stalg.ConsistentPartitionWindowRel( + window_functions=[ + stalg.ConsistentPartitionWindowRel.WindowRelFunction( + function_reference=0 + ) + ] + ) + ) + + out = remap_function_references(rel, remap) + assert out.project.expressions[0].scalar_function.function_reference == 5 + assert out.project.expressions[1].window_function.function_reference == 5 + assert ( + remap_function_references(aggregate, remap) + .aggregate.measures[0] + .measure.function_reference + == 5 + ) + assert ( + remap_function_references(window, remap) + .window.window_functions[0] + .function_reference + == 5 + ) + + +def test_remap_function_references_rewrites_zero_in_the_presence_bearing_fields(): + """The two reference fields that *do* have presence still carry 0 when set to + it, so selecting one must be enough to have it rewritten.""" + remap = {0: 5} + sort = stalg.SortField(comparison_function_reference=0) + key = stalg.ComparisonJoinKey( + comparison=stalg.ComparisonJoinKey.ComparisonType(custom_function_reference=0) + ) + + assert remap_function_references(sort, remap).comparison_function_reference == 5 + assert ( + remap_function_references(key, remap).comparison.custom_function_reference == 5 + ) + + +@pytest.mark.parametrize( + "direction", + [ + stalg.SortField.SORT_DIRECTION_ASC_NULLS_FIRST, + # The trap: a zero-valued oneof member is as invisible to ListFields() as a + # reference of 0, so a walk that gates on "is this field set" the wrong way + # cannot tell the two apart. + stalg.SortField.SORT_DIRECTION_UNSPECIFIED, + ], + ids=["asc", "unspecified"], +) +def test_remap_function_references_does_not_invent_a_sort_comparison(direction): + """``comparison_function_reference`` shares oneof ``sort_kind`` with + ``direction``, so a remap that knows the key 0 must not write it blind: a field + sorted by direction would come out sorted by a comparison function instead. + """ + sort = stalg.SortField( + expr=stalg.Expression(literal=stalg.Expression.Literal(i64=1)), + direction=direction, + ) + before = sort.SerializeToString(deterministic=True) + + out = remap_function_references(sort, {0: 5}) + + assert out.WhichOneof("sort_kind") == "direction" + assert out.SerializeToString(deterministic=True) == before + + +def test_remap_function_references_does_not_invent_a_sort_kind(): + """A SortField with no ``sort_kind`` at all keeps none: the oneof gate reads + which member is selected, not whether the field could hold a reference.""" + sort = stalg.SortField( + expr=stalg.Expression(literal=stalg.Expression.Literal(i64=1)) + ) + out = remap_function_references(sort, {0: 5}) + assert out.WhichOneof("sort_kind") is None + + +@pytest.mark.parametrize( + "simple", + [ + stalg.ComparisonJoinKey.SIMPLE_COMPARISON_TYPE_EQ, + stalg.ComparisonJoinKey.SIMPLE_COMPARISON_TYPE_UNSPECIFIED, + ], + ids=["eq", "unspecified"], +) +def test_remap_function_references_does_not_invent_a_join_key_comparison(simple): + """``custom_function_reference`` shares oneof ``inner_type`` with ``simple``, so + the same gate applies: an equi-join key must not come out comparing by function. + """ + key = stalg.ComparisonJoinKey( + comparison=stalg.ComparisonJoinKey.ComparisonType(simple=simple) + ) + before = key.SerializeToString(deterministic=True) + + out = remap_function_references(key, {0: 5}) + + assert out.comparison.WhichOneof("inner_type") == "simple" + assert out.SerializeToString(deterministic=True) == before + + +def test_remap_function_references_does_not_materialize_an_unset_function(): + """Only *set* submessages are descended into. An ``AggregateRel.Measure`` with no + ``measure`` holds no reference to rewrite, and descending anyway would bring the + ``AggregateFunction`` into existence -- inventing a measure calling function 5 -- + because its ``function_reference`` reads back as the 0 the remap knows. + """ + rel = stalg.Rel( + aggregate=stalg.AggregateRel( + measures=[ + stalg.AggregateRel.Measure( + filter=stalg.Expression( + literal=stalg.Expression.Literal(boolean=True) + ) + ) + ] + ) + ) + before = rel.SerializeToString(deterministic=True) + + out = remap_function_references(rel, {0: 5}) + + assert not out.aggregate.measures[0].HasField("measure") + assert out.SerializeToString(deterministic=True) == before + + +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.