feat: implement enum function arguments

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,6 +107,48 @@ def _function_options(options):
     return result
 
 
+def _function_signature(bound_expressions, registry):
+    """The flattened match signature for a call's positional arguments.
+
+    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.
+    """
+    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(
     expression: ExtendedExpressionOrUnbound,
     base_schema: stp.NamedStruct,
@@ -571,15 +637,13 @@ def resolve(
         base_schema: stp.NamedStruct, registry: ExtensionRegistry
     ) -> stee.ExtendedExpression:
         bound_expressions = [
-            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
         ]
 
-        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)
 
@@ -588,27 +652,20 @@ def resolve(
 
         func_ref = function_reference(urn, str(func[0]))
 
+        arguments, arg_names = _function_arguments(bound_expressions)
+
         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
-                            ],
+                            arguments=arguments,
                             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,
@@ -640,19 +697,17 @@ 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 = [typ for es in expression_schemas for typ in es.types]
+        signature = _function_signature(bound_expressions, registry)
 
         func = registry.lookup_function(urn, function, signature)
 
@@ -661,15 +716,14 @@ def resolve(
 
         func_ref = function_reference(urn, str(func[0]))
 
+        arguments, arg_names = _function_arguments(bound_expressions)
+
         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
-                        ],
+                        arguments=arguments,
                         options=_function_options(options),
                         output_type=func[1],
                         invocation=invocation
@@ -682,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,
@@ -710,19 +760,17 @@ 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)
 
@@ -731,30 +779,23 @@ def resolve(
 
         func_ref = function_reference(urn, str(func[0]))
 
+        arguments, arg_names = _function_arguments(bound_expressions)
+
         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
-                            ],
+                            arguments=arguments,
                             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,

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

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
-    ]
+    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).
 

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,

src/substrait/extension_registry/function_entry.py

@@ -48,7 +48,14 @@ def __init__(
     def __repr__(self) -> str:
         return f"{self.name}:{'_'.join(self.normalized_inputs)}"
 
-    def satisfies_signature(self, signature: tuple | list) -> Optional[str]:
+    def satisfies_signature(self, signature: tuple | list) -> Optional[Type]:
+        """Match ``signature`` against this overload, returning its output type.
+
+        ``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.
+        """
         if self.impl.variadic:
             min_args_allowed = self.impl.variadic.min or 0
             if len(signature) < min_args_allowed:
@@ -61,6 +68,10 @@ def satisfies_signature(self, signature: tuple | list) -> Optional[str]:
         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
@@ -73,7 +84,10 @@ def satisfies_signature(self, signature: tuple | list) -> Optional[str]:
                     == 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
         ):