fix: nw.len().over was unnecessarily raising for pandas-like (#2372)

Author: MarcoGorelli

narwhals/_dask/expr.py

@@ -644,9 +644,18 @@ def func(df: DaskLazyFrame) -> Sequence[dx.Series]:
                         message=".*`meta` is not specified",
                         category=UserWarning,
                     )
-                    res_native = df.native.groupby(partition_by)[
-                        list(output_names)
-                    ].transform(dask_function_name, **self._call_kwargs)
+                    grouped = df.native.groupby(partition_by)
+                    if dask_function_name == "size":
+                        if len(output_names) != 1:  # pragma: no cover
+                            msg = "Safety check failed, please report a bug."
+                            raise AssertionError(msg)
+                        res_native = grouped.transform(
+                            dask_function_name, **self._call_kwargs
+                        ).to_frame(output_names[0])
+                    else:
+                        res_native = grouped[list(output_names)].transform(
+                            dask_function_name, **self._call_kwargs
+                        )
                 result_frame = df._with_native(
                     res_native.rename(columns=dict(zip(output_names, aliases)))
                 ).native

narwhals/_pandas_like/expr.py

@@ -198,7 +198,7 @@ def cum_sum(self: Self, *, reverse: bool) -> Self:
     def shift(self: Self, n: int) -> Self:
         return self._reuse_series("shift", call_kwargs={"n": n})
 
-    def over(
+    def over(  # noqa: PLR0915
         self: Self,
         partition_by: Sequence[str],
         order_by: Sequence[str] | None,
@@ -265,21 +265,25 @@ def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:
                 elif reverse:
                     columns = list(set(partition_by).union(output_names))
                     df = df[columns][::-1]
+                grouped = df._native_frame.groupby(partition_by)
                 if function_name.startswith("rolling"):
-                    rolling = df._native_frame.groupby(partition_by)[
-                        list(output_names)
-                    ].rolling(**pandas_kwargs)
+                    rolling = grouped[list(output_names)].rolling(**pandas_kwargs)
                     assert pandas_function_name is not None  # help mypy  # noqa: S101
                     if pandas_function_name in {"std", "var"}:
                         res_native = getattr(rolling, pandas_function_name)(
                             ddof=self._call_kwargs["ddof"]
                         )
                     else:
                         res_native = getattr(rolling, pandas_function_name)()
+                elif function_name == "len":
+                    if len(output_names) != 1:  # pragma: no cover
+                        msg = "Safety check failed, please report a bug."
+                        raise AssertionError(msg)
+                    res_native = grouped.transform("size").to_frame(aliases[0])
                 else:
-                    res_native = df._native_frame.groupby(partition_by)[
-                        list(output_names)
-                    ].transform(pandas_function_name, **pandas_kwargs)
+                    res_native = grouped[list(output_names)].transform(
+                        pandas_function_name, **pandas_kwargs
+                    )
                 result_frame = df._with_native(res_native).rename(
                     dict(zip(output_names, aliases))
                 )

tests/expr_and_series/over_test.py

@@ -428,3 +428,17 @@ def test_over_without_partition_by(
     )
     expected = {"a": [1, 2, -1], "b": [1, 3, 4], "i": [0, 1, 2]}
     assert_equal_data(result, expected)
+
+
+def test_len_over_2369(constructor: Constructor, request: pytest.FixtureRequest) -> None:
+    if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3):
+        pytest.skip()
+    if "pandas" in str(constructor) and PANDAS_VERSION < (1, 5):
+        pytest.skip()
+    if any(x in str(constructor) for x in ("modin",)):
+        # https://github.com/modin-project/modin/issues/7508
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor({"a": [1, 2, 4], "b": ["x", "x", "y"]}))
+    result = df.with_columns(a_len_per_group=nw.len().over("b")).sort("a")
+    expected = {"a": [1, 2, 4], "b": ["x", "x", "y"], "a_len_per_group": [2, 2, 1]}
+    assert_equal_data(result, expected)