Repository: narwhals-dev/narwhals
Pull request #1677 (2024-12-29 19:49:25)

feat: raise informative error message when grouping by expressions
This has come up a couple of times on discord

it may be possible to support this, but i don't have time right now, so i'd say that for now let's raise a proper error message and we can revisit later if necessary

<!--
# 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 #\<issue number\>
- Closes #\<issue number\>

## 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 8576dd2b..122ece9c 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -597,7 +597,6 @@ class ArrowDataFrame(CompliantDataFrame, CompliantLazyFrame):
         return pa_csv.write_csv(pa_table, file)
 
     def is_duplicated(self: Self) -> ArrowSeries:
-        import numpy as np  # ignore-banned-import
         import pyarrow as pa
         import pyarrow.compute as pc
 
@@ -608,14 +607,18 @@ class ArrowDataFrame(CompliantDataFrame, CompliantLazyFrame):
         columns = self.columns
         col_token = generate_temporary_column_name(n_bytes=8, columns=columns)
         row_count = (
-            df.append_column(col_token, pa.array(np.arange(len(self))))
+            df.append_column(col_token, pa.repeat(pa.scalar(1), len(self)))
             .group_by(columns)
-            .aggregate([(col_token, "count")])
+            .aggregate([(col_token, "sum")])
         )
         is_duplicated = pc.greater(
             df.join(
-                row_count, keys=columns, right_keys=columns, join_type="inner"
-            ).column(f"{col_token}_count"),
+                row_count,
+                keys=columns,
+                right_keys=columns,
+                join_type="inner",
+                use_threads=False,
+            ).column(f"{col_token}_sum"),
             1,
         )
         return ArrowSeries(
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index ea37ae76..99f043eb 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -445,18 +445,18 @@ class ArrowWhen:
         except TypeError:
             # `self._otherwise_value` is a scalar and can't be converted to an expression
             value_series = condition.__class__._from_iterable(
-                [self._then_value] * len(condition),
+                pa.repeat(pa.scalar(self._then_value), len(condition)),
                 name="literal",
                 backend_version=self._backend_version,
                 version=self._version,
             )
 
         value_series_native = value_series._native_series
-        condition_native = condition._native_series.combine_chunks()
+        condition_native = condition._native_series
 
         if self._otherwise_value is None:
-            otherwise_native = pa.array(
-                [None] * len(condition_native), type=value_series_native.type
+            otherwise_native = pa.repeat(
+                pa.scalar(None, type=value_series_native.type), len(condition_native)
             )
             return [
                 value_series._from_native_series(
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index f284ce42..23907465 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -261,7 +261,7 @@ def vertical_concat(dfs: list[pa.Table]) -> pa.Table:
 
     import pyarrow as pa
 
-    return pa.concat_tables(dfs).combine_chunks()
+    return pa.concat_tables(dfs)
 
 
 def diagonal_concat(dfs: list[pa.Table], backend_version: tuple[int, ...]) -> pa.Table:
@@ -276,7 +276,7 @@ def diagonal_concat(dfs: list[pa.Table], backend_version: tuple[int, ...]) -> pa
         if backend_version < (14, 0, 0)
         else {"promote_options": "default"}  # type: ignore[dict-item]
     )
-    return pa.concat_tables(dfs, **kwargs).combine_chunks()
+    return pa.concat_tables(dfs, **kwargs)
 
 
 def floordiv_compat(left: Any, right: Any) -> Any:
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index ccdba365..4bc706e4 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -2209,9 +2209,18 @@ class DataFrame(BaseFrame[DataFrameT]):
             │ c   ┆ 3   ┆ 1   │
             └─────┴─────┴─────┘
         """
+        from narwhals.expr import Expr
         from narwhals.group_by import GroupBy
+        from narwhals.series import Series
 
-        return GroupBy(self, *flatten(keys), drop_null_keys=drop_null_keys)
+        flat_keys = flatten(keys)
+        if any(isinstance(x, (Expr, Series)) for x in flat_keys):
+            msg = (
+                "`group_by` with expression or Series keys is not (yet?) supported.\n\n"
+                "Hint: instead of `df.group_by(nw.col('a'))`, use `df.group_by('a')`."
+            )
+            raise NotImplementedError(msg)
+        return GroupBy(self, *flat_keys, drop_null_keys=drop_null_keys)
 
     def sort(
         self,
@@ -4444,9 +4453,18 @@ class LazyFrame(BaseFrame[FrameT]):
             │ c   ┆ 3   ┆ 1   │
             └─────┴─────┴─────┘
         """
+        from narwhals.expr import Expr
         from narwhals.group_by import LazyGroupBy
+        from narwhals.series import Series
 
-        return LazyGroupBy(self, *flatten(keys), drop_null_keys=drop_null_keys)
+        flat_keys = flatten(keys)
+        if any(isinstance(x, (Expr, Series)) for x in flat_keys):
+            msg = (
+                "`group_by` with expression or Series keys is not (yet?) supported.\n\n"
+                "Hint: instead of `df.group_by(nw.col('a'))`, use `df.group_by('a')`."
+            )
+            raise NotImplementedError(msg)
+        return LazyGroupBy(self, *flat_keys, drop_null_keys=drop_null_keys)
 
     def sort(
         self,


Test patch:
diff --git a/tests/group_by_test.py b/tests/group_by_test.py
index a36c90bf..7ed99f07 100644
--- a/tests/group_by_test.py
+++ b/tests/group_by_test.py
@@ -453,3 +453,9 @@ def test_all_kind_of_aggs(
         "i": [3, 2],
     }
     assert_equal_data(result, expected)
+
+
+def test_group_by_expr(constructor: Constructor) -> None:
+    df = nw.from_native(constructor({"a": [1, 1, 3], "b": [4, 5, 6]}))
+    with pytest.raises(NotImplementedError, match=r"not \(yet\?\) supported"):
+        df.group_by(nw.col("a")).agg(nw.col("b").mean())  # type: ignore[arg-type]