Add _ensure_expr function to validate and convert expressions in DataFrame methods

Author: kosiew

python/datafusion/dataframe.py

@@ -40,7 +40,13 @@
 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, _to_expr_list, sort_or_default
+from datafusion.expr import (
+    Expr,
+    SortExpr,
+    _ensure_expr,
+    _to_expr_list,
+    sort_or_default,
+)
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
 
@@ -424,13 +430,7 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
-            if not isinstance(p, Expr):
-                msg = (
-                    f"Expected Expr, got {type(p).__name__}. "
-                    "Use col() or lit() to construct expressions."
-                )
-                raise TypeError(msg)
-            df = df.filter(p.expr)
+            df = df.filter(_ensure_expr(p))
         return DataFrame(df)
 
     def with_column(self, name: str, expr: Expr) -> DataFrame:
@@ -443,13 +443,7 @@ def with_column(self, name: str, expr: Expr) -> DataFrame:
         Returns:
             DataFrame with the new column.
         """
-        if not isinstance(expr, Expr):
-            msg = (
-                f"Expected Expr, got {type(expr).__name__}. "
-                "Use col() or lit() to construct expressions."
-            )
-            raise TypeError(msg)
-        return DataFrame(self.df.with_column(name, expr.expr))
+        return DataFrame(self.df.with_column(name, _ensure_expr(expr)))
 
     def with_columns(
         self, *exprs: Expr | Iterable[Expr], **named_exprs: Expr
@@ -480,31 +474,13 @@ def _simplify_expression(
         ) -> list[expr_internal.Expr]:
             expr_list = []
             for expr in exprs:
-                if isinstance(expr, Expr):
-                    expr_list.append(expr.expr)
-                elif isinstance(expr, Iterable):
-                    for inner_expr in expr:
-                        if not isinstance(inner_expr, Expr):
-                            msg = (
-                                f"Expected Expr, got {type(inner_expr).__name__}. "
-                                "Use col() or lit() to construct expressions."
-                            )
-                            raise TypeError(msg)
-                        expr_list.append(inner_expr.expr)
+                if isinstance(expr, Iterable) and not isinstance(expr, Expr):
+                    expr_list.extend(_ensure_expr(inner_expr) for inner_expr in expr)
                 else:
-                    msg = (
-                        f"Expected Expr, got {type(expr).__name__}. "
-                        "Use col() or lit() to construct expressions."
-                    )
-                    raise TypeError(msg)
+                    expr_list.append(_ensure_expr(expr))
             if named_exprs:
                 for alias, expr in named_exprs.items():
-                    if not isinstance(expr, Expr):
-                        msg = (
-                            f"Expected Expr, got {type(expr).__name__}. "
-                            "Use col() or lit() to construct expressions."
-                        )
-                        raise TypeError(msg)
+                    _ensure_expr(expr)
                     expr_list.append(expr.alias(alias).expr)
             return expr_list
 
@@ -549,15 +525,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 = []
-        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)
+        aggs_exprs = [_ensure_expr(agg) for agg in aggs_list]
         return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
     def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
@@ -803,15 +771,7 @@ def join_on(
         Returns:
             DataFrame after join.
         """
-        exprs = []
-        for expr in on_exprs:
-            if not isinstance(expr, Expr):
-                msg = (
-                    f"Expected Expr, got {type(expr).__name__}. "
-                    "Use col() or lit() to construct expressions."
-                )
-                raise TypeError(msg)
-            exprs.append(expr.expr)
+        exprs = [_ensure_expr(expr) 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

@@ -220,6 +220,17 @@ def _to_expr_list(exprs: Sequence[Expr | str]) -> list[expr_internal.Expr]:
     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],
 ) -> Optional[list[expr_internal.Expr]]: