Skip to content

Commit

Permalink
loudly raise for over with key overlap
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Jan 26, 2025
1 parent 9f7aa3b commit 37f6315
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 4 deletions.
10 changes: 9 additions & 1 deletion narwhals/_dask/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -661,7 +661,15 @@ def null_count(self: Self) -> Self:

def over(self: Self, keys: list[str]) -> Self:
def func(df: DaskLazyFrame) -> list[Any]:
_, aliases = evaluate_output_names_and_aliases(self, df, [])
output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
if overlap := set(output_names).intersection(keys):
# E.g. `df.select(nw.all().sum().over('a'))`. This is well-defined,
# we just don't support it yet.
msg = (
f"Column names {overlap} appear in both expression output names and in `over` keys.\n"
"This is not yet supported."
)
raise NotImplementedError(msg)
if df._native_frame.npartitions == 1: # pragma: no cover
tmp = df.group_by(*keys, drop_null_keys=False).agg(self)
tmp_native = (
Expand Down
14 changes: 11 additions & 3 deletions tests/expr_and_series/over_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,11 +188,19 @@ def test_over_anonymous_cumulative(constructor_eager: ConstructorEager) -> None:
assert_equal_data(result, expected)


def test_over_anonymous_reduction(constructor_eager: ConstructorEager) -> None:
df = nw.from_native(constructor_eager({"a": [1, 1, 2], "b": [4, 5, 6]}))
def test_over_anonymous_reduction(
constructor: Constructor, request: pytest.FixtureRequest
) -> None:
if "duckdb" in str(constructor) or "pyspark" in str(constructor):
# TODO(unassigned): we should be able to support these
request.applymarker(pytest.mark.xfail)

df = nw.from_native(constructor({"a": [1, 1, 2], "b": [4, 5, 6]}))
context = (
pytest.raises(NotImplementedError)
if df.implementation.is_pyarrow() or df.implementation.is_pandas_like()
if df.implementation.is_pyarrow()
or df.implementation.is_pandas_like()
or df.implementation.is_dask()
else does_not_raise()
)
with context:
Expand Down

0 comments on commit 37f6315

Please sign in to comment.