refactor: add _ensure_expr function for improved expression validation in DataFrame methods

Author: kosiew

python/datafusion/dataframe.py

@@ -292,6 +292,23 @@ def __init__(
         self.bloom_filter_ndv = bloom_filter_ndv
 
 
+def _ensure_expr(value: Expr) -> expr_internal.Expr:
+    """Return the internal expression or raise ``TypeError`` if invalid.
+
+    Args:
+        value: Candidate expression.
+
+    Returns:
+        The internal expression representation.
+
+    Raises:
+        TypeError: If ``value`` is not an instance of :class:`Expr`.
+    """
+    if not isinstance(value, Expr):
+        raise TypeError(EXPR_TYPE_ERROR)
+    return value.expr
+
+
 class DataFrame:
     """Two dimensional table representation of data.
 
@@ -430,9 +447,7 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
-            if not isinstance(p, Expr):
-                raise TypeError(EXPR_TYPE_ERROR)
-            df = df.filter(p.expr)
+            df = df.filter(_ensure_expr(p))
         return DataFrame(df)
 
     def with_column(self, name: str, expr: Expr) -> DataFrame:
@@ -445,9 +460,7 @@ def with_column(self, name: str, expr: Expr) -> DataFrame:
         Returns:
             DataFrame with the new column.
         """
-        if not isinstance(expr, Expr):
-            raise TypeError(EXPR_TYPE_ERROR)
-        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
@@ -778,11 +791,7 @@ def join_on(
         Returns:
             DataFrame after join.
         """
-        exprs = []
-        for expr in on_exprs:
-            if not isinstance(expr, Expr):
-                raise TypeError(EXPR_TYPE_ERROR)
-            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: