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

This reverts commit 5598f90.

Author: kosiew

python/datafusion/dataframe.py

@@ -45,7 +45,6 @@
     Expr,
     SortKey,
     ensure_expr,
-    ensure_expr_list,
     expr_list_to_raw_expr_list,
     sort_list_to_raw_sort_list,
 )
@@ -489,7 +488,17 @@ def with_columns(
         Returns:
             DataFrame with the new columns added.
         """
-        expressions = ensure_expr_list(exprs)
+
+        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)]
         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, Iterable, Optional, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence
 
 import pyarrow as pa
 
@@ -219,7 +219,6 @@
     "WindowFrame",
     "WindowFrameBound",
     "ensure_expr",
-    "ensure_expr_list",
 ]
 
 
@@ -244,31 +243,6 @@ 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,14 +436,6 @@ 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):