feat: enhance aggregate method to support tuple inputs for group_by and aggs

Author: kosiew

python/datafusion/dataframe.py

@@ -22,6 +22,7 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -52,7 +53,7 @@
 
 if TYPE_CHECKING:
     import pathlib
-    from typing import Callable, Sequence
+    from typing import Callable
 
     import pandas as pd
     import polars as pl
@@ -523,20 +524,28 @@ def with_column_renamed(self, old_name: str, new_name: str) -> DataFrame:
 
     def aggregate(
         self,
-        group_by: list[Expr | str] | Expr | str,
-        aggs: list[Expr] | Expr,
+        group_by: Sequence[Expr | str] | Expr | str,
+        aggs: Sequence[Expr] | Expr,
     ) -> DataFrame:
         """Aggregates the rows of the current DataFrame.
 
         Args:
-            group_by: List of expressions or column names to group by.
-            aggs: List of expressions to aggregate.
+            group_by: Sequence of expressions or column names to group by.
+            aggs: Sequence of expressions to aggregate.
 
         Returns:
             DataFrame after aggregation.
         """
-        group_by_list = group_by if isinstance(group_by, list) else [group_by]
-        aggs_list = aggs if isinstance(aggs, list) else [aggs]
+        group_by_list = (
+            list(group_by)
+            if isinstance(group_by, Sequence) and not isinstance(group_by, (Expr, str))
+            else [group_by]
+        )
+        aggs_list = (
+            list(aggs)
+            if isinstance(aggs, Sequence) and not isinstance(aggs, Expr)
+            else [aggs]
+        )
 
         group_by_exprs = expr_list_to_raw_expr_list(group_by_list)
         aggs_exprs = [ensure_expr(agg) for agg in aggs_list]

python/tests/test_dataframe.py

@@ -303,6 +303,18 @@ def test_aggregate_string_and_expression_equivalent(df):
     assert result_str == result_expr
 
 
+def test_aggregate_tuple_group_by(df):
+    result_list = df.aggregate(["a"], [f.count()]).sort("a").to_pydict()
+    result_tuple = df.aggregate(("a",), [f.count()]).sort("a").to_pydict()
+    assert result_tuple == result_list
+
+
+def test_aggregate_tuple_aggs(df):
+    result_list = df.aggregate("a", [f.count()]).sort("a").to_pydict()
+    result_tuple = df.aggregate("a", (f.count(),)).sort("a").to_pydict()
+    assert result_tuple == result_list
+
+
 def test_filter_string_unsupported(df):
     with pytest.raises(TypeError, match=re.escape(EXPR_TYPE_ERROR)):
         df.filter("a > 1")