diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0362b343..ed692035 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -37,7 +37,11 @@ repos: hooks: - id: mypy entry: mypy - files: ^(duckdb/|_duckdb-stubs/) + # Check the package and its stubs together rather than the staged files: a list without a stub file hides + # _duckdb from mypy, and the hook's default --ignore-missing-imports then turns it into Any. + pass_filenames: false + args: [ duckdb, _duckdb-stubs ] + files: ^(duckdb/|_duckdb-stubs/|pyproject\.toml$) exclude: ^duckdb/(experimental|query_graph)/ additional_dependencies: [ numpy, polars ] diff --git a/duckdb/_expression.py b/duckdb/_expression.py index 6677c136..6835bb27 100644 --- a/duckdb/_expression.py +++ b/duckdb/_expression.py @@ -1151,6 +1151,26 @@ def where(self, predicate: Expression) -> Expression: raise TypeError(message) return self._with(_filter=_as_condition(predicate, "where")) + def order_by(self, *keys: Expression) -> Expression: + """Feed the rows to the aggregate in this order, as `first("x" ORDER BY "t")`; goes before `.over()`. + + Give a key a direction and a null order with `.desc()`, `.nulls_last()` and their siblings. A later + `.order_by()` replaces the keys. As with `where()`, a scalar call fails at run time, not here. + """ + if not isinstance(self, (Func, Distinct)): + message = "order_by() applies to an aggregate call, as col('x').first() or fn('first', ...)" + raise TypeError(message) + if not keys: + message = "order_by() takes at least one key" + raise TypeError(message) + return self._with(_order=_as_expressions(keys, "order_by")) + + def _ordered(self, keys: Exprs | None) -> Expression: + """`.order_by()` from the `order_by=` keyword of an aggregate method, where it was given.""" + if keys is None: + return self + return self.order_by(*(keys if isinstance(keys, (list, tuple)) else (keys,))) + def over( self, partition_by: Exprs | None = None, @@ -1279,6 +1299,10 @@ def ignore_nulls(self) -> Expression: """`Expression.ignore_nulls`, applied to the wrapped function call.""" return self._rewrapped(self.inner.ignore_nulls()) + def order_by(self, *keys: Expression) -> Expression: + """`Expression.order_by`, applied to the wrapped aggregate call.""" + return self._rewrapped(self.inner.order_by(*keys)) + def _rewrapped(self, inner: Expression) -> Expression: wrapped = type(self)(inner) wrapped._alias = self._alias @@ -1522,6 +1546,7 @@ def __init__( self.args = args self.named = named or [] self._ignore_nulls = False + self._order: list[Expression] = [] self._filter: Expression | None = None def fragment(self) -> str: @@ -1529,8 +1554,9 @@ def fragment(self) -> str: [_argument(a, fold_name(self.name[-1])) for a in self.args] + [f"{quote(k)} := {v.fragment()}" for k, v in self.named] ) - # IGNORE NULLS goes inside the call, where DuckDB reads it; after the closing parenthesis it is a syntax error. - tail = " IGNORE NULLS" if self._ignore_nulls else "" + # ORDER BY and then IGNORE NULLS go inside the call, where DuckDB reads them; the other order, or either + # after the closing parenthesis, is a syntax error. + tail = _order_clause(self._order) + (" IGNORE NULLS" if self._ignore_nulls else "") return f"{function_name(self.name)}({rendered}{tail})" + _filter_clause(self._filter) @@ -1552,10 +1578,12 @@ def __init__(self, name: str | tuple[str, ...], operand: Expression) -> None: super().__init__() self.name = name_parts(name, "function name") self.operand = operand + self._order: list[Expression] = [] self._filter: Expression | None = None def fragment(self) -> str: - return f"{function_name(self.name)}(DISTINCT {self.operand.fragment()})" + _filter_clause(self._filter) + inner = f"DISTINCT {self.operand.fragment()}{_order_clause(self._order)}" + return f"{function_name(self.name)}({inner})" + _filter_clause(self._filter) #: The functions that read an argument's name, so an alias on one of their arguments is written as `name := expr`. @@ -1576,6 +1604,11 @@ def _filter_clause(predicate: Expression | None) -> str: return f" FILTER (WHERE {predicate.fragment()})" if predicate is not None else "" +def _order_clause(keys: list[Expression]) -> str: + """The ORDER BY inside an aggregate call, if it has one.""" + return " ORDER BY " + ", ".join(k.as_order() for k in keys) if keys else "" + + class Concat(Expression): """Text concatenation, which SQL spells `||`.""" diff --git a/duckdb/_expression_aggregates.py b/duckdb/_expression_aggregates.py index 29349ce5..c3fe956b 100644 --- a/duckdb/_expression_aggregates.py +++ b/duckdb/_expression_aggregates.py @@ -5,11 +5,11 @@ from typing import TYPE_CHECKING, ClassVar if TYPE_CHECKING: - from duckdb._expression import Expression, ValueOrExpr + from duckdb._expression import Expression, Exprs, ValueOrExpr class AggregateMethods: - """The aggregate functions as methods; `.over()` makes each a window function, and `fn()` reaches the rest.""" + """The aggregate functions as methods, with `order_by=` for `.order_by()`; `.over()` makes each a window.""" #: Method name to SQL function, the table these methods were written from. _CALLABLES: ClassVar[dict[str, str]] = { @@ -45,110 +45,110 @@ class AggregateMethods: def _call(self, function: str, *args: object) -> Expression: # pragma: no cover (defined there) raise NotImplementedError - def sum(self, *args: ValueOrExpr) -> Expression: + def sum(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The sum of the values. SQL: `sum`.""" - return self._call("sum", *args) + return self._call("sum", *args)._ordered(order_by) - def mean(self, *args: ValueOrExpr) -> Expression: + def mean(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The average of the values. SQL: `avg`.""" - return self._call("avg", *args) + return self._call("avg", *args)._ordered(order_by) - def avg(self, *args: ValueOrExpr) -> Expression: + def avg(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The average of the values. The same as `mean`. SQL: `avg`.""" - return self._call("avg", *args) + return self._call("avg", *args)._ordered(order_by) - def min(self, *args: ValueOrExpr) -> Expression: + def min(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The smallest value. SQL: `min`.""" - return self._call("min", *args) + return self._call("min", *args)._ordered(order_by) - def max(self, *args: ValueOrExpr) -> Expression: + def max(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The largest value. SQL: `max`.""" - return self._call("max", *args) + return self._call("max", *args)._ordered(order_by) - def count(self, *args: ValueOrExpr) -> Expression: + def count(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """How many values are not NULL. For rows, use `count_all()`. SQL: `count`.""" - return self._call("count", *args) + return self._call("count", *args)._ordered(order_by) - def median(self, *args: ValueOrExpr) -> Expression: + def median(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The middle value. SQL: `median`.""" - return self._call("median", *args) + return self._call("median", *args)._ordered(order_by) - def std(self, *args: ValueOrExpr) -> Expression: + def std(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The sample standard deviation. SQL: `stddev_samp`.""" - return self._call("stddev_samp", *args) + return self._call("stddev_samp", *args)._ordered(order_by) - def var(self, *args: ValueOrExpr) -> Expression: + def var(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The sample variance. SQL: `var_samp`.""" - return self._call("var_samp", *args) + return self._call("var_samp", *args)._ordered(order_by) - def first(self, *args: ValueOrExpr) -> Expression: - """The first value seen. Order it with `.over()` or `sort` to make that mean something. SQL: `first`.""" - return self._call("first", *args) + def first(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: + """The first value; give `order_by=` to say which, else whichever the engine met first. SQL: `first`.""" + return self._call("first", *args)._ordered(order_by) - def last(self, *args: ValueOrExpr) -> Expression: - """The last value seen. SQL: `last`.""" - return self._call("last", *args) + def last(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: + """The last value; give `order_by=` to say which, else whichever the engine met last. SQL: `last`.""" + return self._call("last", *args)._ordered(order_by) - def any_value(self, *args: ValueOrExpr) -> Expression: - """Any one of the values, whichever is cheapest. SQL: `any_value`.""" - return self._call("any_value", *args) + def any_value(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: + """Any one of the values, whichever is cheapest; with `order_by=`, the first non-NULL one. SQL: `any_value`.""" + return self._call("any_value", *args)._ordered(order_by) - def bit_and(self, *args: ValueOrExpr) -> Expression: + def bit_and(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The bitwise AND of the values. SQL: `bit_and`.""" - return self._call("bit_and", *args) + return self._call("bit_and", *args)._ordered(order_by) - def bit_or(self, *args: ValueOrExpr) -> Expression: + def bit_or(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The bitwise OR of the values. SQL: `bit_or`.""" - return self._call("bit_or", *args) + return self._call("bit_or", *args)._ordered(order_by) - def bool_and(self, *args: ValueOrExpr) -> Expression: + def bool_and(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """Whether every value is true. SQL: `bool_and`.""" - return self._call("bool_and", *args) + return self._call("bool_and", *args)._ordered(order_by) - def bool_or(self, *args: ValueOrExpr) -> Expression: + def bool_or(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """Whether any value is true. SQL: `bool_or`.""" - return self._call("bool_or", *args) + return self._call("bool_or", *args)._ordered(order_by) - def product(self, *args: ValueOrExpr) -> Expression: + def product(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The product of the values. SQL: `product`.""" - return self._call("product", *args) + return self._call("product", *args)._ordered(order_by) - def string_agg(self, *args: ValueOrExpr) -> Expression: + def string_agg(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The values joined as text. Pass the separator as an argument. SQL: `string_agg`.""" - return self._call("string_agg", *args) + return self._call("string_agg", *args)._ordered(order_by) - def skewness(self, *args: ValueOrExpr) -> Expression: + def skewness(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The skewness of the distribution. SQL: `skewness`.""" - return self._call("skewness", *args) + return self._call("skewness", *args)._ordered(order_by) - def kurtosis(self, *args: ValueOrExpr) -> Expression: + def kurtosis(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The kurtosis of the distribution. SQL: `kurtosis`.""" - return self._call("kurtosis", *args) + return self._call("kurtosis", *args)._ordered(order_by) - def entropy(self, *args: ValueOrExpr) -> Expression: + def entropy(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The entropy of the distribution. SQL: `entropy`.""" - return self._call("entropy", *args) + return self._call("entropy", *args)._ordered(order_by) - def array_agg(self, *args: ValueOrExpr) -> Expression: + def array_agg(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The values gathered into a list. SQL: `array_agg`.""" - return self._call("array_agg", *args) + return self._call("array_agg", *args)._ordered(order_by) - def mode(self, *args: ValueOrExpr) -> Expression: + def mode(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The most frequent value. SQL: `mode`.""" - return self._call("mode", *args) + return self._call("mode", *args)._ordered(order_by) - def quantile(self, *args: ValueOrExpr) -> Expression: + def quantile(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The value at a quantile, one of the values themselves. Pass the quantile, 0 to 1. SQL: `quantile_disc`.""" - return self._call("quantile_disc", *args) + return self._call("quantile_disc", *args)._ordered(order_by) - def quantile_cont(self, *args: ValueOrExpr) -> Expression: + def quantile_cont(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The value at a quantile, interpolated. Pass the quantile, 0 to 1. SQL: `quantile_cont`.""" - return self._call("quantile_cont", *args) + return self._call("quantile_cont", *args)._ordered(order_by) - def arg_min(self, *args: ValueOrExpr) -> Expression: + def arg_min(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The value of the argument where this is smallest. SQL: `arg_min`.""" - return self._call("arg_min", *args) + return self._call("arg_min", *args)._ordered(order_by) - def arg_max(self, *args: ValueOrExpr) -> Expression: + def arg_max(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression: """The value of the argument where this is largest. SQL: `arg_max`.""" - return self._call("arg_max", *args) + return self._call("arg_max", *args)._ordered(order_by) diff --git a/duckdb/_expression_namespaces.py b/duckdb/_expression_namespaces.py index 6f1115e1..6fbb0ba0 100644 --- a/duckdb/_expression_namespaces.py +++ b/duckdb/_expression_namespaces.py @@ -638,7 +638,7 @@ def aggregate(self, *args: ValueOrExpr) -> Expression: """The named aggregate applied to the elements. Arguments: function_name. SQL: `list_aggregate`.""" return self._call("list_aggregate", *args) - def any_value(self, *args: ValueOrExpr) -> Expression: + def any_value(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """Any one element, whichever is cheapest. SQL: `list_any_value`.""" return self._call("list_any_value", *args) @@ -650,15 +650,15 @@ def approx_count_distinct(self, *args: ValueOrExpr) -> Expression: """An approximate count of the distinct elements. SQL: `list_approx_count_distinct`.""" return self._call("list_approx_count_distinct", *args) - def avg(self, *args: ValueOrExpr) -> Expression: + def avg(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The average of the elements. SQL: `list_avg`.""" return self._call("list_avg", *args) - def bit_and(self, *args: ValueOrExpr) -> Expression: + def bit_and(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The bitwise AND of the elements. SQL: `list_bit_and`.""" return self._call("list_bit_and", *args) - def bit_or(self, *args: ValueOrExpr) -> Expression: + def bit_or(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The bitwise OR of the elements. SQL: `list_bit_or`.""" return self._call("list_bit_or", *args) @@ -666,11 +666,11 @@ def bit_xor(self, *args: ValueOrExpr) -> Expression: """The bitwise XOR of the elements. SQL: `list_bit_xor`.""" return self._call("list_bit_xor", *args) - def bool_and(self, *args: ValueOrExpr) -> Expression: + def bool_and(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """Whether every element is true. SQL: `list_bool_and`.""" return self._call("list_bool_and", *args) - def bool_or(self, *args: ValueOrExpr) -> Expression: + def bool_or(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """Whether any element is true. SQL: `list_bool_or`.""" return self._call("list_bool_or", *args) @@ -686,7 +686,7 @@ def cosine_similarity(self, *args: ValueOrExpr) -> Expression: """The cosine similarity of two lists of numbers. Arguments: list2. SQL: `list_cosine_similarity`.""" return self._call("list_cosine_similarity", *args) - def count(self, *args: ValueOrExpr) -> Expression: + def count(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """How many elements are not NULL. SQL: `list_count`.""" return self._call("list_count", *args) @@ -698,7 +698,7 @@ def distinct(self, *args: ValueOrExpr) -> ListExpr: """Removes all duplicates and `NULL` values from a list. SQL: `list_distinct`.""" return ListExpr(self._call("list_distinct", *args)) - def entropy(self, *args: ValueOrExpr) -> Expression: + def entropy(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The entropy of the elements. SQL: `list_entropy`.""" return self._call("list_entropy", *args) @@ -710,7 +710,7 @@ def filter(self, function: ValueOrExpr, *args: ValueOrExpr) -> ListExpr: """The elements the function holds for: `lambda x: x > 0`. SQL: `list_filter`.""" return ListExpr(self._call("list_filter", function, *args)) - def first(self, *args: ValueOrExpr) -> Expression: + def first(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The first element. SQL: `list_first`.""" return self._call("list_first", *args) @@ -742,7 +742,7 @@ def intersect(self, *args: ValueOrExpr) -> ListExpr: """The distinct elements present in both lists. Arguments: list2. SQL: `list_intersect`.""" return ListExpr(self._call("list_intersect", *args)) - def kurtosis(self, *args: ValueOrExpr) -> Expression: + def kurtosis(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The kurtosis of the elements. SQL: `list_kurtosis`.""" return self._call("list_kurtosis", *args) @@ -750,7 +750,7 @@ def kurtosis_pop(self, *args: ValueOrExpr) -> Expression: """The population kurtosis of the elements. SQL: `list_kurtosis_pop`.""" return self._call("list_kurtosis_pop", *args) - def last(self, *args: ValueOrExpr) -> Expression: + def last(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The last element. SQL: `list_last`.""" return self._call("list_last", *args) @@ -762,19 +762,19 @@ def mask(self, *args: ValueOrExpr) -> ListExpr: """The elements where the matching boolean is true. Arguments: mask_list. SQL: `list_where`.""" return ListExpr(self._call("list_where", *args)) - def max(self, *args: ValueOrExpr) -> Expression: + def max(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The largest element. SQL: `list_max`.""" return self._call("list_max", *args) - def median(self, *args: ValueOrExpr) -> Expression: + def median(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The middle element. SQL: `list_median`.""" return self._call("list_median", *args) - def min(self, *args: ValueOrExpr) -> Expression: + def min(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The smallest element. SQL: `list_min`.""" return self._call("list_min", *args) - def mode(self, *args: ValueOrExpr) -> Expression: + def mode(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The most frequent element. SQL: `list_mode`.""" return self._call("list_mode", *args) @@ -790,7 +790,7 @@ def prepend(self, e: ValueOrExpr, *args: ValueOrExpr) -> ListExpr: """The `list_prepend` function. Arguments: e. SQL: `list_prepend`.""" return ListExpr(self._call_at("list_prepend", 1, e, *args)) - def product(self, *args: ValueOrExpr) -> Expression: + def product(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The product of the elements. SQL: `list_product`.""" return self._call("list_product", *args) @@ -818,7 +818,7 @@ def sem(self, *args: ValueOrExpr) -> Expression: """The standard error of the mean of the elements. SQL: `list_sem`.""" return self._call("list_sem", *args) - def skewness(self, *args: ValueOrExpr) -> Expression: + def skewness(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The skewness of the elements. SQL: `list_skewness`.""" return self._call("list_skewness", *args) @@ -834,11 +834,11 @@ def stddev_samp(self, *args: ValueOrExpr) -> Expression: """The sample standard deviation of the elements. SQL: `list_stddev_samp`.""" return self._call("list_stddev_samp", *args) - def string_agg(self, *args: ValueOrExpr) -> StrExpr: + def string_agg(self, *args: ValueOrExpr) -> StrExpr: # type: ignore[override] """The elements joined as text, separated by commas. SQL: `list_string_agg`.""" return StrExpr(self._call("list_string_agg", *args)) - def sum(self, *args: ValueOrExpr) -> Expression: + def sum(self, *args: ValueOrExpr) -> Expression: # type: ignore[override] """The sum of the elements. SQL: `list_sum`.""" return self._call("list_sum", *args) diff --git a/pyproject.toml b/pyproject.toml index ca3d9fff..af213921 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -403,6 +403,7 @@ ignore_errors = true [[tool.mypy.overrides]] module = [ "fsspec.*", + "IPython.*", "pandas", "polars", "pyarrow.*", diff --git a/scripts/func_namespaces.toml b/scripts/func_namespaces.toml index 9ac5a79f..a2183a51 100644 --- a/scripts/func_namespaces.toml +++ b/scripts/func_namespaces.toml @@ -293,9 +293,9 @@ count = { sql = "count", doc = "How many values are not NULL. For rows, use `cou median = { sql = "median", doc = "The middle value." } std = { sql = "stddev_samp", doc = "The sample standard deviation." } var = { sql = "var_samp", doc = "The sample variance." } -first = { sql = "first", doc = "The first value seen. Order it with `.over()` or `sort` to make that mean something." } -last = { sql = "last", doc = "The last value seen." } -any_value = { sql = "any_value", doc = "Any one of the values, whichever is cheapest." } +first = { sql = "first", doc = "The first value; give `order_by=` to say which, else whichever the engine met first." } +last = { sql = "last", doc = "The last value; give `order_by=` to say which, else whichever the engine met last." } +any_value = { sql = "any_value", doc = "Any one of the values, whichever is cheapest; with `order_by=`, the first non-NULL one." } bit_and = { sql = "bit_and", doc = "The bitwise AND of the values." } bit_or = { sql = "bit_or", doc = "The bitwise OR of the values." } bool_and = { sql = "bool_and", doc = "Whether every value is true." } diff --git a/scripts/gen_aggregates.py b/scripts/gen_aggregates.py index 12ee8f34..521006cb 100644 --- a/scripts/gen_aggregates.py +++ b/scripts/gen_aggregates.py @@ -27,11 +27,11 @@ def load_aggregates() -> dict[str, tuple[str, str]]: from typing import TYPE_CHECKING, ClassVar if TYPE_CHECKING: - from duckdb._expression import Expression, ValueOrExpr + from duckdb._expression import Expression, Exprs, ValueOrExpr class AggregateMethods: - """The aggregate functions as methods; `.over()` makes each a window function, and `fn()` reaches the rest.""" + """The aggregate functions as methods, with `order_by=` for `.order_by()`; `.over()` makes each a window.""" #: Method name to SQL function, the table these methods were written from. _CALLABLES: ClassVar[dict[str, str]] = { @@ -47,9 +47,9 @@ def render() -> str: out.append(" def _call(self, function: str, *args: object) -> Expression: # pragma: no cover (defined there)\n") out.append(" raise NotImplementedError\n") for method, (function, doc) in AGGREGATES.items(): - out.append(f"\n def {method}(self, *args: ValueOrExpr) -> Expression:\n") + out.append(f"\n def {method}(self, *args: ValueOrExpr, order_by: Exprs | None = None) -> Expression:\n") out.append(f' """{doc} SQL: `{function}`."""\n') - out.append(f' return self._call("{function}", *args)\n') + out.append(f' return self._call("{function}", *args)._ordered(order_by)\n') return "".join(out) diff --git a/scripts/gen_func_namespaces.py b/scripts/gen_func_namespaces.py index 332c9c38..2921f461 100644 --- a/scripts/gen_func_namespaces.py +++ b/scripts/gen_func_namespaces.py @@ -252,6 +252,9 @@ def build( def render(resolved: dict[str, Any]) -> str: """The module text.""" + from duckdb._expression_aggregates import AggregateMethods + + aggregates = set(AggregateMethods._CALLABLES) out = [HEADER] for family, cls in CLASSES.items(): section = resolved[family] @@ -264,16 +267,19 @@ def render(resolved: dict[str, Any]) -> str: out.append(f' "{method}": ("{m["function"]}", {m["subject"]}, {types}),\n') out.append(" }\n") for method, m in section["methods"].items(): - out.append(_method(method, m)) + out.append(_method(method, m, aggregates)) return "".join(out) -def _method(method: str, m: dict[str, Any]) -> str: - """One generated method.""" +def _method(method: str, m: dict[str, Any], aggregates: set[str]) -> str: + """One generated method; `aggregates` names the aggregate methods a family method may take over.""" wrap = m["returns"] tail = f" SQL: `{m['function']}`." + # A family method that takes over an aggregate's name, as `.list().sum()` is `list_sum`, is a scalar call with + # no rows to order, so it has no `order_by=`; mypy reads that as a narrower override. + override = " # type: ignore[override]" if method in aggregates else "" if m["lambda"]: - head = f"\n def {method}(self, function: ValueOrExpr, *args: ValueOrExpr) -> {wrap}:\n" + head = f"\n def {method}(self, function: ValueOrExpr, *args: ValueOrExpr) -> {wrap}:{override}\n" call = f'self._call("{m["function"]}", function, *args)' else: if m["params"]: @@ -281,11 +287,11 @@ def _method(method: str, m: dict[str, Any]) -> str: subject = m["subject"] if subject > 0: leading = ", ".join(f"{p}: ValueOrExpr" for p in m["params"][:subject]) - head = f"\n def {method}(self, {leading}, *args: ValueOrExpr) -> {wrap}:\n" + head = f"\n def {method}(self, {leading}, *args: ValueOrExpr) -> {wrap}:{override}\n" names = ", ".join(m["params"][:subject]) call = f'self._call_at("{m["function"]}", {subject}, {names}, *args)' else: - head = f"\n def {method}(self, *args: ValueOrExpr) -> {wrap}:\n" + head = f"\n def {method}(self, *args: ValueOrExpr) -> {wrap}:{override}\n" call = f'self._call("{m["function"]}", *args)' room = LINE_LIMIT - len(' """') - len(tail) - len('"""') doc = f' """{_fit(_sentence(m["doc"]), room)}{tail}"""\n' diff --git a/tests/fast/expression/test_ordered_aggregates.py b/tests/fast/expression/test_ordered_aggregates.py new file mode 100644 index 00000000..424d12f1 --- /dev/null +++ b/tests/fast/expression/test_ordered_aggregates.py @@ -0,0 +1,191 @@ +"""An aggregate call takes an ORDER BY inside its parentheses, so `first(x ORDER BY t)` needs no SQL text.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +import duckdb +from duckdb._expression import col, count_all, fn, lit, row_number + +if TYPE_CHECKING: + from duckdb._expression import Expr + + +@pytest.fixture(scope="module") +def con() -> duckdb.DuckDBPyConnection: + connection = duckdb.connect() + connection.execute( + "CREATE TABLE t AS SELECT * FROM (VALUES " + "(1, 3, 'a'), (1, 1, 'b'), (1, 2, NULL), (1, NULL, 'd'), (2, 5, 'c'), (2, 5, 'c')" + ") v(g, t, x)" + ) + yield connection + connection.close() + + +def grouped(con: duckdb.DuckDBPyConnection, expression: Expr, where: str = "") -> list[tuple]: + return con.execute(f"SELECT g, {expression.fragment()} FROM t {where} GROUP BY g ORDER BY g").fetchall() + + +def answers(con: duckdb.DuckDBPyConnection, sql: str, where: str = "") -> list[tuple]: + return con.execute(f"SELECT g, {sql} FROM t {where} GROUP BY g ORDER BY g").fetchall() + + +class TestRendering: + def test_renders_inside_the_call(self) -> None: + assert col("x").first().order_by(col("t")).fragment() == '"first"("x" ORDER BY "t")' + + def test_renders_with_distinct(self) -> None: + assert col("x").n_unique().order_by(col("x")).fragment() == 'count(DISTINCT "x" ORDER BY "x")' + + def test_renders_before_filter(self) -> None: + call = col("x").first().order_by(col("t")).where(col("t") > 1) + assert call.fragment() == '"first"("x" ORDER BY "t") FILTER (WHERE ("t" > 1))' + reversed_order = col("x").first().where(col("t") > 1).order_by(col("t")) + assert reversed_order.fragment() == call.fragment() + + def test_renders_before_ignore_nulls(self) -> None: + call = fn("first_value", col("x")).order_by(col("t")).ignore_nulls() + assert call.fragment() == 'first_value("x" ORDER BY "t" IGNORE NULLS)' + assert fn("first_value", col("x")).ignore_nulls().order_by(col("t")).fragment() == call.fragment() + + def test_direction_and_null_order(self) -> None: + call = col("x").last().order_by(col("t").desc().nulls_last()) + assert call.fragment() == '"last"("x" ORDER BY "t" DESC NULLS LAST)' + + def test_several_keys(self) -> None: + call = col("x").string_agg(lit(",")).order_by(col("t").desc(), col("x").nulls_first()) + assert call.fragment() == """string_agg("x", ',' ORDER BY "t" DESC, "x" NULLS FIRST)""" + + def test_a_later_order_by_replaces_the_keys(self) -> None: + call = col("x").first().order_by(col("t")).order_by(col("x").desc()) + assert call.fragment() == '"first"("x" ORDER BY "x" DESC)' + assert col("x").first(order_by=col("t")).order_by(col("x")).fragment() == '"first"("x" ORDER BY "x")' + + def test_keyword_on_a_method(self) -> None: + assert col("x").first(order_by=col("t")).fragment() == col("x").first().order_by(col("t")).fragment() + two = col("x").first(order_by=[col("t"), col("x")]) + assert two.fragment() == col("x").first().order_by(col("t"), col("x")).fragment() + assert col("x").first(order_by=None).fragment() == '"first"("x")' + + def test_keyword_beside_positional_arguments(self) -> None: + joined = col("x").string_agg(lit("|"), order_by=col("t").desc()) + assert joined.fragment() == """string_agg("x", '|' ORDER BY "t" DESC)""" + quantile = col("t").quantile(0.5, order_by=col("t")) + assert quantile.fragment() == 'quantile_disc("t", 0.5::DOUBLE ORDER BY "t")' + + def test_fn_reaches_it_through_the_method(self) -> None: + assert fn("first", col("x")).order_by(col("t")).fragment() == '"first"("x" ORDER BY "t")' + # A keyword argument of fn() is a named SQL argument, whatever its name. + assert fn("first", col("x"), order_by=col("t")).fragment() == '"first"("x", "order_by" := "t")' + + def test_through_a_family_wrapper(self) -> None: + wrapped = col("x").first().str().order_by(col("t")) + assert wrapped.fragment() == '"first"("x" ORDER BY "t")' + assert wrapped.lower().fragment() == 'lower("first"("x" ORDER BY "t"))' + + def test_a_family_method_takes_no_keyword(self) -> None: + # `.list().sum()` is the scalar `list_sum`, so it has no rows to order. + with pytest.raises(TypeError, match="unexpected keyword argument 'order_by'"): + col("xs").list().sum(order_by=col("t")) # type: ignore[call-arg] + + def test_the_original_is_unchanged(self) -> None: + plain = col("x").first() + ordered = plain.order_by(col("t")) + assert plain.fragment() == '"first"("x")' + assert ordered.fragment() == '"first"("x" ORDER BY "t")' + + +class TestRefusals: + def test_a_string_key_is_refused(self) -> None: + with pytest.raises(TypeError, match="order_by takes an expression, not the string 't'"): + col("x").first().order_by("t") + with pytest.raises(TypeError, match="order_by takes an expression, not the string 't'"): + col("x").first(order_by="t") + with pytest.raises(TypeError, match="order_by takes an expression, not the string 't'"): + col("x").first(order_by=[col("g"), "t"]) + + def test_refused_on_a_non_aggregate(self) -> None: + with pytest.raises(TypeError, match="order_by\\(\\) applies to an aggregate call"): + col("x").order_by(col("t")) + with pytest.raises(TypeError, match="order_by\\(\\) applies to an aggregate call"): + (col("x") + 1).order_by(col("t")) + + def test_refused_after_over(self) -> None: + with pytest.raises(TypeError, match="order_by\\(\\) applies to an aggregate call"): + col("x").first().over(partition_by=col("g")).order_by(col("t")) + + def test_no_keys_is_refused(self) -> None: + with pytest.raises(TypeError, match="at least one key"): + col("x").first().order_by() + with pytest.raises(TypeError, match="at least one key"): + col("x").first(order_by=[]) + + def test_what_only_the_binder_knows_fails_at_run_time(self, con: duckdb.DuckDBPyConnection) -> None: + """A scalar call, a window-only function, a DISTINCT key outside the arguments and a star all build.""" + scalar = fn("upper", col("x")).order_by(col("t")) + assert scalar.fragment() == 'upper("x" ORDER BY "t")' + with pytest.raises(duckdb.InvalidInputException, match="Scalar Function"): + con.execute(f"SELECT {scalar.fragment()} FROM t") + with pytest.raises(duckdb.BinderException): + con.execute(f"SELECT {row_number().order_by(col('t')).fragment()} FROM t") + with pytest.raises(duckdb.BinderException, match="DISTINCT aggregate"): + con.execute(f"SELECT {col('x').n_unique().order_by(col('t')).fragment()} FROM t") + with pytest.raises(duckdb.BinderException, match="STAR"): + con.execute(f"SELECT {count_all().order_by(col('t')).fragment()} FROM t") + + +class TestAnswers: + """Each ordered call answers what the same SQL answers, on groups with NULLs in the key and in the value.""" + + def test_first_and_last(self, con: duckdb.DuckDBPyConnection) -> None: + assert grouped(con, col("x").first().order_by(col("t"))) == [(1, "b"), (2, "c")] + assert grouped(con, col("x").last().order_by(col("t"))) == [(1, "d"), (2, "c")] + assert grouped(con, col("x").first().order_by(col("t").nulls_first())) == [(1, "d"), (2, "c")] + # The engine puts NULLs last in either direction unless told otherwise. + assert grouped(con, col("x").first().order_by(col("t").desc())) == [(1, "a"), (2, "c")] + assert grouped(con, col("x").first().order_by(col("t").desc().nulls_first())) == [(1, "d"), (2, "c")] + + def test_a_null_value_is_a_value(self, con: duckdb.DuckDBPyConnection) -> None: + picked = col("x").first().order_by(col("t")).where(col("t") > 1) + assert grouped(con, picked) == answers(con, "first(x ORDER BY t) FILTER (WHERE t > 1)") + assert grouped(con, picked) == [(1, None), (2, "c")] + + def test_a_second_key_breaks_a_tie(self, con: duckdb.DuckDBPyConnection) -> None: + # In group 1 three rows tie on `x IS NULL`; the second key decides which of them is first. + lowest = col("t").first().order_by(col("x").is_null(), col("t")) + highest = col("t").first().order_by(col("x").is_null(), col("t").desc()) + assert grouped(con, lowest) == answers(con, "first(t ORDER BY x IS NULL, t)") == [(1, 1), (2, 5)] + assert grouped(con, highest) == answers(con, "first(t ORDER BY x IS NULL, t DESC)") == [(1, 3), (2, 5)] + + def test_string_agg_and_array_agg(self, con: duckdb.DuckDBPyConnection) -> None: + joined = col("x").string_agg(lit(",")).order_by(col("t")) + assert grouped(con, joined) == [(1, "b,a,d"), (2, "c,c")] + listed = col("x").array_agg().order_by(col("t").desc()) + assert grouped(con, listed) == [(1, ["a", None, "b", "d"]), (2, ["c", "c"])] + + def test_with_distinct(self, con: duckdb.DuckDBPyConnection) -> None: + ordered = col("x").n_unique().order_by(col("x")) + assert grouped(con, ordered) == answers(con, "count(DISTINCT x ORDER BY x)") + assert grouped(con, ordered) == [(1, 3), (2, 1)] + + def test_as_a_window(self, con: duckdb.DuckDBPyConnection) -> None: + windowed = col("x").first().order_by(col("t")).over(partition_by=col("g")) + assert windowed.fragment() == '"first"("x" ORDER BY "t") OVER (PARTITION BY "g")' + rows = con.execute(f"SELECT {windowed.fragment()} FROM t ORDER BY g, t").fetchall() + assert [r[0] for r in rows] == ["b", "b", "b", "b", "c", "c"] + + def test_as_a_window_ignoring_nulls(self, con: duckdb.DuckDBPyConnection) -> None: + keys = col("x").nulls_first() + with_nulls = fn("first_value", col("x")).order_by(keys).over(partition_by=col("g")) + without = fn("first_value", col("x")).order_by(keys).ignore_nulls().over(partition_by=col("g")) + source = " FROM t ORDER BY g, t" + assert [r[0] for r in con.execute(f"SELECT {with_nulls.fragment()}{source}").fetchall()][:4] == [None] * 4 + assert [r[0] for r in con.execute(f"SELECT {without.fragment()}{source}").fetchall()] == ["a"] * 4 + ["c"] * 2 + + def test_over_an_empty_group(self, con: duckdb.DuckDBPyConnection) -> None: + assert grouped(con, col("x").first().order_by(col("t")), "WHERE false") == [] + empty = con.execute(f"SELECT {col('x').first().order_by(col('t')).fragment()} FROM t WHERE false").fetchall() + assert empty == [(None,)]