refactor: de-duplicate small helpers in the DataFrame/Expr API

src/substrait/builders/extended_expression.py

@@ -61,6 +61,27 @@ def resolve_expression(
     )
 
 
+def alias(
+    expression: ExtendedExpressionOrUnbound,
+    name: str,
+) -> UnboundExtendedExpression:
+    """Rename the first output column of ``expression`` to ``name``.
+
+    Returns a resolver that binds ``expression`` and rewrites its top-level output
+    name. Shared by the DataFrame/Expr ``.alias`` and the Narwhals expression
+    wrapper so the single-column rename lives in one place.
+    """
+
+    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
+
+
 _EPOCH_DATE = date(1970, 1, 1)
 
 

src/substrait/dataframe/expr.py

@@ -40,6 +40,9 @@
     singular_or_list,
     switch,
 )
+from substrait.builders.extended_expression import (
+    alias as _alias,
+)
 from substrait.builders.extended_expression import (
     dynamic_parameter as _dynamic_parameter,
 )
@@ -62,6 +65,7 @@
     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"
@@ -239,18 +243,21 @@ def _resolve_over_urns(
     raises a uniform error if none do. Shared by the operator path
     (:func:`_numeric_binary`) and the ``f.*`` namespace's multi-URN helper
     (``substrait.dataframe.functions._multi_urn_helper``) so both resolve
-    identically and their error text cannot drift apart.
+    identically and their error text cannot drift apart. The registry finds the
+    winning extension across every candidate URN in one call; ``entry.urn``
+    recovers it so ``builder`` can rebuild against the concrete overload.
     """
     signature = [
         typ
         for b in bound
         for typ in infer_extended_expression_schema(b, registry=registry).types
     ]
-    for urn in urns:
-        if registry.lookup_function(urn, name, signature):
-            return builder(urn, name, expressions=bound, alias=alias, options=options)(
-                base_schema, registry
-            )
+    match = registry.find_function(name, signature, urns)
+    if match is not None:
+        winning_urn = match[0].urn
+        return builder(
+            winning_urn, name, expressions=bound, alias=alias, options=options
+        )(base_schema, registry)
     kinds = [t.WhichOneof("kind") for t in signature]
     raise Exception(
         f"No matching overload for '{name}' across {urns} with signature {kinds}"
@@ -357,28 +364,19 @@ def _plan_of(query: Any):
     return plan
 
 
-def _sort_direction(descending: bool, nulls_last: bool):
-    if descending:
-        return (
-            stalg.SortField.SORT_DIRECTION_DESC_NULLS_LAST
-            if nulls_last
-            else stalg.SortField.SORT_DIRECTION_DESC_NULLS_FIRST
-        )
-    return (
-        stalg.SortField.SORT_DIRECTION_ASC_NULLS_LAST
-        if nulls_last
-        else stalg.SortField.SORT_DIRECTION_ASC_NULLS_FIRST
-    )
+# Sort direction keyed by (descending, nulls_last); the canonical mapping for the
+# DataFrame/Expr layer, shared with ``substrait.dataframe.frame``.
+_SORT_DIRECTIONS = {
+    (False, False): stalg.SortField.SORT_DIRECTION_ASC_NULLS_FIRST,
+    (False, True): stalg.SortField.SORT_DIRECTION_ASC_NULLS_LAST,
+    (True, False): stalg.SortField.SORT_DIRECTION_DESC_NULLS_FIRST,
+    (True, True): stalg.SortField.SORT_DIRECTION_DESC_NULLS_LAST,
+}
 
 
-def _merge_extensions_into(target, source):
-    """Append any extension URNs/declarations from ``source`` not already present."""
-    for urn in source.extension_urns:
-        if urn not in target.extension_urns:
-            target.extension_urns.append(urn)
-    for decl in source.extensions:
-        if decl not in target.extensions:
-            target.extensions.append(decl)
+def sort_direction(descending: bool, nulls_last: bool):
+    """The ``SortField.SortDirection`` for a ``(descending, nulls_last)`` pair."""
+    return _SORT_DIRECTIONS[(bool(descending), bool(nulls_last))]
 
 
 def _window_bound(value):
@@ -751,7 +749,7 @@ def order_by(
         ``keys`` are column names or expressions; ``descending``/``nulls_last``
         apply to all of them. Only meaningful on an aggregate measure.
         """
-        direction = _sort_direction(descending, nulls_last)
+        direction = sort_direction(descending, nulls_last)
         inner = self._unbound
 
         def resolve(base_schema, registry):
@@ -768,7 +766,7 @@ def resolve(base_schema, registry):
                     )
                 )
                 # Carry over any extensions a (function-valued) sort key introduced.
-                _merge_extensions_into(bound, bound_key)
+                merge_extensions_into(bound, bound_key)
             return bound
 
         return Expr(resolve)
@@ -799,7 +797,7 @@ def over(
             else list(partition_by)
         )
         order_keys = [order_by] if isinstance(order_by, (str, Expr)) else list(order_by)
-        direction = _sort_direction(descending, nulls_last)
+        direction = sort_direction(descending, nulls_last)
         inner = self._unbound
 
         def resolve(base_schema, registry):
@@ -812,7 +810,7 @@ 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)
+                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)
@@ -821,7 +819,7 @@ def resolve(base_schema, registry):
                         expr=bound_k.referred_expr[0].expression, direction=direction
                     )
                 )
-                _merge_extensions_into(bound, bound_k)
+                merge_extensions_into(bound, bound_k)
             frame = rows if rows is not None else range
             if frame is not None:
                 wf.bounds_type = (
@@ -860,14 +858,7 @@ def cast(self, type: Any) -> "Expr":
 
     def alias(self, name: str) -> "Expr":
         """Return a copy of this expression with its output name set to ``name``."""
-        inner = self._unbound
-
-        def resolve(base_schema, registry):
-            bound = inner(base_schema, registry)
-            bound.referred_expr[0].output_names[0] = name
-            return bound
-
-        return Expr(resolve)
+        return Expr(_alias(self._unbound, name))
 
     def __repr__(self) -> str:  # pragma: no cover - debugging aid
         return "Expr(<unbound>)"

src/substrait/dataframe/frame.py

@@ -36,7 +36,7 @@
 
 from substrait.builders import plan as _plan
 from substrait.builders import type as _type
-from substrait.dataframe.expr import Expr, Measure, col, lit
+from substrait.dataframe.expr import Expr, Measure, col, lit, sort_direction
 from substrait.extension_registry import ExtensionRegistry
 from substrait.type_inference import infer_plan_schema
 
@@ -74,14 +74,6 @@
     "update": stalg.WriteRel.WRITE_OP_UPDATE,
 }
 
-# Sort direction keyed by (descending, nulls_last).
-_SORT_DIRECTIONS = {
-    (False, False): stalg.SortField.SORT_DIRECTION_ASC_NULLS_FIRST,
-    (False, True): stalg.SortField.SORT_DIRECTION_ASC_NULLS_LAST,
-    (True, False): stalg.SortField.SORT_DIRECTION_DESC_NULLS_FIRST,
-    (True, True): stalg.SortField.SORT_DIRECTION_DESC_NULLS_LAST,
-}
-
 
 def _per_column(value: Any, n: int, name: str) -> "list[bool]":
     """Broadcast a bool, or validate a per-column list of bools, to length ``n``."""
@@ -289,7 +281,7 @@ def sort(
         desc = _per_column(descending, n, "descending")
         nulls = _per_column(nulls_last, n, "nulls_last")
         expressions = [
-            (_unbound(c), _SORT_DIRECTIONS[(desc[i], nulls[i])])
+            (_unbound(c), sort_direction(desc[i], nulls[i]))
             for i, c in enumerate(columns)
         ]
         return self._next(_plan.sort(self._plan, expressions=expressions))
@@ -333,8 +325,7 @@ def top_n(
         desc = _per_column(descending, count, "descending")
         nulls = _per_column(nulls_last, count, "nulls_last")
         sorts = [
-            (_unbound(c), _SORT_DIRECTIONS[(desc[i], nulls[i])])
-            for i, c in enumerate(keys)
+            (_unbound(c), sort_direction(desc[i], nulls[i])) for i, c in enumerate(keys)
         ]
         return self._next(
             _plan.top_n(

src/substrait/extension_registry/registry.py

@@ -154,6 +154,27 @@ def list_functions_across_urns(
         """List all matching functions across all URNs."""
         return self._find_matching_functions(function_name, signature)
 
+    def find_function(
+        self,
+        function_name: str,
+        signature: tuple[Type] | list[Type],
+        urns: Optional[list[str]] = None,
+    ) -> Optional[tuple[FunctionEntry, Type]]:
+        """Find the best-matching function for ``function_name`` across ``urns``.
+
+        Searches ``urns`` in order (every registered URN when ``None``) and returns
+        the first ``(FunctionEntry, output_type)`` whose overload satisfies
+        ``signature``, or ``None``. The winning extension URN is ``entry.urn``.
+
+        Generalizes :meth:`lookup_function` (a single URN) and
+        :meth:`list_functions_across_urns` (every URN) to an ordered subset, so a
+        caller resolving a name that lives in several extensions -- preferring, say,
+        the base arithmetic extension over its decimal variant -- needs one call
+        rather than a per-URN ``lookup_function`` loop.
+        """
+        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)
 

src/substrait/narwhals/expression.py

@@ -1,27 +1,10 @@
-import substrait.extended_expression_pb2 as stee
-import substrait.type_pb2 as stp
-
 from substrait.builders.extended_expression import (
-    ExtendedExpressionOrUnbound,
     UnboundExtendedExpression,
-    resolve_expression,
     scalar_function,
 )
-from substrait.extension_registry import ExtensionRegistry
-
-
-def _alias(
-    expr: ExtendedExpressionOrUnbound,
-    alias: str = None,
-):
-    def resolve(
-        base_schema: stp.NamedStruct, registry: ExtensionRegistry
-    ) -> stee.ExtendedExpression:
-        bound_expression = resolve_expression(expr, base_schema, registry)
-        bound_expression.referred_expr[0].output_names[0] = alias
-        return bound_expression
-
-    return resolve
+from substrait.builders.extended_expression import (
+    alias as _alias,
+)
 
 
 class Expression:

src/substrait/utils/__init__.py

@@ -62,6 +62,38 @@ def merge_extension_declarations(
                     seen_extension_functions.add(ident)
                     ret.append(declaration)
             else:
-                raise Exception("")  # TODO handle extension types
+                # TODO handle extension type / type-variation declarations.
+                mapping_type = declaration.WhichOneof("mapping_type")
+                raise NotImplementedError(
+                    f"cannot merge extension declaration of type {mapping_type!r}; "
+                    f"only 'extension_function' declarations are supported so far"
+                )
 
     return ret
+
+
+def merge_extensions_into(target, *sources):
+    """Merge the extension URNs and declarations of ``sources`` into ``target`` in place.
+
+    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.
+
+    ``target`` and each ``source`` are messages carrying repeated ``extension_urns``
+    and ``extensions`` fields (a ``Plan`` or an ``ExtendedExpression``). Unlike the
+    functions above, this mutates ``target`` rather than returning a new list, for
+    callers that have already materialized the message they are accumulating into.
+    """
+    merged_urns = merge_extension_urns(
+        target.extension_urns, *(s.extension_urns for s in sources)
+    )
+    merged_extensions = merge_extension_declarations(
+        target.extensions, *(s.extensions for s in sources)
+    )
+    target.ClearField("extension_urns")
+    target.extension_urns.extend(merged_urns)
+    target.ClearField("extensions")
+    target.extensions.extend(merged_extensions)

tests/extension_registry/test_registry_lookup.py

@@ -177,3 +177,41 @@ def test_function_with_discrete_nullability_nonexisting(registry):
         )
         is None
     )
+
+
+# ``add`` exists both in the test extension and in the default
+# functions_arithmetic, so it exercises resolution across several URNs.
+_ARITHMETIC = "extension:io.substrait:functions_arithmetic"
+_TEST = "extension:test:functions"
+
+
+def test_find_function_single_urn_matches_lookup(registry):
+    signature = [i8(nullable=False), i8(nullable=False)]
+    match = registry.find_function("add", signature, [_TEST])
+    assert match is not None
+    assert match[0].urn == _TEST
+    assert match == registry.lookup_function(_TEST, "add", signature)
+
+
+def test_find_function_respects_candidate_urn_order(registry):
+    signature = [i8(nullable=False), i8(nullable=False)]
+    # The winning extension is the first candidate URN that has a matching
+    # overload; ``entry.urn`` recovers it.
+    assert registry.find_function("add", signature, [_ARITHMETIC, _TEST])[0].urn == (
+        _ARITHMETIC
+    )
+    assert registry.find_function("add", signature, [_TEST, _ARITHMETIC])[0].urn == (
+        _TEST
+    )
+
+
+def test_find_function_no_matching_overload_returns_none(registry):
+    # ``add`` has no single-argument overload in any extension.
+    assert registry.find_function("add", [i8(nullable=False)], [_ARITHMETIC]) is None
+
+
+def test_find_function_searches_all_urns_when_unspecified(registry):
+    signature = [i8(nullable=False), i8(nullable=False)]
+    match = registry.find_function("add", signature)
+    assert match is not None
+    assert match == registry.list_functions_across_urns("add", signature)[0]