Refactor DataFrame methods to use _to_expr_list for expression handling

Author: kosiew

python/datafusion/dataframe.py

@@ -40,7 +40,7 @@
 from datafusion._internal import DataFrame as DataFrameInternal
 from datafusion._internal import ParquetColumnOptions as ParquetColumnOptionsInternal
 from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal
-from datafusion.expr import Expr, SortExpr, sort_or_default
+from datafusion.expr import Expr, SortExpr, _to_expr_list, sort_or_default
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
 
@@ -394,9 +394,7 @@ def select(self, *exprs: Expr | str) -> DataFrame:
             df = df.select("a", col("b"), col("a").alias("alternate_a"))
 
         """
-        exprs_internal = [
-            Expr.column(arg).expr if isinstance(arg, str) else arg.expr for arg in exprs
-        ]
+        exprs_internal = _to_expr_list(exprs)
         return DataFrame(self.df.select(*exprs_internal))
 
     def drop(self, *columns: str) -> DataFrame:
@@ -548,19 +546,8 @@ def aggregate(
         group_by_list = group_by if isinstance(group_by, list) else [group_by]
         aggs_list = aggs if isinstance(aggs, list) else [aggs]
 
-        group_by_exprs = [
-            Expr.column(e).expr if isinstance(e, str) else e.expr for e in group_by_list
-        ]
-
-        aggs_exprs: list[expr_internal.Expr] = []
-        for agg in aggs_list:
-            if not isinstance(agg, Expr):
-                msg = (
-                    f"Expected Expr, got {type(agg).__name__}. "
-                    "Use col() or lit() to construct expressions."
-                )
-                raise TypeError(msg)
-            aggs_exprs.append(agg.expr)
+        group_by_exprs = _to_expr_list(group_by_list)
+        aggs_exprs = _to_expr_list(aggs_list)
         return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
     def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
@@ -575,10 +562,14 @@ def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
         Returns:
             DataFrame after sorting.
         """
-        exprs_raw = [
-            sort_or_default(Expr.column(expr) if isinstance(expr, str) else expr)
-            for expr in exprs
-        ]
+        expr_seq = [e for e in exprs if not isinstance(e, SortExpr)]
+        raw_exprs_iter = iter(_to_expr_list(expr_seq))
+        exprs_raw = []
+        for e in exprs:
+            if isinstance(e, SortExpr):
+                exprs_raw.append(sort_or_default(e))
+            else:
+                exprs_raw.append(sort_or_default(Expr(next(raw_exprs_iter))))
         return DataFrame(self.df.sort(*exprs_raw))
 
     def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame:

python/datafusion/expr.py

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence
 
 import pyarrow as pa
 
@@ -215,6 +215,11 @@
 ]
 
 
+def _to_expr_list(exprs: Sequence[Expr | str]) -> list[expr_internal.Expr]:
+    """Convert a sequence of expressions or column names to raw expressions."""
+    return [Expr.column(e).expr if isinstance(e, str) else e.expr for e in exprs]
+
+
 def expr_list_to_raw_expr_list(
     expr_list: Optional[list[Expr] | Expr],
 ) -> Optional[list[expr_internal.Expr]]: