Accept string column names in GroupingSet factory methods

GroupingSet.rollup(), .cube(), and .grouping_sets() now accept both
Expr objects and string column names, consistent with DataFrame.aggregate().

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timsaucer

python/datafusion/expr.py

@@ -1445,7 +1445,7 @@ class GroupingSet:
     """
 
     @staticmethod
-    def rollup(*exprs: Expr) -> Expr:
+    def rollup(*exprs: Expr | str) -> Expr:
         """Create a ``ROLLUP`` grouping set for use with ``aggregate()``.
 
         ``ROLLUP`` generates all prefixes of the given column list as
@@ -1455,7 +1455,8 @@ def rollup(*exprs: Expr) -> Expr:
         This is equivalent to ``GROUP BY ROLLUP(a, b)`` in SQL.
 
         Args:
-            *exprs: Column expressions to include in the rollup.
+            *exprs: Column expressions or column name strings to
+                include in the rollup.
 
         Examples:
             >>> import datafusion as dfn
@@ -1474,11 +1475,11 @@ def rollup(*exprs: Expr) -> Expr:
             :py:meth:`cube`, :py:meth:`grouping_sets`,
             :py:func:`~datafusion.functions.grouping`
         """
-        args = [e.expr for e in exprs]
+        args = [_to_raw_expr(e) for e in exprs]
         return Expr(expr_internal.GroupingSet.rollup(*args))
 
     @staticmethod
-    def cube(*exprs: Expr) -> Expr:
+    def cube(*exprs: Expr | str) -> Expr:
         """Create a ``CUBE`` grouping set for use with ``aggregate()``.
 
         ``CUBE`` generates all possible subsets of the given column list
@@ -1488,7 +1489,8 @@ def cube(*exprs: Expr) -> Expr:
         This is equivalent to ``GROUP BY CUBE(a, b)`` in SQL.
 
         Args:
-            *exprs: Column expressions to include in the cube.
+            *exprs: Column expressions or column name strings to
+                include in the cube.
 
         Examples:
             With a single column, ``cube`` behaves identically to
@@ -1510,23 +1512,25 @@ def cube(*exprs: Expr) -> Expr:
             :py:meth:`rollup`, :py:meth:`grouping_sets`,
             :py:func:`~datafusion.functions.grouping`
         """
-        args = [e.expr for e in exprs]
+        args = [_to_raw_expr(e) for e in exprs]
         return Expr(expr_internal.GroupingSet.cube(*args))
 
     @staticmethod
-    def grouping_sets(*expr_lists: list[Expr]) -> Expr:
+    def grouping_sets(*expr_lists: list[Expr | str]) -> Expr:
         """Create explicit grouping sets for use with ``aggregate()``.
 
-        Each argument is a list of column expressions representing one
-        grouping set. For example, ``grouping_sets([a], [b])`` groups
-        by ``a`` alone and by ``b`` alone in a single query.
+        Each argument is a list of column expressions or column name
+        strings representing one grouping set. For example,
+        ``grouping_sets([a], [b])`` groups by ``a`` alone and by ``b``
+        alone in a single query.
 
         This is equivalent to ``GROUP BY GROUPING SETS ((a), (b))`` in
         SQL.
 
         Args:
             *expr_lists: Each positional argument is a list of
-                expressions forming one grouping set.
+                expressions or column name strings forming one
+                grouping set.
 
         Examples:
             >>> import datafusion as dfn
@@ -1552,5 +1556,5 @@ def grouping_sets(*expr_lists: list[Expr]) -> Expr:
             :py:meth:`rollup`, :py:meth:`cube`,
             :py:func:`~datafusion.functions.grouping`
         """
-        raw_lists = [[e.expr for e in lst] for lst in expr_lists]
+        raw_lists = [[_to_raw_expr(e) for e in lst] for lst in expr_lists]
         return Expr(expr_internal.GroupingSet.grouping_sets(*raw_lists))

python/tests/test_functions.py

@@ -1844,8 +1844,10 @@ def test_percentile_cont(func, filter_expr, expected):
     [
         (GroupingSet.rollup(column("a")), [0, 0, 1], [30, 30, 60]),
         (GroupingSet.cube(column("a")), [0, 0, 1], [30, 30, 60]),
+        (GroupingSet.rollup("a"), [0, 0, 1], [30, 30, 60]),
+        (GroupingSet.cube("a"), [0, 0, 1], [30, 30, 60]),
     ],
-    ids=["rollup", "cube"],
+    ids=["rollup", "cube", "rollup_str", "cube_str"],
 )
 def test_grouping_set_single_column(
     grouping_set_expr, expected_grouping, expected_sums
@@ -1870,8 +1872,10 @@ def test_grouping_set_single_column(
         (GroupingSet.rollup(column("a"), column("b")), 6),
         # cube(a, b) => (a,b), (a), (b), () => 3 + 2 + 2 + 1 = 8
         (GroupingSet.cube(column("a"), column("b")), 8),
+        (GroupingSet.rollup("a", "b"), 6),
+        (GroupingSet.cube("a", "b"), 8),
     ],
-    ids=["rollup", "cube"],
+    ids=["rollup", "cube", "rollup_str", "cube_str"],
 )
 def test_grouping_set_multi_column(grouping_set_expr, expected_rows):
     ctx = SessionContext()
@@ -1884,12 +1888,20 @@ def test_grouping_set_multi_column(grouping_set_expr, expected_rows):
     assert total_rows == expected_rows
 
 
-def test_grouping_sets_explicit():
+@pytest.mark.parametrize(
+    "grouping_set_expr",
+    [
+        GroupingSet.grouping_sets([column("a")], [column("b")]),
+        GroupingSet.grouping_sets(["a"], ["b"]),
+    ],
+    ids=["expr", "str"],
+)
+def test_grouping_sets_explicit(grouping_set_expr):
     # Each row's grouping() value tells you which columns are aggregated across.
     ctx = SessionContext()
     df = ctx.from_pydict({"a": ["x", "x", "y"], "b": ["m", "n", "m"], "c": [1, 2, 3]})
     result = df.aggregate(
-        [GroupingSet.grouping_sets([column("a")], [column("b")])],
+        [grouping_set_expr],
         [
             f.sum(column("c")).alias("s"),
             f.grouping(column("a")),