Repository: ibis-project/ibis

Issue:
Enable use of functions in GroupedTableExpr.mutate



Discussion:

Resolving patch:
diff --git a/ibis/expr/groupby.py b/ibis/expr/groupby.py
--- a/ibis/expr/groupby.py
+++ b/ibis/expr/groupby.py
@@ -41,7 +41,7 @@ def __init__(self, table, by, having=None, order_by=None, window=None):
 
     def __getitem__(self, args):
         # Shortcut for projection with window functions
-        return self.projection(args)
+        return self.projection(list(args))
 
     def __getattr__(self, attr):
         if hasattr(self.table, attr):
@@ -99,14 +99,21 @@ def order_by(self, expr):
 
     def mutate(self, exprs=None, **kwds):
         """
-        Returns a table projection with analytic / window functions applied
+        Returns a table projection with analytic / window functions
+        applied. Any arguments can be functions.
+
+        Parameters
+        ----------
+        exprs : list, default None
+        kwds : key=value pairs
 
         Examples
         --------
         expr = (table
                 .group_by('foo')
                 .order_by(ibis.desc('bar'))
-                .mutate(qux=table.baz.lag()))
+                .mutate(qux=lambda x: x.baz.lag(),
+                        qux2=table.baz.lead()))
 
         Returns
         -------
@@ -117,14 +124,22 @@ def mutate(self, exprs=None, **kwds):
         else:
             exprs = util.promote_list(exprs)
 
-        for k, v in sorted(kwds.items()):
+        kwd_names = list(kwds.keys())
+        kwd_values = list(kwds.values())
+        kwd_values = self.table._resolve(kwd_values)
+
+        for k, v in sorted(zip(kwd_names, kwd_values)):
             exprs.append(v.name(k))
 
         return self.projection([self.table] + exprs)
 
     def projection(self, exprs):
+        """
+        Like mutate, but do not include existing table columns
+        """
         w = self._get_window()
         windowed_exprs = []
+        exprs = self.table._resolve(exprs)
         for expr in exprs:
             expr = L.windowize_function(expr, w=w)
             windowed_exprs.append(expr)


Test patch:
diff --git a/ibis/expr/tests/test_table.py b/ibis/expr/tests/test_table.py
--- a/ibis/expr/tests/test_table.py
+++ b/ibis/expr/tests/test_table.py
@@ -1285,6 +1285,28 @@ def g(x):
         expected = self.table.mutate(foo=g)
         assert_equal(result, expected)
 
+    def test_groupby_mutate(self):
+        t = self.table
+
+        g = t.group_by('g').order_by('f')
+        expr = g.mutate(foo=lambda x: x.f.lag(),
+                        bar=lambda x: x.f.rank())
+        expected = g.mutate(foo=t.f.lag(),
+                            bar=t.f.rank())
+
+        assert_equal(expr, expected)
+
+    def test_groupby_projection(self):
+        t = self.table
+
+        g = t.group_by('g').order_by('f')
+        expr = g.projection([lambda x: x.f.lag().name('foo'),
+                             lambda x: x.f.rank().name('bar')])
+        expected = g.projection([t.f.lag().name('foo'),
+                                 t.f.rank().name('bar')])
+
+        assert_equal(expr, expected)
+
     def test_set_column(self):
         def g(x):
             return x.f * 2