Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixup for cudf, use fastpath for modin #24

Merged
merged 6 commits into from
Mar 19, 2024
Merged

fixup for cudf, use fastpath for modin #24

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
185e889
fixup for cudf?
MarcoGorelli Mar 19, 2024
c81acd3
wip
MarcoGorelli Mar 19, 2024
1c55212
wip
MarcoGorelli Mar 19, 2024
1d614da
wip
MarcoGorelli Mar 19, 2024
8b2734d
modin can use fastpath too
MarcoGorelli Mar 19, 2024
4ff71b2
modin can use fastpath too
MarcoGorelli Mar 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions narwhals/pandas_like/group_by.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,9 @@ def agg(
raise ValueError(msg)
output_names.extend(expr._output_names)

if implementation == "pandas" and not os.environ.get("NARWHALS_FORCE_GENERIC"):
if implementation in ("pandas", "modin") and not os.environ.get(
"NARWHALS_FORCE_GENERIC"
):
return agg_pandas(
grouped,
exprs,
Expand Down Expand Up @@ -174,7 +176,7 @@ def agg_generic( # noqa: PLR0913
to_remove: list[int] = []
for i, expr in enumerate(exprs):
if is_simple_aggregation(expr):
dfs.append(evaluate_simple_aggregation(expr, grouped))
dfs.append(evaluate_simple_aggregation(expr, grouped, group_by_keys))
to_remove.append(i)
exprs = [expr for i, expr in enumerate(exprs) if i not in to_remove]

Expand Down
11 changes: 9 additions & 2 deletions narwhals/pandas_like/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ def is_simple_aggregation(expr: PandasExpr) -> bool:
)


def evaluate_simple_aggregation(expr: PandasExpr, grouped: Any) -> Any:
def evaluate_simple_aggregation(expr: PandasExpr, grouped: Any, keys: list[str]) -> Any:
"""
Use fastpath for simple aggregations if possible.

Expand All @@ -232,7 +232,14 @@ def evaluate_simple_aggregation(expr: PandasExpr, grouped: Any) -> Any:
Returns naive DataFrame.
"""
if expr._depth == 0:
return grouped.size()["size"].rename(expr._output_names[0]) # type: ignore[index]
# e.g. agg(pl.len())
df = getattr(grouped, expr._function_name.replace("len", "size"))()
df = (
df.drop(columns=keys)
if len(df.shape) > 1
else df.reset_index(drop=True).to_frame("size")
)
return df.rename(columns={"size": expr._output_names[0]}) # type: ignore[index]
if expr._root_names is None or expr._output_names is None:
msg = "Expected expr to have root_names and output_names set, but they are None. Please report a bug."
raise AssertionError(msg)
Expand Down
Loading