feat!: assign extension anchors per plan, not per registry

[nielspardon]

Removes the extension-merging half of #207 as well (the schema re-inference half is untouched).

Problem

ExtensionRegistry hands out function_anchor / extension_urn_anchor values at registration time, and builders stamp those registry-global numbers into plans. Anchors are plan-local in Substrait, so this 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 — the value encoded how many functions the other default YAMLs defined and the order the functions*.yaml glob happened to return them (filesystem order, not sorted). A substrait-extensions bump, or a different machine, shifted every anchor.

Extending a plan built elsewhere silently corrupted it. merge_extension_urns / merge_extension_declarations dedupe by identity and their docstrings state "Assumes that there are no collisions", with nothing enforcing it. Given a plan already using anchor 10/284 for different entities:

extension_urn_anchor: 10  urn: "extension:acme:custom"
extension_urn_anchor: 10  urn: "extension:io.substrait:functions_arithmetic"
function_anchor: 284  name: "acme_thing:i64"   extension_urn_reference: 10
function_anchor: 284  name: "add:i64_i64"      extension_urn_reference: 10

function_reference: 284 is now ambiguous, and no error is raised.

Approach

ExtensionCollector owns those anchors for the duration of one build. Function references are allocated on first use from 1; URN anchors are derived at emit time, since nothing outside SimpleExtensionDeclaration refers to one. This follows substrait-java's io.substrait.extension.ExtensionCollector, including its first-use numbering and its deferral of URN anchors.

The collector reaches builders through a contextvar, as the builders' other per-build state already does (_rel_anchor_counter, outer_schemas, anchor_scope). Four seams carry the change: build_scoped wrapping each resolver, plus _bind (plans), resolve_expression (expressions) and _inner_rel (subqueries) for adopting inputs. _inner_rel matters because only a bare Rel crosses into Expression.Subquery — a pre-built plan's declarations would otherwise stay behind with the discarded plan and leave its references dangling.

An incoming materialized plan has its declarations read back to (urn, name) identities and its references re-derived rather than trusted, so two independently numbered inputs cannot disagree about what a reference means. This is load-bearing rather than defensive: the SQL translator builds a set operation's two sides as separate plans before merging them, and both number from 1.

Identities come off the declaration, not a catalog lookup, so a plan naming functions absent from the registry still round-trips. An input that declares two different functions at one anchor is refused rather than silently resolved to one of them — that ambiguity is the defect this PR exists to remove, so producing it from a merge and accepting it from an input should not have different answers.

Anchor 0

Anchor 0 is re-derived like any other anchor. The spec marks it a valid anchor/reference (substrait#900, carried by the pinned v0.99.0 protos), and pyarrow's serialize_expressions numbers from 0 — a two-expression serialization emits anchors 0 and 1, the first as a bare extension_function { name: "add" }.

Rewriting a reference of 0 needs remap_function_references to read reference fields off the descriptor rather than ListFields(), which omits default-valued proto3 scalars and so made function_reference: 0 invisible. Presence rules make that exact:

field	presence	handling
ScalarFunction / WindowFunction / AggregateFunction / WindowRelFunction .function_reference	none	remapped unconditionally — a present containing message always implies a real reference
SortField.comparison_function_reference (oneof sort_kind)	yes	gated on WhichOneof, so a SortField using direction never acquires one
ComparisonJoinKey.ComparisonType.custom_function_reference (oneof inner_type)	yes	same

Emission stays 1-based, as those same protos ask producers to prefer non-zero values for ergonomics. This does not extend to type_variation_anchor, where 0 remains reserved for the system-preferred variation.

Treating 0 as ordinary also fixes a case that was silently wrong: two pyarrow expressions folded into one build both emitted at function_anchor: 0, so each reference resolved ambiguously and multiply was readable as add.

Two related corrections

aggregate now refuses a measure that is not an aggregate function. It previously emitted such a measure as a measure that is set but empty — malformed on its own, and the one shape where a present message does not imply a real function reference, so the only shape reference renumbering could not treat correctly.

alias renames on a copy. It returned the message resolve_expression handed it, which is the caller's own object whenever the remap is empty, so it mutated a caller's ExtendedExpression.

ExtensionRegistry becomes a pure catalog. Its urn→function mapping, signature matching and extension-relation registration are unchanged.

Result

Anchors are now dense and plan-local:

extension_urn_anchor: 1  urn: "extension:io.substrait:functions_arithmetic"
function_anchor: 1       name: "add:i64_i64"

Byte-identical across a minimal vs. the full default registry, across repeated builds of the same frame, and under PYTHONHASHSEED 0/1/7/12345.

Because the collector accumulates once per build, per-level extension merging is gone rather than optimized. Measured against a main worktree over an N-verb project chain:

N	main (merge calls / declarations scanned)	this branch
5	10 / 20	0 / 0
10	20 / 50	0 / 0
20	40 / 110	0 / 0
40	80 / 230	0 / 0

Wall-clock build time is essentially unchanged (59 → 61 ms at N=40): schema re-inference still dominates, which is #207's other half and out of scope here.

Examples are not covered by CI, so they were diffed against a main worktree: pyarrow_example and duckdb_example print unchanged output (duckdb executes the re-anchored plan to the same result set), though pyarrow_example's plan moves its declaration off anchor 0 and duckdb_example's moves 86 → 1. builder_example and dataframe_example differ only in anchor values (89/284/479 → 1–4); canonicalizing every reference by the (urn, name) it resolves to makes all of their emitted plans identical to main, with no dangling references.

Test layout

The pyarrow tests here 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 (which is what CI runs). 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. Note -m replaces the default expression rather than narrowing it -- -m "not pyarrow" would re-enable the engines -- which CONTRIBUTING.md now spells out along with the selectors.

BREAKING CHANGE: emitted extension anchors are 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 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

🤖 Generated with AI