Customizing extension support across relations, expressions and types

[nielspardon]

Context

While making extension anchors plan-local (#245), every whole-tree rewrite in the library had to be taught to reconcile an incoming plan's numbering. That surfaced a boundary the current extension machinery does not cross: a google.protobuf.Any payload is opaque to those rewrites, so anything a user packs into one is silently exempt from them. The anchor work itself is unaffected (this library never emits a reference inside an Any), but a plan built elsewhere may, and the same boundary limits every future tree rewrite.

Worth looking at as one question rather than patching per-rewrite, and worth covering relations, expressions and types, since support is currently uneven across the three.

The concrete gap

substrait.utils.remap_function_references walks the protobuf descriptor. It descends into an Any wrapper, finds only type_url and opaque value bytes, and moves on — so a function reference inside a packed payload keeps the incoming plan's numbering after being folded into a build. Rewriting it means parsing it, which needs the schema Any withholds.

Every Any-typed field in the protos, i.e. every place this applies:

field
ExtensionLeafRel.detail, ExtensionSingleRel.detail, ExtensionMultiRel.detail
ExtensionObject.detail
ReadRel.ExtensionTable.detail
ReadRel.LocalFiles.FileOrFiles.extension
ExchangeRel.ExchangeTarget.extended
Expression.Literal.UserDefined.value
AdvancedExtension.optimization, AdvancedExtension.enhancement

The same boundary applies to the library's other whole-tree rewrites, not just anchors: rebase_reference_ordinals, to_id_based_outer_references, and inline_reference_rels all stop at the wrapper too. Each is currently a hand-written recursive walk with its own traversal (_iter_child_rels / _child_rel_fields for relations, _iter_direct_subexpressions for expressions), so an extension author has no single place to participate even if they wanted to.

What exists today

Extension relations have a real extension point. A user implements one of the detail ABCs in substrait.dataframe.extension_relations — ExtensionLeafDetail, ExtensionSingleDetail, ExtensionMultiDetail — providing type_url, to_any(), from_any() and derive_schema(), and registers the class with ExtensionRegistry.register_extension_relation. Schema inference then reconstructs the detail from a serialized plan and calls derive_schema, so a custom relation infers like a built-in one.

Notably, to_any/from_any already make the payload round-trippable. The machinery to rewrite a packed payload is therefore mostly present; nothing calls it for that purpose.

Extension types have a builder (substrait.builders.type.user_defined) but no detail/registry equivalent, and a UserDefined literal's value is itself an opaque Any.

Extension expressions have nothing analogous at all.

What seems missing

• No way for an extension to participate in a tree rewrite. Anchor remapping is the current example; reference rebasing and outer-reference conversion have the same shape. An extension author cannot opt into any of them.
• No parity across relations / expressions / types. Relations have detail classes and inference support; types have a builder only; expressions have neither.
• A second plan-local anchor space is unhandled. Plan.type_aliases[].type_alias_anchor, referenced by Type.alias.type_alias_reference, is plan-local exactly like extension anchors. Relational builders drop it; with_execution_behavior preserves it via CopyFrom. So folding a foreign plan that uses type aliases can leave dangling type references. Verified pre-existing (identical before feat!: assign extension anchors per plan, not per registry #245) and this library never emits type aliases, so it is latent — but it is the same class of problem the extension-anchor work addressed for functions. Same for advanced_extensions, expected_type_urls and parameter_bindings, which builders also drop.
• Silent rather than loud. None of the above reports anything. A warning is awkward because most details presumably carry no references at all, so the useful signal probably has to come from the extension declaring what it holds.

Prior art: substrait-java

substrait-java solves the traversal half with customizable visitors, which is the piece substrait-python does ad hoc per rewrite:

• Relations — RelVisitor, AbstractRelVisitor, RelCopyOnWriteVisitor, ExpressionCopyOnWriteVisitor
• Expressions — ExpressionVisitor, AbstractExpressionVisitor, MaskExpressionVisitor
• Types — TypeVisitor, NamedFieldCountingTypeVisitor, StringTypeVisitor

RelCopyOnWriteVisitor is the closest analogue to what this library keeps rewriting by hand: a copy-on-write rewrite over the tree that a user can subclass. A shared traversal in substrait-python would give the anchor remapper, the ordinal rebaser and the outer-reference converter one place to hook, and give an extension one place to be asked.

Questions to settle

1. Is a shared, customizable traversal (visitor-shaped or otherwise) the right foundation, or is a narrower per-rewrite hook enough?
2. Should participation be opt-in on the detail classes (e.g. an optional method, defaulting to no-op so existing implementors keep working), or driven from the registry?
3. What is the equivalent extension point for expressions and for types, and should the three share one protocol?
4. substrait.utils is deliberately proto-only and imports no registry. Should a hook be a callback supplied by the call sites that already hold a registry (builders.plan._bind, builders.extended_expression.resolve_expression / _inner_rel), or should the dependency direction change?
5. Should type_aliases (and the other dropped plan-level fields) be handled as part of this, given it is the same plan-local-anchor problem?

🤖 Generated with AI