From d990f67f98cff5c0c5715168613f89f2bb0d2cf3 Mon Sep 17 00:00:00 2001 From: Thomas Gschwind Date: Wed, 2 Sep 2026 00:03:20 +0200 Subject: [PATCH 1/4] test(dataframe): demonstrate missing enum-argument function support Enumeration-argument functions like extract (component/indexing declared with an options domain in functions_datetime.yaml) are unusable today: sub.f.extract(col, component="YEAR") raises "Unknown function extract" because the overload's arity counts the enum positions but the resolved signature carries only the value operands. Even once resolved, the enum must serialize into ScalarFunction.arguments as a FunctionArgument.enum, not into options (consumers such as DuckDB read enum selections only from arguments). Adds three xfail(strict=True) tests pinning the observable contract: resolves, and the enum lands in arguments (not options) in signature order. Marked strict so they flip to green and drop the marker when the fix lands. --- tests/dataframe/test_enum_arguments.py | 140 +++++++++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 tests/dataframe/test_enum_arguments.py diff --git a/tests/dataframe/test_enum_arguments.py b/tests/dataframe/test_enum_arguments.py new file mode 100644 index 0000000..7e0c10c --- /dev/null +++ b/tests/dataframe/test_enum_arguments.py @@ -0,0 +1,140 @@ +"""Enumeration-argument functions (``extract``, ``round_temporal``, ...). + +Some standard functions take *enumeration arguments* -- the Substrait spec +(``functions_datetime.yaml``) declares them under ``args:`` with an ``options:`` +domain rather than a ``value:`` type. ``extract`` is the canonical example:: + + - name: "extract" + impls: + - args: + - name: component # enumeration argument + options: [ YEAR, ISO_YEAR, US_YEAR, UNIX_TIME ] + - name: x # value argument + value: date + return: i64 + +Per the spec these serialize into ``ScalarFunction.arguments`` as +``FunctionArgument.enum`` (interleaved with the value arguments, in signature +order) -- NOT into ``ScalarFunction.options``, which carries only *behavioral* +options (``overflow``, ``rounding``, ...). Consumers such as DuckDB read enum +selections exclusively from ``arguments``. + +These tests demonstrate two defects in the current DataFrame / builder pipeline: + +1. **Resolution.** ``sub.f.extract(col, component="YEAR")`` cannot be built at + all: the overload's arity counts the enum positions, but the signature handed + to the registry contains only the value operands, so the arity check never + matches and resolution raises ``Unknown function extract``. +2. **Serialization.** Even once it resolves, the enum selection must land in + ``arguments`` as a ``FunctionArgument.enum``. Routing it into ``options`` + (as the current ``**options`` kwarg path would) produces a plan that omits + the enum from ``arguments`` -- which DuckDB's consumer then reads as an empty + enum vector and crashes on. + +They are marked ``xfail(strict=True)``: they fail today and must flip to passing +(and lose the marker) when enum-argument support lands. The call site here -- +enum selections passed as keyword arguments -- is the user-facing API and is +independent of how resolution is implemented internally. +""" + +import pytest +import substrait.algebra_pb2 as stalg + +import substrait.dataframe as sub +from substrait.builders.type import precision_timestamp + +_NOT_YET = "enumeration-argument functions are not resolved/serialized yet" + + +def _scalar_functions(message) -> list: + """Every ``Expression.ScalarFunction`` reachable anywhere in a proto tree.""" + target = stalg.Expression.ScalarFunction.DESCRIPTOR.full_name + found: list = [] + + def walk(msg): + for field, value in msg.ListFields(): + if field.message_type is None: + continue + items = value if field.is_repeated else [value] + for item in items: + if field.message_type.full_name == target: + found.append(item) + walk(item) + + walk(message) + return found + + +def _function_name(plan, scalar_function) -> str: + """Resolve a ScalarFunction's declared extension name via its anchor.""" + ref = scalar_function.function_reference + for decl in plan.extensions: + fn = decl.extension_function + if fn.function_anchor == ref: + return fn.name + return "" + + +def _arg_kinds(scalar_function) -> list: + return [a.WhichOneof("arg_type") for a in scalar_function.arguments] + + +@pytest.mark.xfail(strict=True, reason=_NOT_YET) +def test_extract_with_a_single_enum_argument_resolves(): + """The two-argument ``extract(component, date)`` overload must build.""" + df = sub.read_named_table("t", {"d": sub.date}) + + # Must not raise "Unknown function extract". + plan = df.with_columns(y=sub.f.extract(sub.col("d"), component="YEAR")).to_plan() + + (extract,) = [ + sf for sf in _scalar_functions(plan) if _function_name(plan, sf).startswith("extract") + ] + assert extract, "extract did not resolve to a registry function" + + +@pytest.mark.xfail(strict=True, reason=_NOT_YET) +def test_extract_serializes_the_enum_as_an_argument_not_an_option(): + """component=YEAR must be a FunctionArgument.enum, in signature position. + + Spec order for this overload is [component (enum), x (value)], so the + positional value operand follows the enum selection. + """ + df = sub.read_named_table("t", {"d": sub.date}) + + plan = df.with_columns(y=sub.f.extract(sub.col("d"), component="YEAR")).to_plan() + + (extract,) = [ + sf for sf in _scalar_functions(plan) if _function_name(plan, sf).startswith("extract") + ] + assert _arg_kinds(extract) == ["enum", "value"], ( + f"enum not carried in arguments in signature order: {_arg_kinds(extract)}" + ) + assert extract.arguments[0].enum == "YEAR" + assert list(extract.options) == [], ( + f"enum selection leaked into options (crashes DuckDB): {extract.options}" + ) + + +@pytest.mark.xfail(strict=True, reason=_NOT_YET) +def test_extract_day_of_month_carries_both_enum_arguments(): + """The exact combination behind the reported DuckDB crash. + + ``extract(component="DAY", indexing="ONE", )`` has two enum + arguments; both must appear in ``arguments`` (order [component, indexing, x]) + and neither in ``options``. + """ + df = sub.read_named_table("t", {"ts": precision_timestamp(6)}) + + plan = df.with_columns( + dom=sub.f.extract(sub.col("ts"), component="DAY", indexing="ONE") + ).to_plan() + + (extract,) = [ + sf for sf in _scalar_functions(plan) if _function_name(plan, sf).startswith("extract") + ] + assert _arg_kinds(extract) == ["enum", "enum", "value"], _arg_kinds(extract) + assert [extract.arguments[0].enum, extract.arguments[1].enum] == ["DAY", "ONE"] + assert list(extract.options) == [], ( + f"enum selections leaked into options (crashes DuckDB): {extract.options}" + ) From e3c0bbe8c56dc854cc6fe22586b0ba66791b4719 Mon Sep 17 00:00:00 2001 From: Thomas Gschwind Date: Wed, 2 Sep 2026 01:15:04 +0200 Subject: [PATCH 2/4] feat(builders): resolve and serialize enumeration-argument functions Functions like `extract`/`round_temporal` declare enumeration arguments (`component`, `indexing`) whose arity counts toward the overload. The DataFrame layer passed only the value operands as the signature, so the arity check never matched and resolution raised "Unknown function extract"; even resolved, the enum selection was routed into `options` as a FunctionOption, which the spec reserves for behavioral options and which DuckDB reads as an empty enum vector (crashing on it). Thread enum selections (given by name, e.g. `component="YEAR"`) through resolution and emission: the registry folds them into the signature it matches, and the builders interleave them into `arguments` as `FunctionArgument.enum` in declared order, leaving only behavioral options for `FunctionOption`. The low-level enum-in-signature registry contract still works. Flips the previously-xfail enum tests to green. --- src/substrait/builders/extended_expression.py | 73 +++++++++++----- src/substrait/dataframe/expr.py | 2 +- .../extension_registry/function_entry.py | 87 ++++++++++++++++++- src/substrait/extension_registry/registry.py | 23 ++++- tests/dataframe/test_enum_arguments.py | 35 +++----- 5 files changed, 167 insertions(+), 53 deletions(-) diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index 2738a55..8b45388 100644 --- a/src/substrait/builders/extended_expression.py +++ b/src/substrait/builders/extended_expression.py @@ -83,6 +83,36 @@ def _function_options(options): return result +def _function_arguments(func_entry, bound_expressions, options): + """Interleave value operands and enumeration selections into FunctionArguments. + + Value operands come from ``bound_expressions`` (positional), enumeration + selections from ``options`` (by argument name). ``func_entry`` weaves them + into declared signature order -- enum selections as ``FunctionArgument.enum``, + values as ``FunctionArgument.value``. Returns ``(arguments, options)`` where + the second element is the options *not* consumed as enumeration arguments + (i.e. the behavioral options destined for ``FunctionOption``). + """ + value_args = ( + stalg.FunctionArgument(value=e.referred_expr[0].expression) + for e in bound_expressions + ) + plan = func_entry.interleave_arguments(value_args, options) + if plan is None: + # Resolution already matched this overload, so interleaving should not + # fail; fall back to the value-only form rather than lose the operands. + return [ + stalg.FunctionArgument(value=e.referred_expr[0].expression) + for e in bound_expressions + ], options + ordered, behavioral_options = plan + arguments = [ + stalg.FunctionArgument(enum=selection) if kind == "enum" else selection + for kind, selection in ordered + ] + return arguments, behavioral_options + + def resolve_expression( expression: ExtendedExpressionOrUnbound, base_schema: stp.NamedStruct, @@ -581,26 +611,25 @@ def resolve( signature = [typ for es in expression_schemas for typ in es.types] - func = registry.lookup_function(urn, function, signature) + func = registry.lookup_function(urn, function, signature, options=options) if not func: raise Exception(f"Unknown function {function} for {signature}") func_ref = function_reference(urn, str(func[0])) + arguments, behavioral_options = _function_arguments( + func[0], bound_expressions, options + ) + return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( expression=stalg.Expression( scalar_function=stalg.Expression.ScalarFunction( function_reference=func_ref, - arguments=[ - stalg.FunctionArgument( - value=e.referred_expr[0].expression - ) - for e in bound_expressions - ], - options=_function_options(options), + arguments=arguments, + options=_function_options(behavioral_options), output_type=func[1], ) ), @@ -654,23 +683,24 @@ def resolve( signature = [typ for es in expression_schemas for typ in es.types] - func = registry.lookup_function(urn, function, signature) + func = registry.lookup_function(urn, function, signature, options=options) if not func: raise Exception(f"Unknown function {function} for {signature}") func_ref = function_reference(urn, str(func[0])) + arguments, behavioral_options = _function_arguments( + func[0], bound_expressions, options + ) + return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( measure=stalg.AggregateFunction( function_reference=func_ref, - arguments=[ - stalg.FunctionArgument(value=e.referred_expr[0].expression) - for e in bound_expressions - ], - options=_function_options(options), + arguments=arguments, + options=_function_options(behavioral_options), output_type=func[1], invocation=invocation if invocation is not None @@ -724,26 +754,25 @@ def resolve( signature = [typ for es in expression_schemas for typ in es.types] - func = registry.lookup_function(urn, function, signature) + func = registry.lookup_function(urn, function, signature, options=options) if not func: raise Exception(f"Unknown function {function} for {signature}") func_ref = function_reference(urn, str(func[0])) + arguments, behavioral_options = _function_arguments( + func[0], bound_expressions, options + ) + return stee.ExtendedExpression( referred_expr=[ stee.ExpressionReference( expression=stalg.Expression( window_function=stalg.Expression.WindowFunction( function_reference=func_ref, - arguments=[ - stalg.FunctionArgument( - value=e.referred_expr[0].expression - ) - for e in bound_expressions - ], - options=_function_options(options), + arguments=arguments, + options=_function_options(behavioral_options), output_type=func[1], partitions=[ e.referred_expr[0].expression for e in bound_partitions diff --git a/src/substrait/dataframe/expr.py b/src/substrait/dataframe/expr.py index c185a37..f2ba9de 100644 --- a/src/substrait/dataframe/expr.py +++ b/src/substrait/dataframe/expr.py @@ -251,7 +251,7 @@ def _resolve_over_urns( for b in bound for typ in infer_extended_expression_schema(b, registry=registry).types ] - match = registry.find_function(name, signature, urns) + match = registry.find_function(name, signature, urns, options=options) if match is not None: winning_urn = match[0].urn return builder( diff --git a/src/substrait/extension_registry/function_entry.py b/src/substrait/extension_registry/function_entry.py index b2bd0d5..9cea4f1 100644 --- a/src/substrait/extension_registry/function_entry.py +++ b/src/substrait/extension_registry/function_entry.py @@ -10,6 +10,8 @@ from .signature_checker_helpers import covers, normalize_substrait_type_names +_MISSING = object() + class FunctionType(Enum): SCALAR = "scalar" @@ -31,6 +33,10 @@ def __init__( self.urn: str = urn self.function_type = function_type self.arguments = [] + # Argument names parallel to ``arguments`` / ``normalized_inputs``. Used to + # match enumeration selections, which arrive by name (e.g. component="YEAR") + # rather than by position, back to their declared signature slot. + self.arg_names: list = [] self.nullability = ( impl.nullability if impl.nullability else se.NullabilityHandling.MIRROR ) @@ -41,24 +47,97 @@ def __init__( self.normalized_inputs.append( normalize_substrait_type_names(arg.value) ) + self.arg_names.append(arg.name) elif isinstance(arg, se.EnumerationArg): self.arguments.append(arg.options) self.normalized_inputs.append("req") + self.arg_names.append(arg.name) def __repr__(self) -> str: return f"{self.name}:{'_'.join(self.normalized_inputs)}" - def satisfies_signature(self, signature: tuple | list) -> Optional[str]: + def interleave_arguments(self, value_items, options): + """Order value operands and enumeration selections per the signature. + + Substrait interleaves enumeration arguments with value arguments in + declared order (``extract`` is ``[component (enum), x (value)]``), but the + two arrive separately: value operands positionally, enumeration selections + by argument name in ``options``. This walks the declared arguments and + weaves them back together. + + ``value_items`` is an iterator of value operands (consumed in order for + each value argument). Returns ``(ordered, remaining_options)`` where each + entry of ``ordered`` is ``("value", item)`` or ``("enum", selection)``, and + ``remaining_options`` are the options *not* consumed as enumeration + arguments -- i.e. the behavioral options. Returns ``None`` if the value + arity is wrong or a required enumeration argument is absent from + ``options``. + """ + remaining = dict(options or {}) + if self.impl.variadic: + # Variadic functions have no enumeration arguments; every operand is a + # value and every option is behavioral. + return [("value", item) for item in value_items], remaining + ordered: list = [] + for kind, name in zip(self.normalized_inputs, self.arg_names): + if kind == "req": # enumeration argument + if name not in remaining: + return None + ordered.append(("enum", str(remaining.pop(name)))) + else: + item = next(value_items, _MISSING) + if item is _MISSING: + return None + ordered.append(("value", item)) + if next(value_items, _MISSING) is not _MISSING: + return None # too many value operands + return ordered, remaining + + def _resolve_signature( + self, signature: tuple | list, options: Optional[dict] + ) -> Optional[list]: + """Interleave ``signature`` with enumeration selections for matching. + + Enumeration selections may be supplied two ways: interleaved into + ``signature`` as strings in declared order (the low-level registry + contract), or by argument name in ``options`` (how the DataFrame builders + pass them, keeping the value-only signature intact). An enum position + prefers ``options`` and otherwise consumes the next ``signature`` item. + Returns the value-and-enum sequence to match against ``self.arguments``, + or ``None`` on an arity mismatch. + """ + remaining = dict(options or {}) + items = iter(signature) + interleaved: list = [] + for kind, name in zip(self.normalized_inputs, self.arg_names): + if kind == "req" and name in remaining: # enum selection by name + interleaved.append(str(remaining.pop(name))) + continue + item = next(items, _MISSING) + if item is _MISSING: + return None + interleaved.append(item) + if next(items, _MISSING) is not _MISSING: + return None # more operands than the signature declares + return interleaved + + def satisfies_signature( + self, signature: tuple | list, options: Optional[dict] = None + ) -> Optional[str]: if self.impl.variadic: min_args_allowed = self.impl.variadic.min or 0 if len(signature) < min_args_allowed: return None inputs = [self.arguments[0]] * len(signature) + interleaved: list = list(signature) else: + interleaved = self._resolve_signature(signature, options) + if interleaved is None: + return None inputs = self.arguments - if len(inputs) != len(signature): + if len(inputs) != len(interleaved): return None - zipped_args = list(zip(inputs, signature)) + zipped_args = list(zip(inputs, interleaved)) parameters = {} for x, y in zipped_args: if isinstance(y, str): @@ -81,7 +160,7 @@ def satisfies_signature(self, signature: tuple | list) -> Optional[str]: [ p.__getattribute__(p.WhichOneof("kind")).nullability == Type.NULLABILITY_NULLABLE - for p in signature + for p in interleaved if isinstance(p, Type) ] ) diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 5ad082f..2d85c19 100644 --- a/src/substrait/extension_registry/registry.py +++ b/src/substrait/extension_registry/registry.py @@ -116,8 +116,15 @@ def _find_matching_functions( function_name: str, signature: tuple[Type] | list[Type], urns: list[str] | None = None, + options: Optional[dict] = None, ) -> list[tuple[FunctionEntry, Type]]: - """Helper method to find matching functions across specified URNs.""" + """Helper method to find matching functions across specified URNs. + + ``options`` carries enumeration-argument selections by name (e.g. + ``{"component": "YEAR"}``); an overload with enumeration arguments only + matches when ``signature`` supplies its value operands and ``options`` + supplies each enumeration selection. + """ matches = [] urns_to_search = ( urns if urns is not None else list(self._function_mapping.keys()) @@ -130,7 +137,7 @@ def _find_matching_functions( continue functions = self._function_mapping[urn][function_name] for f in functions: - rtn = f.satisfies_signature(signature) + rtn = f.satisfies_signature(signature, options) if rtn is not None: matches.append((f, rtn)) return matches @@ -141,9 +148,12 @@ def lookup_function( urn: str, function_name: str, signature: tuple[Type] | list[Type], + options: Optional[dict] = None, ) -> Optional[tuple[FunctionEntry, Type]]: """Look up a function within a specific URN.""" - matches = self._find_matching_functions(function_name, signature, [urn]) + matches = self._find_matching_functions( + function_name, signature, [urn], options + ) return matches[0] if matches else None def list_functions( @@ -163,12 +173,15 @@ def find_function( function_name: str, signature: tuple[Type] | list[Type], urns: Optional[list[str]] = None, + options: Optional[dict] = 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``. + ``options`` carries enumeration-argument selections by name (see + :meth:`_find_matching_functions`). Generalizes :meth:`lookup_function` (a single URN) and :meth:`list_functions_across_urns` (every URN) to an ordered subset, so a @@ -176,7 +189,9 @@ def find_function( 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) + matches = self._find_matching_functions( + function_name, signature, urns, options + ) return matches[0] if matches else None def has_urn(self, urn: str) -> bool: diff --git a/tests/dataframe/test_enum_arguments.py b/tests/dataframe/test_enum_arguments.py index 7e0c10c..6a315de 100644 --- a/tests/dataframe/test_enum_arguments.py +++ b/tests/dataframe/test_enum_arguments.py @@ -19,32 +19,26 @@ options (``overflow``, ``rounding``, ...). Consumers such as DuckDB read enum selections exclusively from ``arguments``. -These tests demonstrate two defects in the current DataFrame / builder pipeline: - -1. **Resolution.** ``sub.f.extract(col, component="YEAR")`` cannot be built at - all: the overload's arity counts the enum positions, but the signature handed - to the registry contains only the value operands, so the arity check never - matches and resolution raises ``Unknown function extract``. -2. **Serialization.** Even once it resolves, the enum selection must land in - ``arguments`` as a ``FunctionArgument.enum``. Routing it into ``options`` - (as the current ``**options`` kwarg path would) produces a plan that omits - the enum from ``arguments`` -- which DuckDB's consumer then reads as an empty - enum vector and crashes on. - -They are marked ``xfail(strict=True)``: they fail today and must flip to passing -(and lose the marker) when enum-argument support lands. The call site here -- -enum selections passed as keyword arguments -- is the user-facing API and is -independent of how resolution is implemented internally. +These tests guard two properties of the DataFrame / builder pipeline: + +1. **Resolution.** ``sub.f.extract(col, component="YEAR")`` must build. The + overload's arity counts the enum positions, so the registry has to fold the + enum selection into the signature it matches -- otherwise the value-only + signature never matches and resolution raises ``Unknown function extract``. +2. **Serialization.** The enum selection must land in ``arguments`` as a + ``FunctionArgument.enum``, not in ``options``. Routing it into ``options`` + produces a plan that omits the enum from ``arguments`` -- which DuckDB's + consumer then reads as an empty enum vector and crashes on. + +The call site here -- enum selections passed as keyword arguments -- is the +user-facing API and is independent of how resolution is implemented internally. """ -import pytest import substrait.algebra_pb2 as stalg import substrait.dataframe as sub from substrait.builders.type import precision_timestamp -_NOT_YET = "enumeration-argument functions are not resolved/serialized yet" - def _scalar_functions(message) -> list: """Every ``Expression.ScalarFunction`` reachable anywhere in a proto tree.""" @@ -79,7 +73,6 @@ def _arg_kinds(scalar_function) -> list: return [a.WhichOneof("arg_type") for a in scalar_function.arguments] -@pytest.mark.xfail(strict=True, reason=_NOT_YET) def test_extract_with_a_single_enum_argument_resolves(): """The two-argument ``extract(component, date)`` overload must build.""" df = sub.read_named_table("t", {"d": sub.date}) @@ -93,7 +86,6 @@ def test_extract_with_a_single_enum_argument_resolves(): assert extract, "extract did not resolve to a registry function" -@pytest.mark.xfail(strict=True, reason=_NOT_YET) def test_extract_serializes_the_enum_as_an_argument_not_an_option(): """component=YEAR must be a FunctionArgument.enum, in signature position. @@ -116,7 +108,6 @@ def test_extract_serializes_the_enum_as_an_argument_not_an_option(): ) -@pytest.mark.xfail(strict=True, reason=_NOT_YET) def test_extract_day_of_month_carries_both_enum_arguments(): """The exact combination behind the reported DuckDB crash. From 77ab2765a098a0f8ad8d471c6301cab4375fa5ef Mon Sep 17 00:00:00 2001 From: Thomas Gschwind Date: Wed, 2 Sep 2026 09:21:45 +0200 Subject: [PATCH 3/4] style: apply ruff format to registry and enum-argument tests --- src/substrait/extension_registry/registry.py | 4 +--- tests/dataframe/test_enum_arguments.py | 12 +++++++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 2d85c19..9412aae 100644 --- a/src/substrait/extension_registry/registry.py +++ b/src/substrait/extension_registry/registry.py @@ -189,9 +189,7 @@ def find_function( 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, options - ) + matches = self._find_matching_functions(function_name, signature, urns, options) return matches[0] if matches else None def has_urn(self, urn: str) -> bool: diff --git a/tests/dataframe/test_enum_arguments.py b/tests/dataframe/test_enum_arguments.py index 6a315de..5cc4087 100644 --- a/tests/dataframe/test_enum_arguments.py +++ b/tests/dataframe/test_enum_arguments.py @@ -81,7 +81,9 @@ def test_extract_with_a_single_enum_argument_resolves(): plan = df.with_columns(y=sub.f.extract(sub.col("d"), component="YEAR")).to_plan() (extract,) = [ - sf for sf in _scalar_functions(plan) if _function_name(plan, sf).startswith("extract") + sf + for sf in _scalar_functions(plan) + if _function_name(plan, sf).startswith("extract") ] assert extract, "extract did not resolve to a registry function" @@ -97,7 +99,9 @@ def test_extract_serializes_the_enum_as_an_argument_not_an_option(): plan = df.with_columns(y=sub.f.extract(sub.col("d"), component="YEAR")).to_plan() (extract,) = [ - sf for sf in _scalar_functions(plan) if _function_name(plan, sf).startswith("extract") + sf + for sf in _scalar_functions(plan) + if _function_name(plan, sf).startswith("extract") ] assert _arg_kinds(extract) == ["enum", "value"], ( f"enum not carried in arguments in signature order: {_arg_kinds(extract)}" @@ -122,7 +126,9 @@ def test_extract_day_of_month_carries_both_enum_arguments(): ).to_plan() (extract,) = [ - sf for sf in _scalar_functions(plan) if _function_name(plan, sf).startswith("extract") + sf + for sf in _scalar_functions(plan) + if _function_name(plan, sf).startswith("extract") ] assert _arg_kinds(extract) == ["enum", "enum", "value"], _arg_kinds(extract) assert [extract.arguments[0].enum, extract.arguments[1].enum] == ["DAY", "ONE"] From a1a2f290aa4d353f3a6a18fa7ba79627395908dc Mon Sep 17 00:00:00 2001 From: Thomas Gschwind Date: Tue, 8 Sep 2026 16:33:47 +0200 Subject: [PATCH 4/4] refactor(dataframe)!: make enum function arguments positional via sub.enum() Enumeration arguments are positional members of a function's argument list in the spec, serialized as FunctionArgument.enum and interleaved with value operands in declared order. Replace the earlier **options kwargs channel (extract(col, component="YEAR")) with a positional marker sub.enum("YEAR") written in its argument slot, and reserve **options for behavioral options only. This mimics substrait-go's types.Enum and substrait-java's EnumArg. The registry matcher reverts to positional-string matching: an enum selection is a token in the signature that must belong to the overload's domain, counting toward arity. This drops the name-based interleaving and the reject-unconsumed-enum disambiguation machinery. Signed-off-by: Thomas Gschwind --- src/substrait/builders/extended_expression.py | 166 ++++++++-------- src/substrait/dataframe/__init__.py | 2 + src/substrait/dataframe/expr.py | 27 ++- src/substrait/dataframe/functions.py | 12 +- .../extension_registry/function_entry.py | 99 ++-------- src/substrait/extension_registry/registry.py | 28 ++- src/substrait/utils/display.py | 10 +- tests/dataframe/test_enum_arguments.py | 187 ++++++++++++++---- 8 files changed, 300 insertions(+), 231 deletions(-) diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index 8b45388..46464a3 100644 --- a/src/substrait/builders/extended_expression.py +++ b/src/substrait/builders/extended_expression.py @@ -58,6 +58,30 @@ def fresh_rel_anchors(): ExtendedExpressionOrUnbound = Union[stee.ExtendedExpression, UnboundExtendedExpression] +class EnumArg: + """A positional enumeration-argument selection (e.g. ``extract``'s component). + + Substrait models an enumeration argument as a positional member of a + function's argument list -- interleaved with value operands in declared + order -- that serializes as ``FunctionArgument.enum`` (not as a behavioral + ``FunctionOption``). Its valid domain is defined per overload by the + extension YAML and only checked at resolve time, so the selection is carried + as a plain string token, mirroring substrait-go's ``types.Enum`` and + substrait-java's ``EnumArg``. Pass one positionally into an ``f.*`` call via + :func:`substrait.dataframe.enum`. + """ + + __slots__ = ("value",) + + def __init__(self, value: str) -> None: + if not isinstance(value, str): + raise ValueError(f"enum() takes a selection string, got {value!r}") + self.value = value + + def __repr__(self) -> str: + return f"enum({self.value!r})" + + def _alias_or_inferred( alias: Union[Iterable[str], str, None], op: str, @@ -83,34 +107,46 @@ def _function_options(options): return result -def _function_arguments(func_entry, bound_expressions, options): - """Interleave value operands and enumeration selections into FunctionArguments. +def _function_signature(bound_expressions, registry): + """The flattened match signature for a call's positional arguments. - Value operands come from ``bound_expressions`` (positional), enumeration - selections from ``options`` (by argument name). ``func_entry`` weaves them - into declared signature order -- enum selections as ``FunctionArgument.enum``, - values as ``FunctionArgument.value``. Returns ``(arguments, options)`` where - the second element is the options *not* consumed as enumeration arguments - (i.e. the behavioral options destined for ``FunctionOption``). + Enumeration arguments (:class:`EnumArg`) contribute their selection *string + token* -- which the registry matches against the overload's domain -- while + value operands contribute their inferred Substrait ``Type``\\ s, interleaved + in call order. This mixed sequence of tokens and types is what + :meth:`FunctionEntry.satisfies_signature` expects. """ - value_args = ( - stalg.FunctionArgument(value=e.referred_expr[0].expression) - for e in bound_expressions - ) - plan = func_entry.interleave_arguments(value_args, options) - if plan is None: - # Resolution already matched this overload, so interleaving should not - # fail; fall back to the value-only form rather than lose the operands. - return [ - stalg.FunctionArgument(value=e.referred_expr[0].expression) - for e in bound_expressions - ], options - ordered, behavioral_options = plan - arguments = [ - stalg.FunctionArgument(enum=selection) if kind == "enum" else selection - for kind, selection in ordered - ] - return arguments, behavioral_options + signature: list = [] + for b in bound_expressions: + if isinstance(b, EnumArg): + signature.append(b.value) + else: + signature.extend( + infer_extended_expression_schema(b, registry=registry).types + ) + return signature + + +def _function_arguments(bound_expressions): + """Build FunctionArguments from a call's positional arguments, in call order. + + An :class:`EnumArg` becomes a ``FunctionArgument.enum`` carrying its selection + token; every other operand becomes a ``FunctionArgument.value``. Returns + ``(arguments, names)`` where ``names`` are the per-argument labels used to + infer an output column name. + """ + arguments = [] + names = [] + for b in bound_expressions: + if isinstance(b, EnumArg): + arguments.append(stalg.FunctionArgument(enum=b.value)) + names.append(b.value) + else: + arguments.append( + stalg.FunctionArgument(value=b.referred_expr[0].expression) + ) + names.append(b.referred_expr[0].output_names[0]) + return arguments, names def resolve_expression( @@ -601,26 +637,22 @@ def resolve( base_schema: stp.NamedStruct, registry: ExtensionRegistry ) -> stee.ExtendedExpression: bound_expressions = [ - resolve_expression(e, base_schema, registry) for e in expressions - ] - - expression_schemas = [ - infer_extended_expression_schema(b, registry=registry) - for b in bound_expressions + e + if isinstance(e, EnumArg) + else resolve_expression(e, base_schema, registry) + for e in expressions ] - signature = [typ for es in expression_schemas for typ in es.types] + signature = _function_signature(bound_expressions, registry) - func = registry.lookup_function(urn, function, signature, options=options) + func = registry.lookup_function(urn, function, signature) if not func: raise Exception(f"Unknown function {function} for {signature}") func_ref = function_reference(urn, str(func[0])) - arguments, behavioral_options = _function_arguments( - func[0], bound_expressions, options - ) + arguments, arg_names = _function_arguments(bound_expressions) return stee.ExtendedExpression( referred_expr=[ @@ -629,15 +661,11 @@ def resolve( scalar_function=stalg.Expression.ScalarFunction( function_reference=func_ref, arguments=arguments, - options=_function_options(behavioral_options), + options=_function_options(options), output_type=func[1], ) ), - output_names=_alias_or_inferred( - alias, - function, - [e.referred_expr[0].output_names[0] for e in bound_expressions], - ), + output_names=_alias_or_inferred(alias, function, arg_names), ) ], base_schema=base_schema, @@ -669,30 +697,26 @@ def resolve( base_schema: stp.NamedStruct, registry: ExtensionRegistry ) -> stee.ExtendedExpression: bound_expressions: Iterable[stee.ExtendedExpression] = [ - resolve_expression(e, base_schema, registry) for e in expressions + e + if isinstance(e, EnumArg) + else resolve_expression(e, base_schema, registry) + for e in expressions ] bound_sorts = [ (resolve_expression(e, base_schema, registry), direction) for e, direction in sorts ] - expression_schemas = [ - infer_extended_expression_schema(b, registry=registry) - for b in bound_expressions - ] + signature = _function_signature(bound_expressions, registry) - signature = [typ for es in expression_schemas for typ in es.types] - - func = registry.lookup_function(urn, function, signature, options=options) + func = registry.lookup_function(urn, function, signature) if not func: raise Exception(f"Unknown function {function} for {signature}") func_ref = function_reference(urn, str(func[0])) - arguments, behavioral_options = _function_arguments( - func[0], bound_expressions, options - ) + arguments, arg_names = _function_arguments(bound_expressions) return stee.ExtendedExpression( referred_expr=[ @@ -700,7 +724,7 @@ def resolve( measure=stalg.AggregateFunction( function_reference=func_ref, arguments=arguments, - options=_function_options(behavioral_options), + options=_function_options(options), output_type=func[1], invocation=invocation if invocation is not None @@ -712,11 +736,7 @@ def resolve( for s, direction in bound_sorts ], ), - output_names=_alias_or_inferred( - alias, - "IfThen", - [e.referred_expr[0].output_names[0] for e in bound_expressions], - ), + output_names=_alias_or_inferred(alias, "IfThen", arg_names), ) ], base_schema=base_schema, @@ -740,30 +760,26 @@ def resolve( base_schema: stp.NamedStruct, registry: ExtensionRegistry ) -> stee.ExtendedExpression: bound_expressions: Iterable[stee.ExtendedExpression] = [ - resolve_expression(e, base_schema, registry) for e in expressions + e + if isinstance(e, EnumArg) + else resolve_expression(e, base_schema, registry) + for e in expressions ] bound_partitions = [ resolve_expression(e, base_schema, registry) for e in partitions ] - expression_schemas = [ - infer_extended_expression_schema(b, registry=registry) - for b in bound_expressions - ] - - signature = [typ for es in expression_schemas for typ in es.types] + signature = _function_signature(bound_expressions, registry) - func = registry.lookup_function(urn, function, signature, options=options) + func = registry.lookup_function(urn, function, signature) if not func: raise Exception(f"Unknown function {function} for {signature}") func_ref = function_reference(urn, str(func[0])) - arguments, behavioral_options = _function_arguments( - func[0], bound_expressions, options - ) + arguments, arg_names = _function_arguments(bound_expressions) return stee.ExtendedExpression( referred_expr=[ @@ -772,18 +788,14 @@ def resolve( window_function=stalg.Expression.WindowFunction( function_reference=func_ref, arguments=arguments, - options=_function_options(behavioral_options), + options=_function_options(options), output_type=func[1], partitions=[ e.referred_expr[0].expression for e in bound_partitions ], ) ), - output_names=_alias_or_inferred( - alias, - function, - [e.referred_expr[0].output_names[0] for e in bound_expressions], - ), + output_names=_alias_or_inferred(alias, function, arg_names), ) ], base_schema=base_schema, diff --git a/src/substrait/dataframe/__init__.py b/src/substrait/dataframe/__init__.py index a3d3f64..e8f9fe2 100644 --- a/src/substrait/dataframe/__init__.py +++ b/src/substrait/dataframe/__init__.py @@ -74,6 +74,7 @@ current_date, current_timestamp, current_timezone, + enum, exists, infer_literal_type, lit, @@ -128,6 +129,7 @@ "update_table", "DataFrame", "col", + "enum", "lit", "outer", "when", diff --git a/src/substrait/dataframe/expr.py b/src/substrait/dataframe/expr.py index f2ba9de..761f2a5 100644 --- a/src/substrait/dataframe/expr.py +++ b/src/substrait/dataframe/expr.py @@ -30,7 +30,9 @@ from substrait.builders import type as _t from substrait.builders.extended_expression import ( + EnumArg, UnboundExtendedExpression, + _function_signature, cast, column, if_then, @@ -246,18 +248,14 @@ def _resolve_over_urns( 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 - ] - match = registry.find_function(name, signature, urns, options=options) + signature = _function_signature(bound, 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] + kinds = [t if isinstance(t, str) else t.WhichOneof("kind") for t in signature] raise Exception( f"No matching overload for '{name}' across {urns} with signature {kinds}" ) @@ -1007,6 +1005,21 @@ def col(name: Union[str, int]) -> Expr: return Expr(column(name)) +def enum(value: str) -> EnumArg: + """A positional enumeration-argument selection for an ``f.*`` call. + + Some standard functions take an enumeration argument -- a positional operand + drawn from a fixed domain (e.g. ``extract``'s ``component``) that serializes + as a ``FunctionArgument.enum`` -- distinct from a behavioral option. Pass it + positionally, in declared argument order:: + + sub.f.extract(sub.enum("YEAR"), sub.col("d")) + + Mirrors substrait-go's ``types.Enum`` and substrait-java's ``EnumArg``. + """ + return EnumArg(value) + + def outer(name: Union[str, int], steps_out: int = 1) -> Expr: """Reference a column from an enclosing query (a correlated reference). diff --git a/src/substrait/dataframe/functions.py b/src/substrait/dataframe/functions.py index 9db205c..136e932 100644 --- a/src/substrait/dataframe/functions.py +++ b/src/substrait/dataframe/functions.py @@ -36,6 +36,7 @@ from typing import Any from substrait.builders.extended_expression import ( + EnumArg, aggregate_function, resolve_expression, scalar_function, @@ -63,7 +64,7 @@ def _urn_priority(urn: str) -> int: def _single_urn_helper(builder, urn: str, name: str): def helper(*args: Any, alias: str | None = None, **options: Any) -> Expr: - exprs = [Expr._coerce(a).unbound for a in args] + exprs = [a if isinstance(a, EnumArg) else Expr._coerce(a).unbound for a in args] return Expr( builder(urn, name, expressions=exprs, alias=alias, options=options or None) ) @@ -73,10 +74,15 @@ def helper(*args: Any, alias: str | None = None, **options: Any) -> Expr: def _multi_urn_helper(builder, urns: list[str], name: str): def helper(*args: Any, alias: str | None = None, **options: Any) -> Expr: - exprs = [Expr._coerce(a).unbound for a in args] + exprs = [a if isinstance(a, EnumArg) else Expr._coerce(a).unbound for a in args] def resolve(base_schema, registry): - bound = [resolve_expression(e, base_schema, registry) for e in exprs] + bound = [ + e + if isinstance(e, EnumArg) + else resolve_expression(e, base_schema, registry) + for e in exprs + ] return _resolve_over_urns( builder, urns, diff --git a/src/substrait/extension_registry/function_entry.py b/src/substrait/extension_registry/function_entry.py index 9cea4f1..d14d4bd 100644 --- a/src/substrait/extension_registry/function_entry.py +++ b/src/substrait/extension_registry/function_entry.py @@ -10,8 +10,6 @@ from .signature_checker_helpers import covers, normalize_substrait_type_names -_MISSING = object() - class FunctionType(Enum): SCALAR = "scalar" @@ -33,10 +31,6 @@ def __init__( self.urn: str = urn self.function_type = function_type self.arguments = [] - # Argument names parallel to ``arguments`` / ``normalized_inputs``. Used to - # match enumeration selections, which arrive by name (e.g. component="YEAR") - # rather than by position, back to their declared signature slot. - self.arg_names: list = [] self.nullability = ( impl.nullability if impl.nullability else se.NullabilityHandling.MIRROR ) @@ -47,99 +41,37 @@ def __init__( self.normalized_inputs.append( normalize_substrait_type_names(arg.value) ) - self.arg_names.append(arg.name) elif isinstance(arg, se.EnumerationArg): self.arguments.append(arg.options) self.normalized_inputs.append("req") - self.arg_names.append(arg.name) def __repr__(self) -> str: return f"{self.name}:{'_'.join(self.normalized_inputs)}" - def interleave_arguments(self, value_items, options): - """Order value operands and enumeration selections per the signature. - - Substrait interleaves enumeration arguments with value arguments in - declared order (``extract`` is ``[component (enum), x (value)]``), but the - two arrive separately: value operands positionally, enumeration selections - by argument name in ``options``. This walks the declared arguments and - weaves them back together. + def satisfies_signature(self, signature: tuple | list) -> Optional[Type]: + """Match ``signature`` against this overload, returning its output type. - ``value_items`` is an iterator of value operands (consumed in order for - each value argument). Returns ``(ordered, remaining_options)`` where each - entry of ``ordered`` is ``("value", item)`` or ``("enum", selection)``, and - ``remaining_options`` are the options *not* consumed as enumeration - arguments -- i.e. the behavioral options. Returns ``None`` if the value - arity is wrong or a required enumeration argument is absent from - ``options``. + ``signature`` interleaves value-operand ``Type``\\ s with enumeration + selections as plain string tokens, in declared argument order (an enum + argument's token must be a member of the overload's option domain). + Returns the derived output ``Type`` on a match, or ``None`` otherwise. """ - remaining = dict(options or {}) - if self.impl.variadic: - # Variadic functions have no enumeration arguments; every operand is a - # value and every option is behavioral. - return [("value", item) for item in value_items], remaining - ordered: list = [] - for kind, name in zip(self.normalized_inputs, self.arg_names): - if kind == "req": # enumeration argument - if name not in remaining: - return None - ordered.append(("enum", str(remaining.pop(name)))) - else: - item = next(value_items, _MISSING) - if item is _MISSING: - return None - ordered.append(("value", item)) - if next(value_items, _MISSING) is not _MISSING: - return None # too many value operands - return ordered, remaining - - def _resolve_signature( - self, signature: tuple | list, options: Optional[dict] - ) -> Optional[list]: - """Interleave ``signature`` with enumeration selections for matching. - - Enumeration selections may be supplied two ways: interleaved into - ``signature`` as strings in declared order (the low-level registry - contract), or by argument name in ``options`` (how the DataFrame builders - pass them, keeping the value-only signature intact). An enum position - prefers ``options`` and otherwise consumes the next ``signature`` item. - Returns the value-and-enum sequence to match against ``self.arguments``, - or ``None`` on an arity mismatch. - """ - remaining = dict(options or {}) - items = iter(signature) - interleaved: list = [] - for kind, name in zip(self.normalized_inputs, self.arg_names): - if kind == "req" and name in remaining: # enum selection by name - interleaved.append(str(remaining.pop(name))) - continue - item = next(items, _MISSING) - if item is _MISSING: - return None - interleaved.append(item) - if next(items, _MISSING) is not _MISSING: - return None # more operands than the signature declares - return interleaved - - def satisfies_signature( - self, signature: tuple | list, options: Optional[dict] = None - ) -> Optional[str]: if self.impl.variadic: min_args_allowed = self.impl.variadic.min or 0 if len(signature) < min_args_allowed: return None inputs = [self.arguments[0]] * len(signature) - interleaved: list = list(signature) else: - interleaved = self._resolve_signature(signature, options) - if interleaved is None: - return None inputs = self.arguments - if len(inputs) != len(interleaved): + if len(inputs) != len(signature): return None - zipped_args = list(zip(inputs, interleaved)) + zipped_args = list(zip(inputs, signature)) parameters = {} for x, y in zipped_args: + if isinstance(x, list) != isinstance(y, str): + # An enumeration slot (domain list) accepts only a selection token, + # and a value slot only a Type -- reject either kind in the other. + return None if isinstance(y, str): if y not in x: return None @@ -152,7 +84,10 @@ def satisfies_signature( == se.NullabilityHandling.DISCRETE, ): return None - output_type = evaluate(self.impl.return_, parameters) + try: + output_type = evaluate(self.impl.return_, parameters) + except Exception: + return None # return type cannot be derived for these arguments if self.nullability == se.NullabilityHandling.MIRROR and isinstance( output_type, Type ): @@ -160,7 +95,7 @@ def satisfies_signature( [ p.__getattribute__(p.WhichOneof("kind")).nullability == Type.NULLABILITY_NULLABLE - for p in interleaved + for p in signature if isinstance(p, Type) ] ) diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 9412aae..7eae5a0 100644 --- a/src/substrait/extension_registry/registry.py +++ b/src/substrait/extension_registry/registry.py @@ -116,14 +116,11 @@ def _find_matching_functions( function_name: str, signature: tuple[Type] | list[Type], urns: list[str] | None = None, - options: Optional[dict] = None, ) -> list[tuple[FunctionEntry, Type]]: """Helper method to find matching functions across specified URNs. - ``options`` carries enumeration-argument selections by name (e.g. - ``{"component": "YEAR"}``); an overload with enumeration arguments only - matches when ``signature`` supplies its value operands and ``options`` - supplies each enumeration selection. + ``signature`` interleaves value-operand ``Type``\\ s with enumeration + selections as plain string tokens, in declared argument order. """ matches = [] urns_to_search = ( @@ -137,7 +134,7 @@ def _find_matching_functions( continue functions = self._function_mapping[urn][function_name] for f in functions: - rtn = f.satisfies_signature(signature, options) + rtn = f.satisfies_signature(signature) if rtn is not None: matches.append((f, rtn)) return matches @@ -148,22 +145,24 @@ def lookup_function( urn: str, function_name: str, signature: tuple[Type] | list[Type], - options: Optional[dict] = None, ) -> Optional[tuple[FunctionEntry, Type]]: """Look up a function within a specific URN.""" - matches = self._find_matching_functions( - function_name, signature, [urn], options - ) + matches = self._find_matching_functions(function_name, signature, [urn]) return matches[0] if matches else None def list_functions( - self, urn: str, function_name: str, signature: tuple[Type] | list[Type] + self, + urn: str, + function_name: str, + signature: tuple[Type] | list[Type], ) -> list[tuple[FunctionEntry, Type]]: """List all matching functions within a specific URN.""" return self._find_matching_functions(function_name, signature, [urn]) def list_functions_across_urns( - self, function_name: str, signature: tuple[Type] | list[Type] + self, + function_name: str, + signature: tuple[Type] | list[Type], ) -> list[tuple[FunctionEntry, Type]]: """List all matching functions across all URNs.""" return self._find_matching_functions(function_name, signature) @@ -173,15 +172,12 @@ def find_function( function_name: str, signature: tuple[Type] | list[Type], urns: Optional[list[str]] = None, - options: Optional[dict] = 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``. - ``options`` carries enumeration-argument selections by name (see - :meth:`_find_matching_functions`). Generalizes :meth:`lookup_function` (a single URN) and :meth:`list_functions_across_urns` (every URN) to an ordered subset, so a @@ -189,7 +185,7 @@ def find_function( 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, options) + matches = self._find_matching_functions(function_name, signature, urns) return matches[0] if matches else None def has_urn(self, urn: str) -> bool: diff --git a/src/substrait/utils/display.py b/src/substrait/utils/display.py index e1e1841..0908e6c 100644 --- a/src/substrait/utils/display.py +++ b/src/substrait/utils/display.py @@ -654,10 +654,11 @@ def _get_function_argument_string(self, arg) -> str: # For nested scalar functions, we'll handle them specially in the main printing # Return a placeholder that indicates it needs recursive expansion return "" - elif arg.value.HasField("enum"): - return f"enum: {arg.value.enum}" else: return "" + elif arg.HasField("enum"): + # enum is a FunctionArgument field, not an Expression field. + return f"enum: {arg.enum}" else: return "" @@ -704,10 +705,11 @@ def _stream_function_argument(self, arg, stream, depth: int): stream.write(f"{indent}field: root\n") elif arg.value.HasField("scalar_function"): self._stream_scalar_function(arg.value.scalar_function, stream, depth) - elif arg.value.HasField("enum"): - stream.write(f"{indent}enum: {arg.value.enum}\n") else: stream.write(f"{indent}\n") + elif arg.HasField("enum"): + # enum is a FunctionArgument field, not an Expression field. + stream.write(f"{indent}enum: {arg.enum}\n") else: stream.write(f"{indent}\n") diff --git a/tests/dataframe/test_enum_arguments.py b/tests/dataframe/test_enum_arguments.py index 5cc4087..4ba09f2 100644 --- a/tests/dataframe/test_enum_arguments.py +++ b/tests/dataframe/test_enum_arguments.py @@ -14,35 +14,37 @@ return: i64 Per the spec these serialize into ``ScalarFunction.arguments`` as -``FunctionArgument.enum`` (interleaved with the value arguments, in signature -order) -- NOT into ``ScalarFunction.options``, which carries only *behavioral* -options (``overflow``, ``rounding``, ...). Consumers such as DuckDB read enum -selections exclusively from ``arguments``. +``FunctionArgument.enum`` (a positional operand, interleaved with the value +arguments in signature order) -- NOT into ``ScalarFunction.options``, which +carries only *behavioral* options (``overflow``, ``rounding``, ...). Consumers +such as DuckDB read enum selections exclusively from ``arguments``. + +The user-facing API mirrors the spec's positional model (and substrait-go's +``types.Enum`` / substrait-java's ``EnumArg``): an enumeration selection is a +positional ``sub.enum("...")`` marker written in declared argument order, while +behavioral options stay on the ``**options`` keyword channel. These tests guard two properties of the DataFrame / builder pipeline: -1. **Resolution.** ``sub.f.extract(col, component="YEAR")`` must build. The - overload's arity counts the enum positions, so the registry has to fold the - enum selection into the signature it matches -- otherwise the value-only - signature never matches and resolution raises ``Unknown function extract``. +1. **Resolution.** ``sub.f.extract(sub.enum("YEAR"), col)`` must build -- the + enum token interleaves into the match signature so the overload (which counts + the enum position toward its arity) resolves. 2. **Serialization.** The enum selection must land in ``arguments`` as a ``FunctionArgument.enum``, not in ``options``. Routing it into ``options`` produces a plan that omits the enum from ``arguments`` -- which DuckDB's consumer then reads as an empty enum vector and crashes on. - -The call site here -- enum selections passed as keyword arguments -- is the -user-facing API and is independent of how resolution is implemented internally. """ +import pytest import substrait.algebra_pb2 as stalg import substrait.dataframe as sub from substrait.builders.type import precision_timestamp -def _scalar_functions(message) -> list: - """Every ``Expression.ScalarFunction`` reachable anywhere in a proto tree.""" - target = stalg.Expression.ScalarFunction.DESCRIPTOR.full_name +def _messages_of(message, message_type) -> list: + """Every message of ``message_type`` reachable anywhere in a proto tree.""" + target = message_type.DESCRIPTOR.full_name found: list = [] def walk(msg): @@ -59,50 +61,58 @@ def walk(msg): return found -def _function_name(plan, scalar_function) -> str: - """Resolve a ScalarFunction's declared extension name via its anchor.""" - ref = scalar_function.function_reference +def _function_name(plan, function) -> str: + """The function's base extension name (the ``foo`` of the ``foo:sig`` anchor).""" + ref = function.function_reference for decl in plan.extensions: fn = decl.extension_function if fn.function_anchor == ref: - return fn.name + return fn.name.split(":", 1)[0] return "" -def _arg_kinds(scalar_function) -> list: - return [a.WhichOneof("arg_type") for a in scalar_function.arguments] +def _functions_named(plan, message_type, name: str) -> list: + """The ``message_type`` calls in ``plan`` whose base name is exactly ``name``. + + Exact on the base name (declarations are ``name:signature``): a prefix match + like ``startswith("extract")`` would also catch ``extract_boolean`` and break + the single-match unpacking below. + """ + return [ + m for m in _messages_of(plan, message_type) if _function_name(plan, m) == name + ] + + +def _arg_kinds(function) -> list: + return [a.WhichOneof("arg_type") for a in function.arguments] + + +def _options(function) -> list: + return [(o.name, list(o.preference)) for o in function.options] def test_extract_with_a_single_enum_argument_resolves(): """The two-argument ``extract(component, date)`` overload must build.""" df = sub.read_named_table("t", {"d": sub.date}) - # Must not raise "Unknown function extract". - plan = df.with_columns(y=sub.f.extract(sub.col("d"), component="YEAR")).to_plan() + # Must not raise "Unknown function extract"; unpacking asserts exactly one. + plan = df.with_columns(y=sub.f.extract(sub.enum("YEAR"), sub.col("d"))).to_plan() - (extract,) = [ - sf - for sf in _scalar_functions(plan) - if _function_name(plan, sf).startswith("extract") - ] - assert extract, "extract did not resolve to a registry function" + (extract,) = _functions_named(plan, stalg.Expression.ScalarFunction, "extract") + assert _arg_kinds(extract) == ["enum", "value"] def test_extract_serializes_the_enum_as_an_argument_not_an_option(): - """component=YEAR must be a FunctionArgument.enum, in signature position. + """sub.enum("YEAR") must be a FunctionArgument.enum, in signature position. Spec order for this overload is [component (enum), x (value)], so the positional value operand follows the enum selection. """ df = sub.read_named_table("t", {"d": sub.date}) - plan = df.with_columns(y=sub.f.extract(sub.col("d"), component="YEAR")).to_plan() + plan = df.with_columns(y=sub.f.extract(sub.enum("YEAR"), sub.col("d"))).to_plan() - (extract,) = [ - sf - for sf in _scalar_functions(plan) - if _function_name(plan, sf).startswith("extract") - ] + (extract,) = _functions_named(plan, stalg.Expression.ScalarFunction, "extract") assert _arg_kinds(extract) == ["enum", "value"], ( f"enum not carried in arguments in signature order: {_arg_kinds(extract)}" ) @@ -115,23 +125,116 @@ def test_extract_serializes_the_enum_as_an_argument_not_an_option(): def test_extract_day_of_month_carries_both_enum_arguments(): """The exact combination behind the reported DuckDB crash. - ``extract(component="DAY", indexing="ONE", )`` has two enum + ``extract(sub.enum("DAY"), sub.enum("ONE"), )`` has two enum arguments; both must appear in ``arguments`` (order [component, indexing, x]) and neither in ``options``. """ df = sub.read_named_table("t", {"ts": precision_timestamp(6)}) plan = df.with_columns( - dom=sub.f.extract(sub.col("ts"), component="DAY", indexing="ONE") + dom=sub.f.extract(sub.enum("DAY"), sub.enum("ONE"), sub.col("ts")) ).to_plan() - (extract,) = [ - sf - for sf in _scalar_functions(plan) - if _function_name(plan, sf).startswith("extract") - ] + (extract,) = _functions_named(plan, stalg.Expression.ScalarFunction, "extract") assert _arg_kinds(extract) == ["enum", "enum", "value"], _arg_kinds(extract) assert [extract.arguments[0].enum, extract.arguments[1].enum] == ["DAY", "ONE"] assert list(extract.options) == [], ( f"enum selections leaked into options (crashes DuckDB): {extract.options}" ) + + +def test_two_enum_selections_yield_distinct_inferred_output_names(): + """Enum selections must feed the inferred alias, else projections collide. + + ``extract(sub.enum("YEAR"), d)`` and ``extract(sub.enum("ISO_YEAR"), d)`` + differ only in the enum selection; if the inferred name ignored it both + columns would be named ``extract(d)`` and clash in one projection. + """ + df = sub.read_named_table("t", {"d": sub.date}) + + plan = df.select( + sub.f.extract(sub.enum("YEAR"), sub.col("d")), + sub.f.extract(sub.enum("ISO_YEAR"), sub.col("d")), + ).to_plan() + + names = list(plan.relations[0].root.names) + assert len(set(names)) == len(names), f"inferred output names collided: {names}" + + +def test_std_dev_resolves_to_enum_overload_not_deprecated_value_only(): + """sub.enum("POPULATION") must select the enum overload. + + ``std_dev`` declares both a deprecated value-only overload (``std_dev:fp64``, + one argument) and an enum-argument one (``std_dev:req_fp64``, two arguments). + Supplying the enum selection positionally makes the call two operands wide, so + it matches the enum overload by arity and the selection lands in ``arguments`` + as an enum rather than becoming a behavioral option. + """ + df = sub.read_named_table("t", {"x": sub.fp64}) + + plan = ( + df.group_by() + .agg(sub.f.std_dev(sub.enum("POPULATION"), sub.col("x")).alias("s")) + .to_plan() + ) + + (std_dev,) = _functions_named(plan, stalg.AggregateFunction, "std_dev") + assert _arg_kinds(std_dev) == ["enum", "value"], _arg_kinds(std_dev) + assert std_dev.arguments[0].enum == "POPULATION" + assert list(std_dev.options) == [], ( + f"enum selection leaked into options: {std_dev.options}" + ) + + +def test_median_splits_enum_argument_from_behavioral_option(): + """A call carrying both kinds at once -- the premise of this change. + + ``median`` takes an enum argument ``precision`` *and* a behavioral option + ``rounding``. The positional ``sub.enum("EXACT")`` must land in ``arguments`` + as an enum, the keyword ``rounding`` in ``options`` as a FunctionOption; they + must not swap channels. + """ + df = sub.read_named_table("t", {"x": sub.fp64}) + + plan = ( + df.group_by() + .agg( + sub.f.median(sub.enum("EXACT"), sub.col("x"), rounding="TRUNCATE").alias( + "m" + ) + ) + .to_plan() + ) + + (median,) = _functions_named(plan, stalg.AggregateFunction, "median") + assert _arg_kinds(median) == ["enum", "value"], _arg_kinds(median) + assert median.arguments[0].enum == "EXACT" + assert _options(median) == [("rounding", ["TRUNCATE"])], _options(median) + + +def test_bare_string_is_not_an_enum_selection(): + """Only ``sub.enum(...)`` marks an enum; a bare string is a literal operand. + + ``extract(sub.col("d"), "YEAR")`` coerces ``"YEAR"`` to a string literal, so + no overload matches ([date, string] value signature) and resolution raises + -- documenting why the positional marker is required. + """ + df = sub.read_named_table("t", {"d": sub.date}) + + with pytest.raises(Exception, match="[Uu]nknown function|No matching overload"): + df.with_columns(y=sub.f.extract(sub.col("d"), "YEAR")).to_plan() + + +def test_window_call_site_still_emits_value_arguments(): + """The window builder shares argument assembly; a value operand must survive. + + No standard window function takes an enumeration argument, so this guards the + value-only path through the changed ``_function_arguments`` return for the + window call site rather than enum handling. + """ + df = sub.read_named_table("t", {"a": sub.i64}) + + plan = df.select(sub.f.ntile(sub.lit(4)).over(order_by="a")).to_plan() + + (ntile,) = _functions_named(plan, stalg.Expression.WindowFunction, "ntile") + assert _arg_kinds(ntile) == ["value"], _arg_kinds(ntile)