From 2421a991e118cbbf8bcbbcfe41f5d0a5f9b2443b Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Thu, 16 Jul 2026 19:58:48 +0200 Subject: [PATCH] refactor: de-duplicate small helpers in the DataFrame/Expr API Follow-up cleanup from the review of #204 (DRY, no behavior change): consolidate four small helpers in the substrait.dataframe / substrait.narwhals layer that re-implemented logic living elsewhere, so the two copies cannot drift apart. - Extension merge: replace expr._merge_extensions_into (which deduped on strict proto equality) with a shared utils.merge_extensions_into that reuses merge_extension_urns / merge_extension_declarations, keying on the same anchor/name identity as builders.plan._merge_extensions. - Sort direction: drop the duplicate (descending, nulls_last) -> SortDirection mapping; keep a single _SORT_DIRECTIONS + sort_direction() in dataframe.expr, used by both dataframe.expr and dataframe.frame. - Multi-URN lookup: add ExtensionRegistry.find_function(name, signature, urns), a generalization of lookup_function / list_functions_across_urns over an ordered URN subset, so _resolve_over_urns resolves in one registry call and recovers the winning extension via FunctionEntry.urn instead of looping. - Alias: extract the referred_expr[0].output_names[0] rewrite into a shared builders.extended_expression.alias, used by both Expr.alias and the Narwhals expression wrapper. Also replace the empty Exception("") in merge_extension_declarations with an informative NotImplementedError that names the unsupported mapping type. Closes #208 --- src/substrait/builders/extended_expression.py | 21 ++++++ src/substrait/dataframe/expr.py | 69 ++++++++----------- src/substrait/dataframe/frame.py | 15 +--- src/substrait/extension_registry/registry.py | 21 ++++++ src/substrait/narwhals/expression.py | 23 +------ src/substrait/utils/__init__.py | 34 ++++++++- .../test_registry_lookup.py | 38 ++++++++++ tests/test_utils.py | 63 ++++++++++++++++- 8 files changed, 211 insertions(+), 73 deletions(-) diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index d99c5c08..059e2776 100644 --- a/src/substrait/builders/extended_expression.py +++ b/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) diff --git a/src/substrait/dataframe/expr.py b/src/substrait/dataframe/expr.py index c48aec44..012ad231 100644 --- a/src/substrait/dataframe/expr.py +++ b/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()" diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index 2b7cd6ff..0517a8b9 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -37,7 +37,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 @@ -75,14 +75,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, -} - # How often execution context variables (current_timestamp, ...) are evaluated. _VARIABLE_EVAL_MODES = { "per_plan": stplan.ExecutionBehavior.VARIABLE_EVALUATION_MODE_PER_PLAN, @@ -296,7 +288,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)) @@ -340,8 +332,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( diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 136de015..41c3ff0a 100644 --- a/src/substrait/extension_registry/registry.py +++ b/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) diff --git a/src/substrait/narwhals/expression.py b/src/substrait/narwhals/expression.py index 8f977373..cd02c92e 100644 --- a/src/substrait/narwhals/expression.py +++ b/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: diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index 37179996..79026fd7 100644 --- a/src/substrait/utils/__init__.py +++ b/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) diff --git a/tests/extension_registry/test_registry_lookup.py b/tests/extension_registry/test_registry_lookup.py index 273dab0a..32d42a4a 100644 --- a/tests/extension_registry/test_registry_lookup.py +++ b/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] diff --git a/tests/test_utils.py b/tests/test_utils.py index 002aa7a4..a2a599e2 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,7 +1,12 @@ +import substrait.extended_expression_pb2 as stee import substrait.extensions.extensions_pb2 as ste import substrait.type_pb2 as stt -from substrait.utils import merge_extension_urns, type_num_names +from substrait.utils import ( + merge_extension_urns, + merge_extensions_into, + type_num_names, +) def test_type_num_names_flat_struct(): @@ -102,3 +107,59 @@ def test_merge_extension_urns_deduplicates(): assert len(merged_urns) == 2 assert merged_urns[0].urn == "extension:example:test" assert merged_urns[1].urn == "extension:example:other" + + +def _extension_function(urn_reference, function_anchor, name): + return ste.SimpleExtensionDeclaration( + extension_function=ste.SimpleExtensionDeclaration.ExtensionFunction( + extension_urn_reference=urn_reference, + function_anchor=function_anchor, + name=name, + ) + ) + + +def test_merge_extensions_into_appends_new_and_dedupes_on_identity(): + """merge_extensions_into keeps target's entries and appends only novel ones, + keying on the same anchor/name identity as the merge_* helpers.""" + target = stee.ExtendedExpression( + extension_urns=[ste.SimpleExtensionURN(extension_urn_anchor=1, urn="A")], + extensions=[_extension_function(1, 10, "f:i8")], + ) + source = stee.ExtendedExpression( + extension_urns=[ + ste.SimpleExtensionURN(extension_urn_anchor=1, urn="A"), # dup URN string + ste.SimpleExtensionURN(extension_urn_anchor=2, urn="B"), + ], + extensions=[ + # Same (urn reference, name) as target's -- a duplicate by identity even + # though the function anchor differs, so it must not be appended (this is + # what distinguishes identity dedup from strict proto equality). + _extension_function(1, 99, "f:i8"), + _extension_function(2, 11, "g:i8"), + ], + ) + + merge_extensions_into(target, source) + + assert [u.urn for u in target.extension_urns] == ["A", "B"] + assert [d.extension_function.name for d in target.extensions] == ["f:i8", "g:i8"] + # target's original declaration is kept; source's identity-duplicate is dropped. + assert target.extensions[0].extension_function.function_anchor == 10 + + +def test_merge_extensions_into_merges_multiple_sources(): + target = stee.ExtendedExpression() + source1 = stee.ExtendedExpression( + extension_urns=[ste.SimpleExtensionURN(extension_urn_anchor=1, urn="A")], + extensions=[_extension_function(1, 10, "f:i8")], + ) + source2 = stee.ExtendedExpression( + extension_urns=[ste.SimpleExtensionURN(extension_urn_anchor=2, urn="B")], + extensions=[_extension_function(2, 11, "g:i8")], + ) + + merge_extensions_into(target, source1, source2) + + assert [u.urn for u in target.extension_urns] == ["A", "B"] + assert [d.extension_function.name for d in target.extensions] == ["f:i8", "g:i8"]