Reapply "feat: add ensure_expr_list function to flatten and validate nested expressions"

This reverts commit 5964e47.

Author: kosiew

python/datafusion/dataframe.py

@@ -45,6 +45,7 @@
     Expr,
     SortKey,
     ensure_expr,
+    ensure_expr_list,
     expr_list_to_raw_expr_list,
     sort_list_to_raw_sort_list,
 )
@@ -488,17 +489,7 @@ def with_columns(
         Returns:
             DataFrame with the new columns added.
         """
-
-        def _iter_exprs(items: Iterable[Expr | Iterable[Expr]]) -> Iterable[Expr | str]:
-            for expr in items:
-                if isinstance(expr, str):
-                    yield expr
-                elif isinstance(expr, Iterable) and not isinstance(expr, Expr):
-                    yield from _iter_exprs(expr)
-                else:
-                    yield expr
-
-        expressions = [ensure_expr(e) for e in _iter_exprs(exprs)]
+        expressions = ensure_expr_list(exprs)
         for alias, expr in named_exprs.items():
             ensure_expr(expr)
             expressions.append(expr.alias(alias).expr)

python/datafusion/expr.py

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Optional, Sequence
 
 import pyarrow as pa
 
@@ -219,6 +219,7 @@
     "WindowFrame",
     "WindowFrameBound",
     "ensure_expr",
+    "ensure_expr_list",
 ]
 
 
@@ -243,6 +244,31 @@ def ensure_expr(value: Expr | Any) -> expr_internal.Expr:
     return value.expr
 
 
+def ensure_expr_list(
+    exprs: Iterable[Expr | Iterable[Expr]],
+) -> list[expr_internal.Expr]:
+    """Flatten an iterable of expressions, validating each via ``ensure_expr``.
+
+    Args:
+        exprs: Possibly nested iterable containing expressions.
+
+    Returns:
+        A flat list of raw expressions.
+
+    Raises:
+        TypeError: If any item is not an instance of :class:`Expr`.
+    """
+
+    def _iter(items: Iterable[Expr | Iterable[Expr]]) -> Iterable[expr_internal.Expr]:
+        for expr in items:
+            if isinstance(expr, Iterable) and not isinstance(expr, Expr):
+                yield from _iter(expr)
+            else:
+                yield ensure_expr(expr)
+
+    return list(_iter(exprs))
+
+
 def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
     """Convert a Python expression or column name to its raw variant.
 

python/tests/test_dataframe.py

@@ -436,6 +436,14 @@ def test_with_columns_invalid_expr(df):
         TypeError, match=r"Use col\(\)/column\(\) or lit\(\)/literal\(\)"
     ):
         df.with_columns(c="a")
+    with pytest.raises(
+        TypeError, match=r"Use col\(\)/column\(\) or lit\(\)/literal\(\)"
+    ):
+        df.with_columns(["a"])
+    with pytest.raises(
+        TypeError, match=r"Use col\(\)/column\(\) or lit\(\)/literal\(\)"
+    ):
+        df.with_columns(c=["a"])
 
 
 def test_cast(df):