Allow any length-preserving `Expr` in `group_by(*keys)`

[dhirschfeld]

Describe the bug

Note

I think this is more of a missing feature than a bug.

I have code that works fine with polars but gives a ComputeError when run with narwhals

Steps or code to reproduce the bug

from collections.abc import Generator
import narwhals as nw
import numpy as np
import pandas as pd
import polars as pl
from narwhals.typing import IntoDataFrameT



def partition_fn(
    data: IntoDataFrameT,
) -> Generator[tuple[str, pa.Table], None, None]:
    """Group by unique quarters in `value_timestamp` and yield `(quarter, group)`
    tuples where the quarters are formatted as `{year}Q{quarter}`.
    """
    ns = nw  # It works if we use `ns = pl` and call `.to_native()` on the df below
    df = nw.from_arrow(data, backend=pl)
    year_expr = ns.col("value_timestamp").dt.year()
    month_expr = ns.col("value_timestamp").dt.month()
    quarter_expr = ((month_expr - 1) // 3 + 1)
    expr = (
        year_expr.cast(ns.String)
        + ns.lit("Q")
        + quarter_expr.cast(ns.String)
    ).alias("_quarter")
    for key, group in df.group_by(expr):
        quarter = key[0] if isinstance(key, tuple) else key
        yield quarter, group.to_arrow()


identifiers = list("ABCDE")
start_date = pd.Timestamp.now().date() + pd.DateOffset(days=1)
end_date = start_date + pd.DateOffset(years=1)
index = pd.date_range(start_date, end_date, freq="D", name="value_timestamp")
rng = np.random.default_rng()
data = rng.standard_normal((len(index), len(identifiers))).cumsum(axis=0)
data += rng.uniform(40, 60, size=(1, len(identifiers)))
data = pd.DataFrame(data, index=index, columns=identifiers).reset_index()


quarters = {
    quarter: group
    for quarter, group
    in partition_fn(data)
}
quarters

Cell In[25], line 26, in partition_fn(data)
     20 quarter_expr = ((month_expr - 1) // 3 + 1)
     21 expr = (
     22     year_expr.cast(ns.String)
     23     + ns.lit("Q")
     24     + quarter_expr.cast(ns.String)
     25 ).alias("_quarter")
---> 26 for key, group in df.group_by(expr):
     27     quarter = key[0] if isinstance(key, tuple) else key
     28     yield quarter, group.to_arrow()

File [.pixi/envs/py312/lib/python3.12/site-packages/narwhals/dataframe.py:1726], in DataFrame.group_by(self, drop_null_keys, *keys)
   1721     from narwhals.exceptions import ComputeError
   1723     msg = (
   1724         "Group by is not supported with keys that are not elementwise expressions"
   1725     )
-> 1726     raise ComputeError(msg)
   1728 return GroupBy(self, expr_flat_keys, drop_null_keys=drop_null_keys)

ComputeError: Group by is not supported with keys that are not elementwise expressions

Expected results

It would be great if this also worked with narwhals.

Actual results

See error above

Please run narwhals.show_version() and enter the output below.

System:
    python: 3.12.11 | packaged by conda-forge | (main, Jun  4 2025, 14:45:31) [GCC 13.3.0]
executable: ~/.pixi/envs/py312/bin/python
   machine: Linux-6.14.0-1012-aws-x86_64-with-glibc2.39

Python dependencies:
     narwhals: 2.5.0
        numpy: 1.26.4
       pandas: 2.2.3
        modin: 
         cudf: 
      pyarrow: 21.0.0
      pyspark: 
       polars: 1.33.1
         dask: 2025.9.1
       duckdb: 1.3.2+g2063dda
         ibis: 
     sqlframe:

Relevant log output

[MarcoGorelli]

thanks for the report - we can definitely relax this restriction for the eager case (which is what you're using by the looks of it?)

for the lazy case, sql forbids group by clauses from containing window functions 😩 , but we may be able to come up with a workaround later

[dhirschfeld]

Yeah, eager is fine.

[dangotbanned]

Current workaround

This'll essentially be what we have to do for pyarrow and pandas anyway

from collections.abc import Generator
import narwhals as nw
import numpy as np
import pandas as pd
import polars as pl
from narwhals.typing import IntoDataFrameT



def partition_fn(
    data: IntoDataFrameT,
) -> Generator[tuple[str, pa.Table], None, None]:
    """Group by unique quarters in `value_timestamp` and yield `(quarter, group)`
    tuples where the quarters are formatted as `{year}Q{quarter}`.
    """
    ns = nw  # It works if we use `ns = pl` and call `.to_native()` on the df below
    df = nw.from_arrow(data, backend=pl)
    year_expr = ns.col("value_timestamp").dt.year()
    month_expr = ns.col("value_timestamp").dt.month()
    quarter_expr = ((month_expr - 1) // 3 + 1)
    expr = (
        year_expr.cast(ns.String)
        + ns.lit("Q")
        + quarter_expr.cast(ns.String)
    ).alias("_quarter")
-   for key, group in df.group_by(expr):
+   for key, group in df.with_columns(expr).group_by("_quarter"):
        quarter = key[0] if isinstance(key, tuple) else key
        yield quarter, group.to_arrow()


identifiers = list("ABCDE")
start_date = pd.Timestamp.now().date() + pd.DateOffset(days=1)
end_date = start_date + pd.DateOffset(years=1)
index = pd.date_range(start_date, end_date, freq="D", name="value_timestamp")
rng = np.random.default_rng()
data = rng.standard_normal((len(index), len(identifiers))).cumsum(axis=0)
data += rng.uniform(40, 60, size=(1, len(identifiers)))
data = pd.DataFrame(data, index=index, columns=identifiers).reset_index()


quarters = {
    quarter: group
    for quarter, group
    in partition_fn(data)
}
quarters

Hiding as it is HUGE

Show output

{'2026Q1': pyarrow.Table
 value_timestamp: timestamp[ns]
 A: double
 B: double
 C: double
 D: double
 E: double
 _quarter: large_string
 ----
 value_timestamp: [[2026-01-01 00:00:00.000000000,2026-01-02 00:00:00.000000000,2026-01-03 00:00:00.000000000,2026-01-04 00:00:00.000000000,2026-01-05 00:00:00.000000000,...,2026-03-27 00:00:00.000000000,2026-03-28 00:00:00.000000000,2026-03-29 00:00:00.000000000,2026-03-30 00:00:00.000000000,2026-03-31 00:00:00.000000000]]
 A: [[72.98393459892353,70.62304585703541,72.38274256860844,71.84801391031067,71.41420526413366,...,68.294475069738,68.5062297307556,69.03220180556598,68.48427274402947,69.56917186542348]]
 B: [[55.68604454982709,54.62403710916328,56.144645189369406,56.38731518194045,57.45037318853473,...,54.23977316992278,53.812227780165195,53.14284423471326,52.47416491635697,53.57975652461302]]
 C: [[39.537732642566596,36.945734693180306,35.52103721999027,36.22410360457894,37.76010117733965,...,15.829149868507024,14.889584712361085,15.05408397703296,14.839119111604191,14.144705397665085]]
 D: [[69.8608501933011,68.98393968098871,70.05143289503968,68.90893196346936,69.2472858275398,...,76.64737208595344,75.85230908987513,77.53403596314074,77.61436836691368,77.52813834462734]]
 E: [[59.69960133855783,59.86763984367212,62.27805723868345,61.972351286782654,61.227529970374015,...,56.013375401931455,53.51394812095682,54.57315536198863,52.71995394203606,51.91229751408749]]
 _quarter: [["2026Q1","2026Q1","2026Q1","2026Q1","2026Q1",...,"2026Q1","2026Q1","2026Q1","2026Q1","2026Q1"]],
 '2025Q3': pyarrow.Table
 value_timestamp: timestamp[ns]
 A: double
 B: double
 C: double
 D: double
 E: double
 _quarter: large_string
 ----
 value_timestamp: [[2025-09-26 00:00:00.000000000,2025-09-27 00:00:00.000000000,2025-09-28 00:00:00.000000000,2025-09-29 00:00:00.000000000,2025-09-30 00:00:00.000000000]]
 A: [[58.045389799969044,58.48328130988173,59.02726912085635,59.48796460648526,60.761100922379356]]
 B: [[56.53355911857172,54.83825432816137,55.124404855092756,55.54743662422056,55.42958145635614]]
 C: [[53.63518143756124,53.06620852936836,53.31236378162644,53.74814848408022,53.14337320469233]]
 D: [[47.9120624345825,47.330924334395654,47.69975463100836,47.80653254902246,47.95961424095559]]
 E: [[47.7603746701041,48.9043788571,48.23852406763661,49.782598703884915,49.208143387691656]]
 _quarter: [["2025Q3","2025Q3","2025Q3","2025Q3","2025Q3"]],
 '2025Q4': pyarrow.Table
 value_timestamp: timestamp[ns]
 A: double
 B: double
 C: double
 D: double
 E: double
 _quarter: large_string
 ----
 value_timestamp: [[2025-10-01 00:00:00.000000000,2025-10-02 00:00:00.000000000,2025-10-03 00:00:00.000000000,2025-10-04 00:00:00.000000000,2025-10-05 00:00:00.000000000,...,2025-12-27 00:00:00.000000000,2025-12-28 00:00:00.000000000,2025-12-29 00:00:00.000000000,2025-12-30 00:00:00.000000000,2025-12-31 00:00:00.000000000]]
 A: [[59.23015211679572,60.72561815480897,60.85142240055712,60.8137878744721,61.532052839611275,...,73.50114341428373,74.88371099852166,75.21772281261696,74.98570111685697,72.11995919351389]]
 B: [[55.30978959816378,54.587937600205784,55.08585376909712,55.11728615717832,56.82027661522173,...,57.33180182231007,55.830185060915454,56.279938122772464,55.05094111309544,56.788396788139536]]
 C: [[52.18606480272245,53.08471809804199,52.565569729957566,51.22731571568726,52.55092805332306,...,44.187288215621805,42.49591690535935,41.972983053418076,41.20808399260679,41.39605034780491]]
 D: [[49.02380949105102,50.50928603818241,51.477998105804005,51.879304046890866,52.988271445479334,...,65.64124348806887,66.28610753325809,68.97472160291991,69.40968309014862,70.67365686025688]]
 E: [[49.609894549366324,52.111112525329304,54.427540051356694,55.328972565501616,55.55353567998921,...,56.2356925401125,57.88149615317577,57.51961896727023,57.332162598127965,57.604572740779936]]
 _quarter: [["2025Q4","2025Q4","2025Q4","2025Q4","2025Q4",...,"2025Q4","2025Q4","2025Q4","2025Q4","2025Q4"]],
 '2026Q3': pyarrow.Table
 value_timestamp: timestamp[ns]
 A: double
 B: double
 C: double
 D: double
 E: double
 _quarter: large_string
 ----
 value_timestamp: [[2026-07-01 00:00:00.000000000,2026-07-02 00:00:00.000000000,2026-07-03 00:00:00.000000000,2026-07-04 00:00:00.000000000,2026-07-05 00:00:00.000000000,...,2026-09-22 00:00:00.000000000,2026-09-23 00:00:00.000000000,2026-09-24 00:00:00.000000000,2026-09-25 00:00:00.000000000,2026-09-26 00:00:00.000000000]]
 A: [[75.6691008860273,76.07880644866424,76.62869801090352,77.61152937748528,76.61020422779734,...,65.54210022923634,64.65528448931593,65.93648054356683,65.81918763598806,66.03623282170295]]
 B: [[43.68686449056686,43.737993473815045,43.76655079319435,43.57024049388463,41.55562444239107,...,42.4737263823149,42.339636529432234,41.115097557983546,41.24586032053971,40.81285227870682]]
 C: [[38.04223135998241,38.37965047917905,38.049055073082954,38.037124032119436,37.55576805457731,...,32.06567113830978,33.8827995784638,34.51100018692816,33.94151897749688,33.709264177649835]]
 D: [[71.43355450261294,70.43229716310915,69.46345956740394,68.13209786134794,67.72162995643235,...,65.81556325206161,65.81879846870194,65.34153059324532,64.74497654827752,65.87648644304576]]
 E: [[68.62630662459415,68.81721553404805,70.22208046401657,71.02916193759759,70.94102864995565,...,80.91721888002348,81.42495561558164,79.21985834444237,78.68450076820409,80.23459357374544]]
 _quarter: [["2026Q3","2026Q3","2026Q3","2026Q3","2026Q3",...,"2026Q3","2026Q3","2026Q3","2026Q3","2026Q3"]],
 '2026Q2': pyarrow.Table
 value_timestamp: timestamp[ns]
 A: double
 B: double
 C: double
 D: double
 E: double
 _quarter: large_string
 ----
 value_timestamp: [[2026-04-01 00:00:00.000000000,2026-04-02 00:00:00.000000000,2026-04-03 00:00:00.000000000,2026-04-04 00:00:00.000000000,2026-04-05 00:00:00.000000000,...,2026-06-26 00:00:00.000000000,2026-06-27 00:00:00.000000000,2026-06-28 00:00:00.000000000,2026-06-29 00:00:00.000000000,2026-06-30 00:00:00.000000000]]
 A: [[68.54397474638454,66.40112719509571,65.69090893135095,64.9156473101267,64.44833464632198,...,75.46281736800428,74.62541879380375,74.31321978474686,74.25917987906733,75.22194491562198]]
 B: [[52.53802533746148,52.19116338374985,52.6813798447387,51.58430525874382,50.11691949078114,...,45.438240934695095,44.08017639303533,44.05275063100605,45.283484711274355,44.410363037291205]]
 C: [[14.709743450588192,12.94008928538829,13.295215758916676,13.953093171323232,11.711827562207382,...,38.43595927991365,38.32757652008906,37.94037987630125,38.7959685286914,38.838573273671344]]
 D: [[78.90689176282692,78.4830030743672,79.61114476641539,79.03550464084162,77.71307073452164,...,72.23547973556381,70.9545670916331,70.41506641146951,70.09486418333685,70.91779746586799]]
 E: [[52.63087610889883,52.600794709852266,52.94024498797269,53.72268145044532,53.46501616631832,...,67.23071587847569,65.74184616153742,66.787390693596,68.11052079595147,70.51691602899254]]
 _quarter: [["2026Q2","2026Q2","2026Q2","2026Q2","2026Q2",...,"2026Q2","2026Q2","2026Q2","2026Q2","2026Q2"]]}

Side note

I didn't expect this to work for all eager backends, but it looks like you can simplify expr to:

-   expr = (year_expr.cast(ns.String) + ns.lit("Q") + quarter_expr.cast(ns.String)).alias(
-       "_quarter"
-   )
+   expr = ns.concat_str(year_expr, ns.lit("Q"), quarter_expr).alias("_quarter")

[dhirschfeld]

If you add the column with .with_columns(expr) you then also have to remove it before returning the group if you don't want it to appear in the output. That's a bit less nice than being able to group by the expression directly without mutating the DataFrame.

[MarcoGorelli]

This'll essentially be what we have to do for pyarrow and pandas anyway

we won't need to do this for pandas nor pyarrow, we can just remove the check

[MarcoGorelli]

just something like

diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 7a3e4ce90..e5301dfff 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -23,6 +23,7 @@ from narwhals._expression_parsing import (
     is_scalar_like,
 )
 from narwhals._typing import Arrow, Pandas, _LazyAllowedImpl, _LazyFrameCollectImpl
+from narwhals.exceptions import ComputeError
 from narwhals._utils import (
     Implementation,
     Version,
@@ -1715,14 +1716,10 @@ class DataFrame(BaseFrame[DataFrameT]):
             k if is_expr else col(k)
             for k, is_expr in zip_strict(flat_keys, key_is_expr_or_series)
         ]
-        expr_flat_keys, kinds = self._flatten_and_extract(*_keys)
-
-        if not all(kind is ExprKind.ELEMENTWISE for kind in kinds):
-            from narwhals.exceptions import ComputeError
+        expr_flat_keys, _kinds = self._flatten_and_extract(*_keys)
 
-            msg = (
-                "Group by is not supported with keys that are not elementwise expressions"
-            )
+        if not all(x._metadata.preserves_length for x in _keys):
+            msg = "Group by is not supported with keys that don't preserve length."
             raise ComputeError(msg)
 
         return GroupBy(self, expr_flat_keys, drop_null_keys=drop_null_keys)

and then adjust the tests

[dangotbanned]

This'll essentially be what we have to do for pyarrow and pandas anyway

we won't need to do this for pandas nor pyarrow, we can just remove the check

Apologies for swooping in with the workaround 😅

A more complete version is that would be equivalent to what we do here:

narwhals/narwhals/_compliant/group_by.py

Lines 118 to 119 in 63c5022

	return (
	compliant_frame.with_columns(*safe_keys),

Which happens before the group_by:

narwhals/narwhals/_arrow/group_by.py

Lines 49 to 60 in 63c5022

	def __init__(
	self,
	df: ArrowDataFrame,
	keys: Sequence[ArrowExpr] | Sequence[str],
	/,
	*,
	drop_null_keys: bool,
	) -> None:
	self._df = df
	frame, self._keys, self._output_key_names = self._parse_keys(df, keys=keys)
	self._compliant_frame = frame.drop_nulls(self._keys) if drop_null_keys else frame
	self._grouped = pa.TableGroupBy(self.compliant.native, self._keys)

But as @MarcoGorelli said, we are raising here prematurely

[dangotbanned]

If you add the column with .with_columns(expr) you then also have to remove it before returning the group if you don't want it to appear in the output.

@dhirschfeld would you be able to make a smaller repro of (#3151 (comment)) so this difference is more obvious please?

I'm about to start working on GroupBy.__iter__ in (#3143) and it would be super helpful to make sure this is avoided there 🙂