refactor: scope extension-relation derivers to the ExtensionRegistry#220
Merged
nielspardon merged 2 commits intoJul 16, 2026
Conversation
User-defined extension relations registered their schema deriver in a process-global dict in type_inference, while ExtensionRegistry also kept a per-instance dict that inference never read (dead state). As a result two registries silently shared one deriver map, registrations were last-writer-wins across type_urls, and state leaked across otherwise independent registries and across tests. Store derivers only on the ExtensionRegistry and thread that registry through inference so _derive_extension_schema consults it, then drop the module-global dict and the module-level register_extension_relation shim. - ExtensionRegistry.register_extension_relation now writes only self._extension_relations; adds lookup_extension_relation(type_url). - infer_plan_schema / infer_rel_schema / infer_expression_type / infer_nested_type / infer_extended_expression_schema / _join_output_struct take a keyword-only registry=None and forward it to every recursive call and to _derive_extension_schema. When registry is None the behavior is unchanged (leaf/multi raise, single passes through). - Builder / DataFrame / expr call sites pass registry=registry (each is already inside a resolve(registry) closure). - test_extension_relations: _kinds threads the registry, plus a new test asserting a deriver registered on one registry does not derive on an independent one. Closes substrait-io#206.
nielspardon
force-pushed
the
refactor/scope-extension-relation-derivers
branch
from
July 16, 2026 13:13
6018b21 to
d95813b
Compare
tokoko
pushed a commit
that referenced
this pull request
Jul 16, 2026
## Problem `type_inference.infer_rel_schema`'s `set` branch derived a `SetRel`'s output schema from **`rel.set.inputs[0]` only** and never read `rel.set.op`. Per the Substrait spec, set inputs share field *types* but may differ in *nullability*, and the output nullability must be combined across all inputs according to the operation. So for `UNION` / `INTERSECTION`, the inferred nullability could be wrong (too strict) when inputs disagree. Reachable through this library's own builders: `plan.set(...)` and the DataFrame `union` / `intersect` / `except_` verbs build a `SetRel` with no stored schema, so a later `infer_plan_schema` over a union of two sub-plans whose shared columns differ in nullability under-reported nullability. ## Fix Infer every input and combine each field's nullability per the operation, matching the spec's [set-operation output-type derivation](https://substrait.io/relations/logical_relations/#set-operation): - **UNION_DISTINCT / UNION_ALL** — nullable if the field is nullable in any input. - **INTERSECTION_PRIMARY** — nullable only when the field is nullable in the primary and in at least one secondary; required otherwise. - **INTERSECTION_MULTISET / _ALL** — required if the field is required in any input. - **MINUS_PRIMARY / _PRIMARY_ALL / _MULTISET** (and unspecified) — same as the primary. Field *types* are taken from the primary input (the spec requires identical field types across inputs); only the top-level nullability is recombined, and `CopyFrom` preserves type parameters and nested element types. ## Testing - `test_inference_set_nullability` — parametrized over all 8 concrete set ops using the **spec's own worked example** (three inputs, one column per required/nullable combination) and asserting the spec's exact output columns. - `test_inference_set_nullability_preserves_field_types` — a `UNION_ALL` over `decimal(10,2)` / `varchar(5)` / `list<string>` columns, asserting the full `Type.Struct` so type parameters, nested element types, and struct-level nullability are all verified. - Full suite: **521 passed, 30 skipped**; `ruff format --check` and `ruff check` clean. An adversarial cross-check (independent re-derivation of all 9 `SetOp` cases against the spec, plus edge-case and test-completeness passes) confirmed the rules; a defensive guard was added so a field `Type` with no `kind` set passes through unchanged rather than raising. ## Relationship to #206 / #220 Independent. [#220](#220) (registry scoping) also touches the `set` branch line to thread `registry=registry`; whichever merges second is a trivial rebase (keep both changes). Closes #219. 🤖 Generated with AI
…nsion-relation-derivers # Conflicts: # src/substrait/type_inference.py
nielspardon
enabled auto-merge (squash)
July 16, 2026 15:14
nielspardon
disabled auto-merge
July 16, 2026 15:14
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
User-defined extension relations registered their schema deriver in a process-global dict in
type_inference(_extension_relation_derivers), whileExtensionRegistry.register_extension_relationalso wrote a per-instance dict (self._extension_relations) that inference never read — dead state.Consequences (from #206):
ExtensionRegistryinstances silently shared one deriver map.type_url.Every other registry capability is instance-scoped; extension-relation derivation should be too.
Change
Store derivers only on the
ExtensionRegistryand thread that registry through inference so_derive_extension_schemaconsults it, then drop the module-global dict and the module-levelregister_extension_relationshim.ExtensionRegistry.register_extension_relationnow writes onlyself._extension_relations; addslookup_extension_relation(type_url).infer_plan_schema/infer_rel_schema/infer_expression_type/infer_nested_type/infer_extended_expression_schema/_join_output_structtake a keyword-onlyregistry=Noneand forward it down to_derive_extension_schema. Whenregistry is Nonethe behavior is unchanged (leaf/multi raise "no schema deriver"; single passes through), so the many non-extension call sites are unaffected.exprcall sites passregistry=registry— each is already inside aresolve(registry)closure, so chained extension-relation inference resolves against the right registry.Behavior change
Registration is now scoped to the
ExtensionRegistryinstance. Code following the documented pattern (register on a registry, then build/infer through that same registry — as theDataFrameand builders do) is unaffected. Two independent registries no longer share deriver state.Testing
test_deriver_registration_is_scoped_to_registry: a detail class registered on one registry does not derive on an independent one (and still derives on the one it was registered on)._kindshelper now threads the registry (it re-infers a finished plan whosedetailis an opaqueAny).ruff format --checkandruff checkclean.Closes #206.
🤖 Generated with AI