fix(polars): project first when creating computed grouping keys

ibis-project · Dec 6, 2023 · 7f9fdd4 · 7f9fdd4
1 parent 6817981
commit 7f9fdd4
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/ibis/backends/polars/compiler.py b/ibis/backends/polars/compiler.py
@@ -267,9 +267,11 @@ def aggregation(op, **kw):
             )
         )
 
+    # project first to handle computed group by columns
+    lf = lf.with_columns([translate(arg, **kw) for arg in op.by])
+
     if op.by:
-        group_by = [translate(arg, **kw) for arg in op.by]
-        lf = lf.group_by(group_by).agg
+        lf = lf.group_by([pl.col(by.name) for by in op.by]).agg
     else:
         lf = lf.select