Repository: narwhals-dev/narwhals
Pull request #721 (2024-08-05 14:24:06)

feat: arrow expr `over` method
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #372 

## Checklist

- [x] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes


Patch:
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 1fcfd528..cef2ca35 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -284,6 +284,28 @@ class ArrowExpr:
     def gather_every(self: Self, n: int, offset: int = 0) -> Self:
         return reuse_series_implementation(self, "gather_every", n=n, offset=offset)
 
+    def over(self: Self, keys: list[str]) -> Self:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            if self._output_names is None:
+                msg = (
+                    "Anonymous expressions are not supported in over.\n"
+                    "Instead of `nw.all()`, try using a named expression, such as "
+                    "`nw.col('a', 'b')`\n"
+                )
+                raise ValueError(msg)
+            tmp = df.group_by(*keys).agg(self)
+            tmp = df.select(*keys).join(tmp, how="left", left_on=keys, right_on=keys)
+            return [tmp[name] for name in self._output_names]
+
+        return self.__class__(
+            func,
+            depth=self._depth + 1,
+            function_name=self._function_name + "->over",
+            root_names=self._root_names,
+            output_names=self._output_names,
+            backend_version=self._backend_version,
+        )
+
     @property
     def dt(self: Self) -> ArrowExprDateTimeNamespace:
         return ArrowExprDateTimeNamespace(self)


Test patch:
diff --git a/tests/expr_and_series/over_test.py b/tests/expr_and_series/over_test.py
index 63b212ed..d9f6d607 100644
--- a/tests/expr_and_series/over_test.py
+++ b/tests/expr_and_series/over_test.py
@@ -1,6 +1,5 @@
 from typing import Any
 
-import pandas as pd
 import pytest
 
 import narwhals.stable.v1 as nw
@@ -16,8 +15,6 @@ data = {
 def test_over_single(request: Any, constructor: Any) -> None:
     if "dask" in str(constructor):
         request.applymarker(pytest.mark.xfail)
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
 
     df = nw.from_native(constructor(data))
     result = df.with_columns(c_max=nw.col("c").max().over("a"))
@@ -33,8 +30,6 @@ def test_over_single(request: Any, constructor: Any) -> None:
 def test_over_multiple(request: Any, constructor: Any) -> None:
     if "dask" in str(constructor):
         request.applymarker(pytest.mark.xfail)
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
 
     df = nw.from_native(constructor(data))
     result = df.with_columns(c_min=nw.col("c").min().over("a", "b"))
@@ -47,7 +42,10 @@ def test_over_multiple(request: Any, constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
-def test_over_invalid() -> None:
-    df = nw.from_native(pd.DataFrame(data))
+def test_over_invalid(request: Any, constructor: Any) -> None:
+    if "polars" in str(constructor) or "dask" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor(data))
     with pytest.raises(ValueError, match="Anonymous expressions"):
         df.with_columns(c_min=nw.all().min().over("a", "b"))