1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
|
from __future__ import annotations
from itertools import chain
from typing import TYPE_CHECKING, Sequence
from narwhals._compliant import LazyGroupBy
if TYPE_CHECKING:
from duckdb import Expression # noqa: F401
from narwhals._duckdb.dataframe import DuckDBLazyFrame
from narwhals._duckdb.expr import DuckDBExpr
class DuckDBGroupBy(LazyGroupBy["DuckDBLazyFrame", "DuckDBExpr", "Expression"]):
def __init__(
self,
df: DuckDBLazyFrame,
keys: Sequence[DuckDBExpr] | Sequence[str],
/,
*,
drop_null_keys: bool,
) -> None:
frame, self._keys, self._output_key_names = self._parse_keys(df, keys=keys)
self._compliant_frame = frame.drop_nulls(self._keys) if drop_null_keys else frame
def agg(self, *exprs: DuckDBExpr) -> DuckDBLazyFrame:
agg_columns = list(chain(self._keys, self._evaluate_exprs(exprs)))
return self.compliant._with_native(
self.compliant.native.aggregate(agg_columns) # type: ignore[arg-type]
).rename(dict(zip(self._keys, self._output_key_names)))
|
