Enhance DataFrame.aggregate method to accept column names as group by expressions

Author: kosiew

docs/source/user-guide/dataframe/index.rst

@@ -142,14 +142,19 @@ For such methods, you can pass column names directly:
 
 .. code-block:: python
 
+    from datafusion import col, functions as f
+
     df.sort('id')
+    df.aggregate('id', [f.count(col('value'))])
 
 The same operation can also be written with an explicit column expression:
 
 .. code-block:: python
 
-    from datafusion import col
+    from datafusion import col, functions as f
+
     df.sort(col('id'))
+    df.aggregate(col('id'), [f.count(col('value'))])
 
 Whenever an argument represents an expression—such as in
 :py:meth:`~datafusion.DataFrame.filter` or

python/datafusion/dataframe.py

@@ -503,23 +503,27 @@ def with_column_renamed(self, old_name: str, new_name: str) -> DataFrame:
         return DataFrame(self.df.with_column_renamed(old_name, new_name))
 
     def aggregate(
-        self, group_by: list[Expr] | Expr, aggs: list[Expr] | Expr
+        self,
+        group_by: list[Expr | str] | Expr | str,
+        aggs: list[Expr] | Expr,
     ) -> DataFrame:
         """Aggregates the rows of the current DataFrame.
 
         Args:
-            group_by: List of expressions to group by.
+            group_by: List of expressions or column names to group by.
             aggs: List of expressions to aggregate.
 
         Returns:
             DataFrame after aggregation.
         """
-        group_by = group_by if isinstance(group_by, list) else [group_by]
-        aggs = aggs if isinstance(aggs, list) else [aggs]
+        group_by_list = group_by if isinstance(group_by, list) else [group_by]
+        aggs_list = aggs if isinstance(aggs, list) else [aggs]
 
-        group_by = [e.expr for e in group_by]
-        aggs = [e.expr for e in aggs]
-        return DataFrame(self.df.aggregate(group_by, aggs))
+        group_by_exprs = [
+            Expr.column(e).expr if isinstance(e, str) else e.expr for e in group_by_list
+        ]
+        aggs_exprs = [e.expr for e in aggs_list]
+        return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
     def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
         """Sort the DataFrame by the specified sorting expressions or column names.