Skip to content

Commit

Permalink
enh: pandas-like group-by complex-agg
Browse files Browse the repository at this point in the history
  • Loading branch information
FBruzzesi committed Jul 12, 2024
1 parent 057c00c commit 6fbd640
Showing 1 changed file with 8 additions and 10 deletions.
18 changes: 8 additions & 10 deletions narwhals/_pandas_like/group_by.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ def agg(
**named_aggs,
)
implementation: Implementation = self._df._implementation
backend = self._df.__native_namespace__()
output_names: list[str] = copy(self._keys)
for expr in exprs:
if expr._output_names is None:
Expand All @@ -64,7 +65,7 @@ def agg(
self._from_native_dataframe,
dataframe_is_empty=self._df._native_dataframe.empty,
implementation=implementation,
backend_version=self._df._backend_version,
backend=backend,
)

def _from_native_dataframe(self, df: PandasLikeDataFrame) -> PandasLikeDataFrame:
Expand Down Expand Up @@ -99,8 +100,8 @@ def agg_pandas(
from_dataframe: Callable[[Any], PandasLikeDataFrame],
*,
implementation: Any,
backend_version: tuple[int, ...],
dataframe_is_empty: bool,
backend: Any,
) -> PandasLikeDataFrame:
"""
This should be the fastpath, but cuDF is too far behind to use it.
Expand Down Expand Up @@ -189,11 +190,8 @@ def func(df: Any) -> Any:
implementation=implementation,
)

if implementation is Implementation.PANDAS and backend_version >= (2, 2):
result_complex = grouped.apply(func, include_groups=False)
else: # pragma: no cover
result_complex = grouped.apply(func)

result = result_complex.reset_index()

return from_dataframe(result.loc[:, output_names])
output_values = [
dict(zip(output_names, (*grp_name, *func(grp_values))))
for grp_name, grp_values in grouped
]
return from_dataframe(backend.DataFrame(output_values))

0 comments on commit 6fbd640

Please sign in to comment.