32 lines
1 KiB
Python
32 lines
1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from narwhals._sql.group_by import SQLGroupBy
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
import ibis.expr.types as ir # noqa: F401
|
|
|
|
from narwhals._ibis.dataframe import IbisLazyFrame
|
|
from narwhals._ibis.expr import IbisExpr
|
|
|
|
|
|
class IbisGroupBy(SQLGroupBy["IbisLazyFrame", "IbisExpr", "ir.Value"]):
|
|
def __init__(
|
|
self,
|
|
df: IbisLazyFrame,
|
|
keys: Sequence[str] | Sequence[IbisExpr],
|
|
/,
|
|
*,
|
|
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: IbisExpr) -> IbisLazyFrame:
|
|
native = self.compliant.native
|
|
return self.compliant._with_native(
|
|
native.group_by(self._keys).aggregate(*self._evaluate_exprs(exprs))
|
|
).rename(dict(zip(self._keys, self._output_key_names)))
|