`DuckDBPyRelation.aggregate` has incorrect typing

[m15a]

What happens?

According to the documentation, DuckDBPyRelation.aggregate() accepts str or list[Expression] as the aggr_expr parameter. However, the stub has the signature (line 453 in _duckdb-stubs/__init__.pyi, the main branch)

def aggregate(self, aggr_expr: Expression | str, group_expr: Expression | str = "") -> DuckDBPyRelation

To Reproduce

For example, type checking the following code fails (duckdb 1.4.2). In example.py:

import duckdb as dd
import pandas as pd
from duckdb import ColumnExpression as col
from duckdb import FunctionExpression as fn

cats = dd.from_df(pd.DataFrame({"family": [1, 1], "name": ["nyan", "meow"]}))
dogs = dd.from_df(pd.DataFrame({"family": [1, 2], "name": ["woof", "wan"]}))

print(
    cats.union(dogs)
    .aggregate(
        [
            col("family"),
            fn("list", col("name")).alias("names"),
        ],
        "family",
    )
)

❯ uv run ty check example.py
error[invalid-argument-type]: Argument to bound method `aggregate` is incorrect
  --> example.py:12:9
   |
10 |       cats.union(dogs)
11 |       .aggregate(
12 | /         [
13 | |             col("family"),
14 | |             fn("list", col("name")).alias("names"),
15 | |         ],
   | |_________^ Expected `Expression | str`, found `list[Unknown | Expression]`
16 |           "family",
17 |       )
   |
info: Method defined here
   --> .venv/lib/python3.13/site-packages/_duckdb-stubs/__init__.pyi:464:9
    |
462 |     def __getitem__(self, name: str) -> DuckDBPyRelation: ...
463 |     def __len__(self) -> int: ...
464 |     def aggregate(self, aggr_expr: Expression | str, group_expr: Expression | str = "") -> DuckDBPyRelation: ...
    |         ^^^^^^^^^       --------------------------- Parameter declared here
465 |     def any_value(
466 |         self, column: str, groups: str = "", window_spec: str = "", projected_columns: str = ""
    |
info: rule `invalid-argument-type` is enabled by default

Found 1 diagnostic

OS:

aarch64_darwin

DuckDB Package Version:

1.4.2

Python Version:

3.13.9

Full Name:

Mitsuhiro Nacamura

Affiliation:

Nacamura Data Analysis

What is the latest build you tested with? If possible, we recommend testing with the latest nightly build.

I have tested with a stable release

Did you include all relevant data sets for reproducing the issue?

Yes

Did you include all code required to reproduce the issue?

• Yes, I have

Did you include all relevant configuration to reproduce the issue?

• Yes, I have

[OutSquareCapital]

I stumbled upon this problem and took a look at the merged pull request, however unfortunately the type is still too narrow.
It should be Iterable[Expression | str] instead of list[Expression].
I'm not sure why list[int] don't crash or even what it does, but it does work too apparently.

see below:

Code

from functools import partial

import duckdb
import polars as pl
import pyochain as pc


def _get_rel() -> duckdb.DuckDBPyRelation:
    df = pl.DataFrame({"foo": ["a"], "bar": [1]})
    return duckdb.from_arrow(df)


def _try_agg(rel: duckdb.DuckDBPyRelation, arg: object) -> None:
    def _name() -> str:
        repr = f"{type(arg).__name__}"
        if isinstance(arg, list):
            repr += f"[{type(arg[0]).__name__}]"
        return f"*{repr}*"

    try:
        res = rel.aggregate(arg).fetchone()
        return print(f"- **Ok** with arg {_name()}: `{res}`")
    except Exception as e:
        return print(f"- **Error** with arg {_name()}: `{e}`")


def main() -> None:
    data = ["foo", "bar"]
    # Generator of ColumExpressions
    col_args = (duckdb.ColumnExpression(str(x)) for x in data)
    # custom object, Iterator protocol implemented.
    iter_arg = pc.Iter(data)

    col = duckdb.ColumnExpression("foo")
    tst_on = partial(_try_agg, _get_rel())
    tst_on(data)
    tst_on(col_args)
    tst_on(iter_arg)
    tst_on(6)
    tst_on([6])
    tst_on("foo")
    tst_on([col])
    tst_on(col)


if __name__ == "__main__":
    main()

Output

• Ok with arg list[str]: ('a', 1)
• Ok with arg generator: ('a', 1)
• Ok with arg Iter: ('a', 1)
• Error with arg int: Invalid Input Error: Please provide either a string or list of Expression objects, not <class 'int'>
• Ok with arg list[int]: (6,)
• Ok with arg str: ('a',)
• Ok with arg list[Expression]: ('a',)
• Error with arg Expression: Invalid Input Error: Please provide either a string or list of Expression objects, not <class '_duckdb.Expression'>

[evertlammerts]

Fixed in #323. Widened aggr_expr on DuckDBPyRelation.aggregate() from Expression | str | list[Expression] to str | Expression | list[Expression | str].

[OutSquareCapital]

Sorry to be pedantic but again list is wrong here as I showed above.
It's too narrow, and I'm very rarely using plain lists in my day-to-day usage.
So either I use cast and type ignore everywhere, or I inefficiently call list on my Iterator/Sequence/tuple before passing it to the method.