diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index 2738a55..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,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,18 +779,15 @@ 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=[ @@ -750,11 +795,7 @@ def resolve( ], ) ), - 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 c185a37..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 - ] + 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 b2bd0d5..d14d4bd 100644 --- a/src/substrait/extension_registry/function_entry.py +++ b/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 ): diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 5ad082f..7eae5a0 100644 --- a/src/substrait/extension_registry/registry.py +++ b/src/substrait/extension_registry/registry.py @@ -117,7 +117,11 @@ def _find_matching_functions( signature: tuple[Type] | list[Type], urns: list[str] | None = None, ) -> list[tuple[FunctionEntry, Type]]: - """Helper method to find matching functions across specified URNs.""" + """Helper method to find matching functions across specified URNs. + + ``signature`` interleaves value-operand ``Type``\\ s with enumeration + selections as plain string tokens, in declared argument order. + """ matches = [] urns_to_search = ( urns if urns is not None else list(self._function_mapping.keys()) @@ -147,13 +151,18 @@ def lookup_function( 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) 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 new file mode 100644 index 0000000..4ba09f2 --- /dev/null +++ b/tests/dataframe/test_enum_arguments.py @@ -0,0 +1,240 @@ +"""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`` (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(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. +""" + +import pytest +import substrait.algebra_pb2 as stalg + +import substrait.dataframe as sub +from substrait.builders.type import precision_timestamp + + +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): + 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, 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.split(":", 1)[0] + return "" + + +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"; unpacking asserts exactly one. + plan = df.with_columns(y=sub.f.extract(sub.enum("YEAR"), sub.col("d"))).to_plan() + + (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(): + """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.enum("YEAR"), sub.col("d"))).to_plan() + + (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)}" + ) + assert extract.arguments[0].enum == "YEAR" + assert list(extract.options) == [], ( + f"enum selection leaked into options (crashes DuckDB): {extract.options}" + ) + + +def test_extract_day_of_month_carries_both_enum_arguments(): + """The exact combination behind the reported DuckDB crash. + + ``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.enum("DAY"), sub.enum("ONE"), sub.col("ts")) + ).to_plan() + + (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)