Repository: narwhals-dev/narwhals
Pull request #513 (2024-07-13 14:42:10)

feat: add dataframe.group_by for Arrow
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 920a662a..dffa5500 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -17,6 +17,7 @@ from narwhals.utils import flatten
 if TYPE_CHECKING:
     from typing_extensions import Self
 
+    from narwhals._arrow.group_by import ArrowGroupBy
     from narwhals._arrow.namespace import ArrowNamespace
     from narwhals._arrow.series import ArrowSeries
     from narwhals._arrow.typing import IntoArrowExpr
@@ -191,6 +192,14 @@ class ArrowDataFrame:
         df = self._native_dataframe.__class__.from_arrays(to_concat, names=output_names)
         return self._from_native_dataframe(df)
 
+    def group_by(self, *keys: str | Iterable[str]) -> ArrowGroupBy:
+        from narwhals._arrow.group_by import ArrowGroupBy
+
+        return ArrowGroupBy(
+            self,
+            flatten(keys),
+        )
+
     def join(
         self,
         other: Self,
diff --git a/narwhals/_arrow/group_by.py b/narwhals/_arrow/group_by.py
new file mode 100644
index 00000000..b3764cba
--- /dev/null
+++ b/narwhals/_arrow/group_by.py
@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from copy import copy
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Callable
+
+from narwhals._expression_parsing import is_simple_aggregation
+from narwhals._expression_parsing import parse_into_exprs
+from narwhals.dependencies import get_pyarrow
+from narwhals.utils import remove_prefix
+
+if TYPE_CHECKING:
+    from narwhals._arrow.dataframe import ArrowDataFrame
+    from narwhals._arrow.expr import ArrowExpr
+    from narwhals._arrow.typing import IntoArrowExpr
+
+POLARS_TO_ARROW_AGGREGATIONS = {
+    "len": "size",
+}
+
+
+class ArrowGroupBy:
+    def __init__(self, df: ArrowDataFrame, keys: list[str]) -> None:
+        pa = get_pyarrow()
+        self._df = df
+        self._keys = list(keys)
+        self._grouped = pa.TableGroupBy(self._df._native_dataframe, list(self._keys))
+
+    def agg(
+        self,
+        *aggs: IntoArrowExpr,
+        **named_aggs: IntoArrowExpr,
+    ) -> ArrowDataFrame:
+        exprs = parse_into_exprs(
+            *aggs,
+            namespace=self._df.__narwhals_namespace__(),
+            **named_aggs,
+        )
+        output_names: list[str] = copy(self._keys)
+        for expr in exprs:
+            if expr._output_names is None:
+                msg = (
+                    "Anonymous expressions are not supported in group_by.agg.\n"
+                    "Instead of `nw.all()`, try using a named expression, such as "
+                    "`nw.col('a', 'b')`\n"
+                )
+                raise ValueError(msg)
+            output_names.extend(expr._output_names)
+
+        return agg_arrow(
+            self._grouped,
+            exprs,
+            output_names,
+            self._df._from_native_dataframe,
+        )
+
+
+def agg_arrow(
+    grouped: Any,
+    exprs: list[ArrowExpr],
+    output_names: list[str],
+    from_dataframe: Callable[[Any], ArrowDataFrame],
+) -> ArrowDataFrame:
+    all_simple_aggs = True
+    for expr in exprs:
+        if not is_simple_aggregation(expr):
+            all_simple_aggs = False
+            break
+
+    if all_simple_aggs:
+        simple_aggregations: dict[str, tuple[str, str]] = {}
+        for expr in exprs:
+            # e.g. agg(nw.mean('a')) # noqa: ERA001
+            assert expr._depth == 1
+            assert expr._root_names is not None
+            assert expr._output_names is not None
+            function_name = remove_prefix(expr._function_name, "col->")
+            function_name = POLARS_TO_ARROW_AGGREGATIONS.get(function_name, function_name)
+            for root_name, output_name in zip(expr._root_names, expr._output_names):
+                simple_aggregations[output_name] = (root_name, function_name)
+
+        aggs = []
+        name_mapping = {}
+        for output_name, named_agg in simple_aggregations.items():
+            aggs.append((named_agg[0], named_agg[1]))
+            name_mapping[f"{named_agg[0]}_{named_agg[1]}"] = output_name
+        result_simple = grouped.aggregate(aggs)
+        result_simple = result_simple.rename_columns(
+            [name_mapping.get(col, col) for col in result_simple.column_names]
+        ).select(output_names)
+        return from_dataframe(result_simple)
+
+    msg = (
+        "Non-trivial complex found.\n\n"
+        "Hint: you were probably trying to apply a non-elementary aggregation with a "
+        "pyarrow table.\n"
+        "Please rewrite your query such that group-by aggregations "
+        "are elementary. For example, instead of:\n\n"
+        "    df.group_by('a').agg(nw.col('b').round(2).mean())\n\n"
+        "use:\n\n"
+        "    df.with_columns(nw.col('b').round(2)).group_by('a').agg(nw.col('b').mean())\n\n"
+    )
+    raise ValueError(msg)
diff --git a/narwhals/_expression_parsing.py b/narwhals/_expression_parsing.py
index 8a05c8b5..50aefb55 100644
--- a/narwhals/_expression_parsing.py
+++ b/narwhals/_expression_parsing.py
@@ -241,3 +241,20 @@ def reuse_series_namespace_implementation(
         root_names=expr._root_names,
         output_names=expr._output_names,
     )
+
+
+def is_simple_aggregation(expr: CompliantExpr) -> bool:
+    """
+    Check if expr is a very simple one, such as:
+
+    - nw.col('a').mean()  # depth 1
+    - nw.mean('a')  # depth 1
+    - nw.len()  # depth 0
+
+    as opposed to, say
+
+    - nw.col('a').filter(nw.col('b')>nw.col('c')).max()
+
+    because then, we can use a fastpath in pandas.
+    """
+    return expr._depth < 2
diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py
index f4e3ef95..6895cdc5 100644
--- a/narwhals/_pandas_like/group_by.py
+++ b/narwhals/_pandas_like/group_by.py
@@ -8,9 +8,9 @@ from typing import Any
 from typing import Callable
 from typing import Iterator
 
+from narwhals._expression_parsing import is_simple_aggregation
 from narwhals._expression_parsing import parse_into_exprs
 from narwhals._pandas_like.utils import Implementation
-from narwhals._pandas_like.utils import is_simple_aggregation
 from narwhals._pandas_like.utils import native_series_from_iterable
 from narwhals.utils import remove_prefix
 
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index b9f6050f..3548966c 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -114,23 +114,6 @@ def create_native_series(
     )
 
 
-def is_simple_aggregation(expr: PandasLikeExpr) -> bool:
-    """
-    Check if expr is a very simple one, such as:
-
-    - nw.col('a').mean()  # depth 1
-    - nw.mean('a')  # depth 1
-    - nw.len()  # depth 0
-
-    as opposed to, say
-
-    - nw.col('a').filter(nw.col('b')>nw.col('c')).max()
-
-    because then, we can use a fastpath in pandas.
-    """
-    return expr._depth < 2
-
-
 def horizontal_concat(
     dfs: list[Any], *, implementation: Implementation, backend_version: tuple[int, ...]
 ) -> Any:
diff --git a/utils/check_backend_completeness.py b/utils/check_backend_completeness.py
index 47ac709b..cb9a6f39 100644
--- a/utils/check_backend_completeness.py
+++ b/utils/check_backend_completeness.py
@@ -12,7 +12,6 @@ import narwhals as nw
 from narwhals._arrow.dataframe import ArrowDataFrame
 
 MISSING = [
-    "DataFrame.group_by",
     "DataFrame.is_duplicated",
     "DataFrame.is_empty",
     "DataFrame.is_unique",


Test patch:
diff --git a/tests/test_group_by.py b/tests/test_group_by.py
index ee08d2cf..16798ed1 100644
--- a/tests/test_group_by.py
+++ b/tests/test_group_by.py
@@ -4,6 +4,7 @@ from typing import Any
 
 import pandas as pd
 import polars as pl
+import pyarrow as pa
 import pytest
 
 import narwhals.stable.v1 as nw
@@ -40,6 +41,14 @@ def test_invalid_group_by() -> None:
         ValueError, match=r"Anonymous expressions are not supported in group_by\.agg"
     ):
         df.group_by("a").agg(nw.all().mean())
+    with pytest.raises(
+        ValueError, match=r"Anonymous expressions are not supported in group_by\.agg"
+    ):
+        nw.from_native(pa.table({"a": [1, 2, 3]})).group_by("a").agg(nw.all().mean())
+    with pytest.raises(ValueError, match=r"Non-trivial complex found"):
+        nw.from_native(pa.table({"a": [1, 2, 3]})).group_by("a").agg(
+            nw.col("b").mean().min()
+        )
 
 
 def test_group_by_iter(constructor: Any) -> None:
@@ -79,3 +88,61 @@ def test_group_by_empty_result_pandas() -> None:
         df.filter(nw.col("a") < 0).group_by("a").agg(
             nw.col("b").sum().round(2).alias("c")
         )
+
+
+def test_group_by_simple_named(constructor_with_pyarrow: Any) -> None:
+    data = {"a": [1, 1, 2], "b": [4, 5, 6], "c": [7, 2, 1]}
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
+    result = (
+        df.group_by("a")
+        .agg(
+            b_min=nw.col("b").min(),
+            b_max=nw.col("b").max(),
+        )
+        .sort("a")
+    )
+    expected = {
+        "a": [1, 2],
+        "b_min": [4, 6],
+        "b_max": [5, 6],
+    }
+    compare_dicts(result, expected)
+
+
+def test_group_by_simple_unnamed(constructor_with_pyarrow: Any) -> None:
+    data = {"a": [1, 1, 2], "b": [4, 5, 6], "c": [7, 2, 1]}
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
+    result = (
+        df.group_by("a")
+        .agg(
+            nw.col("b").min(),
+            nw.col("c").max(),
+        )
+        .sort("a")
+    )
+    expected = {
+        "a": [1, 2],
+        "b": [4, 6],
+        "c": [7, 1],
+    }
+    compare_dicts(result, expected)
+
+
+def test_group_by_multiple_keys(constructor_with_pyarrow: Any) -> None:
+    data = {"a": [1, 1, 2], "b": [4, 4, 6], "c": [7, 2, 1]}
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
+    result = (
+        df.group_by("a", "b")
+        .agg(
+            c_min=nw.col("c").min(),
+            c_max=nw.col("c").max(),
+        )
+        .sort("a")
+    )
+    expected = {
+        "a": [1, 2],
+        "b": [4, 6],
+        "c_min": [2, 1],
+        "c_max": [7, 1],
+    }
+    compare_dicts(result, expected)
diff --git a/tests/tpch_q1_test.py b/tests/tpch_q1_test.py
index d73400ee..37216971 100644
--- a/tests/tpch_q1_test.py
+++ b/tests/tpch_q1_test.py
@@ -2,10 +2,12 @@ from __future__ import annotations
 
 import os
 from datetime import datetime
+from typing import Any
 from unittest import mock
 
 import pandas as pd
 import polars as pl
+import pyarrow.parquet as pq
 import pytest
 
 import narwhals.stable.v1 as nw
@@ -15,18 +17,20 @@ from tests.utils import compare_dicts
 
 @pytest.mark.parametrize(
     "library",
-    ["pandas", "polars"],
+    ["pandas", "polars", "pyarrow"],
 )
 @pytest.mark.filterwarnings("ignore:.*Passing a BlockManager.*:DeprecationWarning")
-@pytest.mark.skipif(
-    parse_version(pd.__version__) < parse_version("1.0.0"), reason="too old for pyarrow"
-)
-def test_q1(library: str) -> None:
-    if library == "pandas":
+def test_q1(library: str, request: Any) -> None:
+    if library == "pandas" and parse_version(pd.__version__) < (1, 5):
+        request.applymarker(pytest.mark.xfail)
+    elif library == "pandas":
         df_raw = pd.read_parquet("tests/data/lineitem.parquet")
         df_raw["l_shipdate"] = pd.to_datetime(df_raw["l_shipdate"])
-    else:
+    elif library == "polars":
         df_raw = pl.scan_parquet("tests/data/lineitem.parquet")
+    else:
+        request.applymarker(pytest.mark.xfail)
+        df_raw = pq.read_table("tests/data/lineitem.parquet")
     var_1 = datetime(1998, 9, 2)
     df = nw.from_native(df_raw).lazy()
     query_result = (
@@ -88,11 +92,10 @@ def test_q1(library: str) -> None:
     "ignore:.*Passing a BlockManager.*:DeprecationWarning",
     "ignore:.*Complex.*:UserWarning",
 )
-@pytest.mark.skipif(
-    parse_version(pd.__version__) < parse_version("1.0.0"), reason="too old for pyarrow"
-)
-def test_q1_w_generic_funcs(library: str) -> None:
-    if library == "pandas":
+def test_q1_w_generic_funcs(library: str, request: Any) -> None:
+    if library == "pandas" and parse_version(pd.__version__) < (1, 5):
+        request.applymarker(pytest.mark.xfail)
+    elif library == "pandas":
         df_raw = pd.read_parquet("tests/data/lineitem.parquet")
         df_raw["l_shipdate"] = pd.to_datetime(df_raw["l_shipdate"])
     else: