feat: add ensure_expr helper function for internal expression validation

Author: kosiew

python/datafusion/dataframe.py

@@ -41,10 +41,9 @@
 from datafusion._internal import ParquetColumnOptions as ParquetColumnOptionsInternal
 from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal
 from datafusion.expr import (
-    EXPR_TYPE_ERROR,
     Expr,
     SortKey,
-    _to_raw_expr,
+    ensure_expr,
     expr_list_to_raw_expr_list,
     sort_list_to_raw_sort_list,
 )
@@ -59,8 +58,6 @@
     import polars as pl
     import pyarrow as pa
 
-    from datafusion._internal import expr as expr_internal
-
 from enum import Enum
 
 
@@ -293,27 +290,6 @@ def __init__(
         self.bloom_filter_ndv = bloom_filter_ndv
 
 
-def _ensure_expr(value: Expr | str) -> expr_internal.Expr:
-    """Return the internal expression from ``Expr`` or raise ``TypeError``.
-
-    This helper rejects plain strings so higher level APIs consistently
-    require explicit :func:`~datafusion.col` or :func:`~datafusion.lit`
-    expressions.
-
-    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 _to_raw_expr(value)
-
-
 class DataFrame:
     """Two dimensional table representation of data.
 
@@ -460,7 +436,7 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
-            df = df.filter(_ensure_expr(p))
+            df = df.filter(ensure_expr(p))
         return DataFrame(df)
 
     def with_column(self, name: str, expr: Expr) -> DataFrame:
@@ -482,7 +458,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, ensure_expr(expr)))
 
     def with_columns(
         self, *exprs: Expr | Iterable[Expr], **named_exprs: Expr
@@ -512,30 +488,19 @@ def with_columns(
             DataFrame with the new columns added.
         """
 
-        def _simplify_expression(
-            *exprs: Expr | Iterable[Expr], **named_exprs: Expr
-        ) -> list[expr_internal.Expr]:
-            expr_list: list[expr_internal.Expr] = []
-            for expr in exprs:
+        def _iter_exprs(items: Iterable[Expr | Iterable[Expr]]) -> Iterable[Expr | str]:
+            for expr in items:
                 if isinstance(expr, str):
-                    raise TypeError(EXPR_TYPE_ERROR)
-                if isinstance(expr, Iterable) and not isinstance(expr, Expr):
-                    expr_value = list(expr)
-                    if any(isinstance(inner, str) for inner in expr_value):
-                        raise TypeError(EXPR_TYPE_ERROR)
+                    yield expr
+                elif isinstance(expr, Iterable) and not isinstance(expr, Expr):
+                    yield from _iter_exprs(expr)
                 else:
-                    expr_value = expr
-                try:
-                    expr_list.extend(expr_list_to_raw_expr_list(expr_value))
-                except TypeError as err:
-                    raise TypeError(EXPR_TYPE_ERROR) from err
-            for alias, expr in named_exprs.items():
-                if not isinstance(expr, Expr):
-                    raise TypeError(EXPR_TYPE_ERROR)
-                expr_list.append(expr.alias(alias).expr)
-            return expr_list
-
-        expressions = _simplify_expression(*exprs, **named_exprs)
+                    yield expr
+
+        expressions = [ensure_expr(e) for e in _iter_exprs(exprs)]
+        for alias, expr in named_exprs.items():
+            ensure_expr(expr)
+            expressions.append(expr.alias(alias).expr)
 
         return DataFrame(self.df.with_columns(expressions))
 
@@ -574,11 +539,7 @@ def aggregate(
         aggs_list = aggs if isinstance(aggs, list) else [aggs]
 
         group_by_exprs = expr_list_to_raw_expr_list(group_by_list)
-        aggs_exprs = []
-        for agg in aggs_list:
-            if not isinstance(agg, Expr):
-                raise TypeError(EXPR_TYPE_ERROR)
-            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: SortKey) -> DataFrame:
@@ -824,7 +785,7 @@ def join_on(
         Returns:
             DataFrame after join.
         """
-        exprs = [_ensure_expr(expr) for expr in on_exprs]
+        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

@@ -218,9 +218,31 @@
     "WindowExpr",
     "WindowFrame",
     "WindowFrameBound",
+    "ensure_expr",
 ]
 
 
+def ensure_expr(value: Expr) -> expr_internal.Expr:
+    """Return the internal expression from ``Expr`` or raise ``TypeError``.
+
+    This helper rejects plain strings so higher level APIs consistently
+    require explicit :func:`~datafusion.col` or :func:`~datafusion.lit`
+    expressions.
+
+    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
+
+
 def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
     """Convert a Python expression or column name to its raw variant.
 

python/tests/test_expr.py

@@ -47,6 +47,7 @@
     TransactionEnd,
     TransactionStart,
     Values,
+    ensure_expr,
 )
 
 
@@ -880,3 +881,10 @@ def test_literal_metadata(ctx):
     for expected_field in expected_schema:
         actual_field = result[0].schema.field(expected_field.name)
         assert expected_field.metadata == actual_field.metadata
+
+
+def test_ensure_expr():
+    e = col("a")
+    assert ensure_expr(e) is e.expr
+    with pytest.raises(TypeError, match=r"Use col\(\) or lit\(\)"):
+        ensure_expr("a")