refactor: enhance expression handling by introducing _to_raw_expr function and updating _ensure_expr

Author: kosiew

python/datafusion/dataframe.py

@@ -44,6 +44,7 @@
     EXPR_TYPE_ERROR,
     Expr,
     SortKey,
+    _to_raw_expr,
     expr_list_to_raw_expr_list,
     sort_list_to_raw_sort_list,
 )
@@ -292,21 +293,19 @@ def __init__(
         self.bloom_filter_ndv = bloom_filter_ndv
 
 
-def _ensure_expr(value: Expr) -> expr_internal.Expr:
+def _ensure_expr(value: Expr | str) -> expr_internal.Expr:
     """Return the internal expression or raise ``TypeError`` if invalid.
 
     Args:
-        value: Candidate expression.
+        value: Candidate expression or column name.
 
     Returns:
         The internal expression representation.
 
     Raises:
-        TypeError: If ``value`` is not an instance of :class:`Expr`.
+        TypeError: If ``value`` is not an instance of :class:`Expr` or ``str``.
     """
-    if not isinstance(value, Expr):
-        raise TypeError(EXPR_TYPE_ERROR)
-    return value.expr
+    return _to_raw_expr(value)
 
 
 class DataFrame:

python/datafusion/expr.py

@@ -222,6 +222,29 @@
 ]
 
 
+def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
+    """Convert a Python expression or column name to its raw variant.
+
+    Args:
+        value: Candidate expression or column name.
+
+    Returns:
+        The internal :class:`~datafusion._internal.expr.Expr` representation.
+
+    Raises:
+        TypeError: If ``value`` is neither an :class:`Expr` nor ``str``.
+    """
+    if isinstance(value, str):
+        return Expr.column(value).expr
+    if isinstance(value, Expr):
+        return value.expr
+    error = (
+        "Expected Expr or column name, found:"
+        f" {type(value).__name__}. {EXPR_TYPE_ERROR}."
+    )
+    raise TypeError(error)
+
+
 def expr_list_to_raw_expr_list(
     expr_list: Optional[Sequence[Expr | str] | Expr | str],
 ) -> Optional[list[expr_internal.Expr]]:
@@ -230,23 +253,11 @@ def expr_list_to_raw_expr_list(
         expr_list = [expr_list]
     if expr_list is None:
         return None
-    raw_exprs: list[expr_internal.Expr] = []
-    for e in expr_list:
-        if isinstance(e, str):
-            raw_exprs.append(Expr.column(e).expr)
-        elif isinstance(e, Expr):
-            raw_exprs.append(e.expr)
-        else:
-            error = (
-                "Expected Expr or column name, found:"
-                f" {type(e).__name__}. {EXPR_TYPE_ERROR}."
-            )
-            raise TypeError(error)
-    return raw_exprs
+    return [_to_raw_expr(e) for e in expr_list]
 
 
 def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
-    """Helper function to return a default Sort if an Expr is provided."""
+    """Return a :class:`SortExpr`, defaulting attributes when necessary."""
     if isinstance(e, SortExpr):
         return e.raw_sort
     return SortExpr(e, ascending=True, nulls_first=True).raw_sort
@@ -262,17 +273,11 @@ def sort_list_to_raw_sort_list(
         return None
     raw_sort_list = []
     for item in sort_list:
-        if isinstance(item, str):
-            expr_obj = Expr.column(item)
-        elif isinstance(item, (Expr, SortExpr)):
-            expr_obj = item
+        if isinstance(item, SortExpr):
+            raw_sort_list.append(sort_or_default(item))
         else:
-            error = (
-                "Expected Expr or column name, found:"
-                f" {type(item).__name__}. {EXPR_TYPE_ERROR}."
-            )
-            raise TypeError(error)
-        raw_sort_list.append(sort_or_default(expr_obj))
+            raw_expr = _to_raw_expr(item)  # may raise ``TypeError``
+            raw_sort_list.append(sort_or_default(Expr(raw_expr)))
     return raw_sort_list