Refactor expression handling by replacing _ensure_expr and _to_expr_list with expr_list_to_raw_expr_list for improved clarity and consistency in DataFrame methods.

Author: kosiew

python/datafusion/dataframe.py

@@ -40,13 +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,
-    _ensure_expr,
-    _to_expr_list,
-    sort_or_default,
-)
+from datafusion.expr import Expr, SortExpr, expr_list_to_raw_expr_list, sort_or_default
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
 
@@ -400,19 +394,7 @@ def select(self, *exprs: Expr | str) -> DataFrame:
             df = df.select("a", col("b"), col("a").alias("alternate_a"))
 
         """
-        checked_exprs: list[Expr | str] = []
-        for expr in exprs:
-            if isinstance(expr, SortExpr):
-                checked_exprs.append(expr.expr())
-            elif isinstance(expr, (Expr, str)):
-                checked_exprs.append(expr)
-            else:
-                msg = (
-                    f"Expected Expr or column name, got {type(expr).__name__}. "
-                    "Use col() or lit() to construct expressions."
-                )
-                raise TypeError(msg)
-        exprs_internal = _to_expr_list(checked_exprs)
+        exprs_internal = expr_list_to_raw_expr_list(exprs)
         return DataFrame(self.df.select(*exprs_internal))
 
     def drop(self, *columns: str) -> DataFrame:
@@ -442,7 +424,7 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
-            df = df.filter(_ensure_expr(p))
+            df = df.filter(expr_list_to_raw_expr_list(p)[0])
         return DataFrame(df)
 
     def with_column(self, name: str, expr: Expr) -> DataFrame:
@@ -455,7 +437,7 @@ def with_column(self, name: str, expr: Expr) -> DataFrame:
         Returns:
             DataFrame with the new column.
         """
-        return DataFrame(self.df.with_column(name, _ensure_expr(expr)))
+        return DataFrame(self.df.with_column(name, expr_list_to_raw_expr_list(expr)[0]))
 
     def with_columns(
         self, *exprs: Expr | Iterable[Expr], **named_exprs: Expr
@@ -487,12 +469,12 @@ def _simplify_expression(
             expr_list = []
             for expr in exprs:
                 if isinstance(expr, Iterable) and not isinstance(expr, Expr):
-                    expr_list.extend(_ensure_expr(inner_expr) for inner_expr in expr)
+                    expr_list.extend(expr_list_to_raw_expr_list(inner_expr)[0] for inner_expr in expr)
                 else:
-                    expr_list.append(_ensure_expr(expr))
+                    expr_list.append(expr_list_to_raw_expr_list(expr)[0])
             if named_exprs:
                 for alias, expr in named_exprs.items():
-                    _ensure_expr(expr)
+                    expr_list_to_raw_expr_list(expr)[0]
                     expr_list.append(expr.alias(alias).expr)
             return expr_list
 
@@ -537,7 +519,7 @@ def aggregate(
         group_by_exprs = [
             Expr.column(e).expr if isinstance(e, str) else e.expr for e in group_by_list
         ]
-        aggs_exprs = [_ensure_expr(agg) for agg in aggs_list]
+        aggs_exprs = [expr_list_to_raw_expr_list(agg)[0] for agg in aggs_list]
         return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
     def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
@@ -552,20 +534,9 @@ def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
         Returns:
             DataFrame after sorting.
         """
-        expr_seq: list[Expr | str] = []
-        for e in exprs:
-            if isinstance(e, SortExpr):
-                continue
-            if isinstance(e, (Expr, str)):
-                expr_seq.append(e)
-            else:
-                msg = (
-                    f"Expected Expr or column name, got {type(e).__name__}. "
-                    "Use col() or lit() to construct expressions."
-                )
-                raise TypeError(msg)
-        raw_exprs_iter = iter(_to_expr_list(expr_seq))
-        exprs_raw: list[Any] = []
+        expr_seq = [e for e in exprs if not isinstance(e, SortExpr)]
+        raw_exprs_iter = iter(expr_list_to_raw_expr_list(expr_seq))
+        exprs_raw = []
         for e in exprs:
             if isinstance(e, SortExpr):
                 exprs_raw.append(sort_or_default(e))
@@ -794,7 +765,7 @@ def join_on(
         Returns:
             DataFrame after join.
         """
-        exprs = [_ensure_expr(expr) for expr in on_exprs]
+        exprs = [expr_list_to_raw_expr_list(expr)[0] for expr in on_exprs]
         return DataFrame(self.df.join_on(right.df, exprs, how))
 
     def explain(self, verbose: bool = False, analyze: bool = False) -> None:

python/datafusion/expr.py

@@ -215,29 +215,15 @@
 ]
 
 
-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 _ensure_expr(value: Any) -> expr_internal.Expr:
-    """Return the internal expression or raise if the value is not an Expr."""
-    if isinstance(value, Expr):
-        return value.expr
-    msg = (
-        f"Expected Expr, got {type(value).__name__}. "
-        "Use col() or lit() to construct expressions."
-    )
-    raise TypeError(msg)
-
-
 def expr_list_to_raw_expr_list(
-    expr_list: Optional[list[Expr] | Expr],
+    expr_list: Optional[Sequence[Expr | str] | Expr | str],
 ) -> Optional[list[expr_internal.Expr]]:
-    """Helper function to convert an optional list to raw expressions."""
-    if isinstance(expr_list, Expr):
+    """Convert a sequence of expressions or column names to raw expressions."""
+    if isinstance(expr_list, (Expr, str)):
         expr_list = [expr_list]
-    return [e.expr for e in expr_list] if expr_list is not None else None
+    if expr_list is None:
+        return None
+    return [Expr.column(e).expr if isinstance(e, str) else e.expr for e in expr_list]
 
 
 def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr: