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

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 = (

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",

src/substrait/extension_registry/collector.py

@@ -0,0 +1,221 @@
+"""Per-build collection of plan-local extension anchors.
+
+Extension anchors (``function_anchor``, ``extension_urn_anchor``) are *plan-local*
+in Substrait: they are an artifact of serializing a plan, not durable identifiers.
+The durable identity of a function is its ``(urn, name)`` pair. This module owns
+the mapping between the two for the duration of a single build, which is what
+keeps :class:`~substrait.extension_registry.ExtensionRegistry` free to be a pure
+catalog.
+
+Mirrors ``ExtensionCollector`` in substrait-java (
+``io.substrait.extension.ExtensionCollector``), including its first-use reference
+numbering and its deferral of URN anchors to emit time.
+"""
+
+import contextlib
+import contextvars
+import functools
+import itertools
+from typing import Optional, Union
+
+import substrait.extended_expression_pb2 as stee
+import substrait.extensions.extensions_pb2 as ste
+import substrait.plan_pb2 as stplan
+
+# Identity of a function as declared in a plan: (extension URN, function name).
+# The name is the compound form carried by SimpleExtensionDeclaration (e.g.
+# "add:i64_i64"), which is what makes an identity resolvable without the catalog.
+# The URN is None for a declaration that names no resolvable extension URN.
+FunctionIdentity = tuple[Optional[str], str]
+
+ExtensionCarrier = Union[stplan.Plan, stee.ExtendedExpression]
+
+
+def _is_underspecified(function, urns_by_anchor: dict) -> bool:
+    """Whether a declaration carries too little to be re-derived, so must be kept as is.
+
+    Some producers emit a bare ``extension_function { name: "add" }`` -- no anchor
+    and no URN (pyarrow's ``serialize_expressions`` does exactly this). There is
+    nothing to renumber: anchor 0 is not a reference any expression can name under
+    the spec's 1-based numbering, so re-assigning it would desynchronize the
+    declaration from expressions that still say 0. Such declarations are preserved
+    verbatim instead.
+    """
+    return (
+        function.function_anchor == 0
+        and function.extension_urn_reference not in urns_by_anchor
+    )
+
+
+class ExtensionCollector:
+    """Owns the plan-local extension anchors for a single build.
+
+    Function references are allocated on first use, numbered from 1 in the order
+    they are first referenced. URN anchors are *not* allocated during the build:
+    nothing outside ``SimpleExtensionDeclaration`` refers to one, so they are
+    derived in :meth:`write_into` from the order the declarations were collected.
+    """
+
+    def __init__(self) -> None:
+        self._references: dict[FunctionIdentity, int] = {}
+        self._reference_generator = itertools.count(1)
+        # Declarations kept verbatim because they carry neither a URN nor an anchor
+        # to re-derive from; see _is_underspecified. Deduplicated on serialized bytes.
+        self._passthrough: "list[ste.SimpleExtensionDeclaration]" = []
+        self._passthrough_keys: set = set()
+
+    def function_reference(self, urn: str, name: str) -> int:
+        """The reference for ``(urn, name)``, allocating one on first use."""
+        identity = (urn, name)
+        reference = self._references.get(identity)
+        if reference is None:
+            reference = next(self._reference_generator)
+            self._references[identity] = reference
+        return reference
+
+    def adopt(self, carrier: ExtensionCarrier) -> dict[int, int]:
+        """Take over the extension declarations of an incoming plan or expression.
+
+        Reads ``carrier``'s declarations back to ``(urn, name)`` identities and
+        allocates this build's reference for each, returning the
+        ``{old reference: new reference}`` remap its relations/expressions need.
+        Pass the result to
+        :func:`substrait.utils.remap_function_references`; it is empty (identity)
+        whenever the incoming numbering already agrees with ours, which is the
+        common case.
+
+        Every incoming reference is re-derived rather than trusted, so two
+        independently built inputs meeting at a multi-input relation cannot
+        disagree about what a reference number means. Identities come from the
+        declaration itself, so a function absent from the catalog is carried
+        through unchanged rather than rejected.
+        """
+        urns_by_anchor = {u.extension_urn_anchor: u.urn for u in carrier.extension_urns}
+        remap = {}
+        for declaration in carrier.extensions:
+            if declaration.WhichOneof("mapping_type") != "extension_function":
+                # Type / type-variation declarations are not collected yet; see
+                # merge_extension_declarations for the same gap.
+                mapping_type = declaration.WhichOneof("mapping_type")
+                raise NotImplementedError(
+                    f"cannot collect extension declaration of type {mapping_type!r}; "
+                    f"only 'extension_function' declarations are supported so far"
+                )
+            function = declaration.extension_function
+            if _is_underspecified(function, urns_by_anchor):
+                key = declaration.SerializeToString(deterministic=True)
+                if key not in self._passthrough_keys:
+                    self._passthrough_keys.add(key)
+                    self._passthrough.append(declaration)
+                continue
+            # An unresolvable URN reference still yields a usable identity: the
+            # function name alone. Anchors stay collision-free either way.
+            urn = urns_by_anchor.get(function.extension_urn_reference)
+            new = self.function_reference(urn, function.name)
+            if new != function.function_anchor:
+                remap[function.function_anchor] = new
+        return remap
+
+    def write_into(self, carrier: ExtensionCarrier) -> None:
+        """Emit the collected extensions onto ``carrier``, replacing what is there.
+
+        URN anchors are assigned here, numbered from 1 in the order each URN was
+        first referenced, so a plan's URN anchors are as dense and plan-local as
+        its function references.
+        """
+        urn_anchors: dict[str, int] = {}
+        urn_anchor_generator = itertools.count(1)
+        # Pass-through declarations first: they keep whatever anchor they arrived
+        # with, and cannot clash with the 1-based numbering below.
+        declarations = list(self._passthrough)
+        for (urn, name), reference in self._references.items():
+            function = ste.SimpleExtensionDeclaration.ExtensionFunction(
+                function_anchor=reference, name=name
+            )
+            if urn is not None:
+                urn_anchor = urn_anchors.get(urn)
+                if urn_anchor is None:
+                    urn_anchor = next(urn_anchor_generator)
+                    urn_anchors[urn] = urn_anchor
+                function.extension_urn_reference = urn_anchor
+            declarations.append(
+                ste.SimpleExtensionDeclaration(extension_function=function)
+            )
+
+        carrier.ClearField("extension_urns")
+        carrier.extension_urns.extend(
+            ste.SimpleExtensionURN(extension_urn_anchor=anchor, urn=urn)
+            for urn, anchor in urn_anchors.items()
+        )
+        carrier.ClearField("extensions")
+        carrier.extensions.extend(declarations)
+
+
+# The ExtensionCollector for the build currently in progress, or None outside a
+# build. Ambient rather than threaded through resolver signatures, following the
+# other per-build state the builders already carry this way (_rel_anchor_counter
+# in builders.extended_expression, outer_schemas / anchor_scope in type_inference).
+_collector: contextvars.ContextVar = contextvars.ContextVar("_collector", default=None)
+
+
+def current_collector() -> Optional[ExtensionCollector]:
+    """The collector for the build in progress, or None outside a build."""
+    return _collector.get()
+
+
+def function_reference(urn: str, name: str) -> int:
+    """This build's reference for the function ``(urn, name)``.
+
+    Convenience over ``current_collector().function_reference(...)`` for builders,
+    which always run inside a scope.
+    """
+    collector = _collector.get()
+    if collector is None:
+        raise RuntimeError(
+            "no build in progress: extension anchors are plan-local, so a builder "
+            "must resolve inside a build_scope() (builders are wrapped in "
+            "build_scoped(), which enters one)"
+        )
+    return collector.function_reference(urn, name)
+
+
+@contextlib.contextmanager
+def build_scope():
+    """Enter the current build, creating a collector if this is the outermost one.
+
+    Yields ``(collector, owns_scope)``. ``owns_scope`` is True only for the
+    outermost resolver of a build -- the one responsible for writing the collected
+    extensions onto its output. Nested resolvers get the same collector and write
+    nothing, which is what lets a build accumulate extensions once instead of
+    re-merging them at every level.
+    """
+    collector = _collector.get()
+    if collector is not None:
+        yield collector, False
+        return
+    collector = ExtensionCollector()
+    token = _collector.set(collector)
+    try:
+        yield collector, True
+    finally:
+        _collector.reset(token)
+
+
+def build_scoped(resolve):
+    """Wrap a builder's resolver so it participates in a build scope.
+
+    The outermost resolver of a build writes the collected extension declarations
+    onto whatever it returns; nested ones leave those fields empty for it to fill.
+    Signature-agnostic, since plan resolvers take ``(registry)`` and expression
+    resolvers take ``(base_schema, registry)``.
+    """
+
+    @functools.wraps(resolve)
+    def wrapper(*args, **kwargs):
+        with build_scope() as (collector, owns_scope):
+            resolved = resolve(*args, **kwargs)
+            if owns_scope:
+                collector.write_into(resolved)
+            return resolved
+
+    return wrapper

src/substrait/extension_registry/function_entry.py

@@ -1,5 +1,6 @@
 """Function entry class for extension registry."""
 
+import warnings
 from enum import Enum
 from typing import Optional, Union
 
@@ -23,14 +24,12 @@ def __init__(
         urn: str,
         name: str,
         impl: Union[se.Impl, se.Impl1, se.Impl2],
-        anchor: int,
         function_type: FunctionType = FunctionType.SCALAR,
     ) -> None:
         self.name = name
         self.impl = impl
         self.normalized_inputs: list = []
         self.urn: str = urn
-        self.anchor = anchor
         self.function_type = function_type
         self.arguments = []
         self.nullability = (
@@ -50,6 +49,23 @@ def __init__(
     def __repr__(self) -> str:
         return f"{self.name}:{'_'.join(self.normalized_inputs)}"
 
+    @property
+    def anchor(self) -> int:
+        """Deprecated. Function anchors are plan-local, not catalog state.
+
+        A function's durable identity is ``(urn, str(entry))``; the numeric anchor
+        it gets in a plan is assigned per build by
+        :class:`~substrait.extension_registry.ExtensionCollector`.
+        """
+        warnings.warn(
+            "FunctionEntry.anchor is deprecated: function anchors are plan-local and "
+            "are assigned per build by ExtensionCollector. Identify a function by "
+            "(entry.urn, str(entry)) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return 1
+
     def satisfies_signature(self, signature: tuple | list) -> Optional[str]:
         if self.impl.variadic:
             min_args_allowed = self.impl.variadic.min or 0