Repository: narwhals-dev/narwhals
Pull request #540 (2024-07-16 17:01:40)

feat: allow nw.len in groupby 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/.github/workflows/extremes.yml b/.github/workflows/extremes.yml
index bb3d02fb..17f3c0fa 100644
--- a/.github/workflows/extremes.yml
+++ b/.github/workflows/extremes.yml
@@ -90,6 +90,10 @@ jobs:
         run: pip install polars
       - name: install-reqs
         run: python -m pip install --upgrade tox virtualenv setuptools pip -r requirements-dev.txt
+      - name: uninstall pyarrow
+        run: python -m pip uninstall pyarrow -y
+      - name: install pyarrow nightly
+        run: python -m pip install --extra-index-url https://pypi.fury.io/arrow-nightlies/ --prefer-binary --pre pyarrow
       - name: uninstall pandas
         run: python -m pip uninstall pandas -y
       - name: install-pandas-nightly
diff --git a/narwhals/_arrow/group_by.py b/narwhals/_arrow/group_by.py
index 5e135907..279d32b7 100644
--- a/narwhals/_arrow/group_by.py
+++ b/narwhals/_arrow/group_by.py
@@ -8,6 +8,7 @@ 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.dependencies import get_pyarrow_compute
 from narwhals.utils import remove_prefix
 
 if TYPE_CHECKING:
@@ -51,6 +52,7 @@ class ArrowGroupBy:
         return agg_arrow(
             self._grouped,
             exprs,
+            self._keys,
             output_names,
             self._df._from_native_dataframe,
         )
@@ -59,9 +61,11 @@ class ArrowGroupBy:
 def agg_arrow(
     grouped: Any,
     exprs: list[ArrowExpr],
+    keys: list[str],
     output_names: list[str],
     from_dataframe: Callable[[Any], ArrowDataFrame],
 ) -> ArrowDataFrame:
+    pc = get_pyarrow_compute()
     all_simple_aggs = True
     for expr in exprs:
         if not is_simple_aggregation(expr):
@@ -69,8 +73,24 @@ def agg_arrow(
             break
 
     if all_simple_aggs:
-        simple_aggregations: dict[str, tuple[str, str]] = {}
+        # Mapping from output name to
+        # (input_column_name, function_name, pyarrow_output_name)  # noqa: ERA001
+        simple_aggregations: dict[str, tuple[Any, str, str]] = {}
         for expr in exprs:
+            if expr._depth == 0:
+                # e.g. agg(nw.len()) # noqa: ERA001
+                if (
+                    expr._output_names is None or expr._function_name != "len"
+                ):  # pragma: no cover
+                    msg = "Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues"
+                    raise AssertionError(msg)
+                simple_aggregations[expr._output_names[0]] = (
+                    keys[0],
+                    "count",
+                    f"{keys[0]}_count",
+                )
+                continue
+
             # e.g. agg(nw.mean('a')) # noqa: ERA001
             if (
                 expr._depth != 1 or expr._root_names is None or expr._output_names is None
@@ -81,13 +101,20 @@ def agg_arrow(
             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)
+                simple_aggregations[output_name] = (
+                    root_name,
+                    function_name,
+                    f"{root_name}_{function_name}",
+                )
 
-        aggs = []
+        aggs: list[Any] = []
         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
+            if named_agg[1] == "count":
+                aggs.append((named_agg[0], named_agg[1], pc.CountOptions(mode="all")))
+            else:
+                aggs.append((named_agg[0], named_agg[1]))
+            name_mapping[named_agg[2]] = 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]
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index 6b7ca2fe..40e1d1b7 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -107,6 +107,23 @@ class ArrowNamespace:
             *flatten(column_names), backend_version=self._backend_version
         )
 
+    def len(self) -> ArrowExpr:
+        # coverage bug? this is definitely hit
+        return ArrowExpr(  # pragma: no cover
+            lambda df: [
+                ArrowSeries._from_iterable(
+                    [len(df._native_dataframe)],
+                    name="len",
+                    backend_version=self._backend_version,
+                )
+            ],
+            depth=0,
+            function_name="len",
+            root_names=None,
+            output_names=["len"],
+            backend_version=self._backend_version,
+        )
+
     def all(self) -> ArrowExpr:
         from narwhals._arrow.expr import ArrowExpr
         from narwhals._arrow.series import ArrowSeries
diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py
index 0059ba7f..22a2d862 100644
--- a/narwhals/_pandas_like/group_by.py
+++ b/narwhals/_pandas_like/group_by.py
@@ -32,6 +32,7 @@ class PandasLikeGroupBy:
             list(self._keys),
             sort=False,
             as_index=True,
+            dropna=False,
         )
 
     def agg(
diff --git a/tpch/notebooks/q1/execute.ipynb b/tpch/notebooks/q1/execute.ipynb
index fdd34d36..4760d2df 100755
--- a/tpch/notebooks/q1/execute.ipynb
+++ b/tpch/notebooks/q1/execute.ipynb
@@ -16,7 +16,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals>=0.9.5 ibis-framework"
+    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow ibis-framework narwhals"
    ]
   },
   {
@@ -122,13 +122,13 @@
     "        .group_by([\"l_returnflag\", \"l_linestatus\"])\n",
     "        .agg(\n",
     "            [\n",
-    "                nw.sum(\"l_quantity\").alias(\"sum_qty\"),\n",
-    "                nw.sum(\"l_extendedprice\").alias(\"sum_base_price\"),\n",
-    "                nw.sum(\"disc_price\").alias(\"sum_disc_price\"),\n",
+    "                nw.col(\"l_quantity\").sum().alias(\"sum_qty\"),\n",
+    "                nw.col(\"l_extendedprice\").sum().alias(\"sum_base_price\"),\n",
+    "                nw.col(\"disc_price\").sum().alias(\"sum_disc_price\"),\n",
     "                nw.col(\"charge\").sum().alias(\"sum_charge\"),\n",
-    "                nw.mean(\"l_quantity\").alias(\"avg_qty\"),\n",
-    "                nw.mean(\"l_extendedprice\").alias(\"avg_price\"),\n",
-    "                nw.mean(\"l_discount\").alias(\"avg_disc\"),\n",
+    "                nw.col(\"l_quantity\").mean().alias(\"avg_qty\"),\n",
+    "                nw.col(\"l_extendedprice\").mean().alias(\"avg_price\"),\n",
+    "                nw.col(\"l_discount\").mean().alias(\"avg_disc\"),\n",
     "                nw.len().alias(\"count_order\"),\n",
     "            ],\n",
     "        )\n",
@@ -220,6 +220,7 @@
    "outputs": [],
    "source": [
     "import ibis\n",
+    "import pyarrow.parquet as pq\n",
     "\n",
     "con_pd = ibis.pandas.connect()\n",
     "con_pl = ibis.polars.connect()\n",
@@ -231,6 +232,7 @@
     "    'polars[eager]': lambda x: pl.read_parquet(x),\n",
     "    'polars[lazy]': lambda x: pl.scan_parquet(x),\n",
     "    'polars[lazy][ibis]': lambda x: con_pl.read_parquet(x),\n",
+    "    'pyarrow': lambda x: pq.read_table(x),\n",
     "}"
    ]
   },
@@ -249,7 +251,7 @@
    "id": "8",
    "metadata": {},
    "source": [
-    "## pandas, pyarrow dtypes, ibis"
+    "## PyArrow.table"
    ]
   },
   {
@@ -259,9 +261,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tool = 'pandas[pyarrow][ibis]'\n",
+    "tool = 'pyarrow'\n",
     "fn = IO_FUNCS[tool]\n",
-    "timings = %timeit -o q1_ibis(fn(lineitem), tool='pandas')\n",
+    "timings = %timeit -o q1(fn(lineitem))\n",
     "results[tool] = timings.all_runs"
    ]
   },
@@ -270,7 +272,7 @@
    "id": "10",
    "metadata": {},
    "source": [
-    "## pandas, pyarrow dtypes, native"
+    "## pandas, pyarrow dtypes, ibis"
    ]
   },
   {
@@ -279,6 +281,29 @@
    "id": "11",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "# very slow, comment-out at your own risk\n",
+    "\n",
+    "# tool = 'pandas[pyarrow][ibis]'\n",
+    "# fn = IO_FUNCS[tool]\n",
+    "# timings = %timeit -o q1_ibis(fn(lineitem), tool='pandas')\n",
+    "# results[tool] = timings.all_runs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "12",
+   "metadata": {},
+   "source": [
+    "## pandas, pyarrow dtypes, native"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "tool = 'pandas[pyarrow]'\n",
     "fn = IO_FUNCS[tool]\n",
@@ -288,7 +313,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "12",
+   "id": "14",
    "metadata": {
     "papermill": {
      "duration": 0.005113,
@@ -306,7 +331,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "13",
+   "id": "15",
    "metadata": {
     "papermill": {
      "duration": 196.786925,
@@ -327,7 +352,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "14",
+   "id": "16",
    "metadata": {
     "papermill": {
      "duration": 0.005184,
@@ -345,7 +370,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "15",
+   "id": "17",
    "metadata": {
     "papermill": {
      "duration": 158.748353,
@@ -366,7 +391,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "16",
+   "id": "18",
    "metadata": {
     "papermill": {
      "duration": 0.005773,
@@ -384,7 +409,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "17",
+   "id": "19",
    "metadata": {
     "papermill": {
      "duration": 37.821116,
@@ -405,7 +430,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "18",
+   "id": "20",
    "metadata": {
     "papermill": {
      "duration": 0.005515,
@@ -423,7 +448,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "19",
+   "id": "21",
    "metadata": {
     "papermill": {
      "duration": 4.800698,
@@ -444,7 +469,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "20",
+   "id": "22",
    "metadata": {},
    "source": [
     "## Polars scan_parquet ibis"
@@ -453,7 +478,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "21",
+   "id": "23",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -465,7 +490,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "22",
+   "id": "24",
    "metadata": {},
    "source": [
     "## Save"
@@ -474,7 +499,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "23",
+   "id": "25",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -525,7 +550,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.11.9"
   },
   "papermill": {
    "default_parameters": {},
diff --git a/tpch/notebooks/q2/execute.ipynb b/tpch/notebooks/q2/execute.ipynb
index 45cb23a3..b4e59307 100755
--- a/tpch/notebooks/q2/execute.ipynb
+++ b/tpch/notebooks/q2/execute.ipynb
@@ -16,7 +16,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals>=0.9.5 ibis-framework "
+    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework "
    ]
   },
   {
diff --git a/tpch/notebooks/q3/execute.ipynb b/tpch/notebooks/q3/execute.ipynb
index 81369bf9..f289ea91 100755
--- a/tpch/notebooks/q3/execute.ipynb
+++ b/tpch/notebooks/q3/execute.ipynb
@@ -16,7 +16,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals>=0.9.5 ibis-framework "
+    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework "
    ]
   },
   {
diff --git a/tpch/notebooks/q4/execute.ipynb b/tpch/notebooks/q4/execute.ipynb
index 18b55a02..f5d1b97b 100755
--- a/tpch/notebooks/q4/execute.ipynb
+++ b/tpch/notebooks/q4/execute.ipynb
@@ -16,7 +16,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals>=0.9.5 ibis-framework "
+    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework "
    ]
   },
   {
diff --git a/tpch/notebooks/q5/execute.ipynb b/tpch/notebooks/q5/execute.ipynb
index 80655ac0..a56ae03d 100755
--- a/tpch/notebooks/q5/execute.ipynb
+++ b/tpch/notebooks/q5/execute.ipynb
@@ -16,7 +16,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals>=0.9.5 ibis-framework "
+    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework "
    ]
   },
   {
diff --git a/tpch/notebooks/q6/execute.ipynb b/tpch/notebooks/q6/execute.ipynb
index 0c38ab6a..0f8d6ce5 100755
--- a/tpch/notebooks/q6/execute.ipynb
+++ b/tpch/notebooks/q6/execute.ipynb
@@ -16,7 +16,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals>=0.9.5 ibis-framework "
+    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework "
    ]
   },
   {
diff --git a/tpch/notebooks/q7/execute.ipynb b/tpch/notebooks/q7/execute.ipynb
index 4ec7beac..3b64df2f 100755
--- a/tpch/notebooks/q7/execute.ipynb
+++ b/tpch/notebooks/q7/execute.ipynb
@@ -16,7 +16,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals>=0.9.5 ibis-framework"
+    "!pip uninstall apache-beam -y && pip install -U pandas polars pyarrow narwhals ibis-framework"
    ]
   },
   {


Test patch:
diff --git a/tests/test_group_by.py b/tests/test_group_by.py
index a8f3659c..ff2a8422 100644
--- a/tests/test_group_by.py
+++ b/tests/test_group_by.py
@@ -152,3 +152,16 @@ def test_group_by_multiple_keys(constructor: Any) -> None:
         "c_max": [7, 1],
     }
     compare_dicts(result, expected)
+
+
+def test_key_with_nulls(constructor: Any) -> None:
+    data = {"b": [4, 5, None], "a": [1, 2, 3]}
+    result = (
+        nw.from_native(constructor(data))
+        .group_by("b")
+        .agg(nw.len(), nw.col("a").min())
+        .sort("a")
+        .with_columns(nw.col("b").cast(nw.Float64))
+    )
+    expected = {"b": [4.0, 5, float("nan")], "len": [1, 1, 1], "a": [1, 2, 3]}
+    compare_dicts(result, expected)
diff --git a/tests/tpch_q1_test.py b/tests/tpch_q1_test.py
index 37216971..da540e46 100644
--- a/tests/tpch_q1_test.py
+++ b/tests/tpch_q1_test.py
@@ -29,7 +29,6 @@ def test_q1(library: str, request: Any) -> None:
     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()
@@ -46,13 +45,13 @@ def test_q1(library: str, request: Any) -> None:
         .group_by(["l_returnflag", "l_linestatus"])
         .agg(
             [
-                nw.sum("l_quantity").alias("sum_qty"),
-                nw.sum("l_extendedprice").alias("sum_base_price"),
-                nw.sum("disc_price").alias("sum_disc_price"),
+                nw.col("l_quantity").sum().alias("sum_qty"),
+                nw.col("l_extendedprice").sum().alias("sum_base_price"),
+                nw.col("disc_price").sum().alias("sum_disc_price"),
                 nw.col("charge").sum().alias("sum_charge"),
-                nw.mean("l_quantity").alias("avg_qty"),
-                nw.mean("l_extendedprice").alias("avg_price"),
-                nw.mean("l_discount").alias("avg_disc"),
+                nw.col("l_quantity").mean().alias("avg_qty"),
+                nw.col("l_extendedprice").mean().alias("avg_price"),
+                nw.col("l_discount").mean().alias("avg_disc"),
                 nw.len().alias("count_order"),
             ],
         )