refactor: improve type checking in DataFrame expression handling

Author: kosiew

python/datafusion/dataframe.py

@@ -478,14 +478,16 @@ def _simplify_expression(
         ) -> list[expr_internal.Expr]:
             expr_list: list[expr_internal.Expr] = []
             for expr in exprs:
-                if isinstance(expr, str) or (
-                    isinstance(expr, Iterable)
-                    and not isinstance(expr, Expr)
-                    and any(isinstance(inner, str) for inner in expr)
-                ):
+                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)
+                else:
+                    expr_value = expr
                 try:
-                    expr_list.extend(expr_list_to_raw_expr_list(expr))
+                    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():

python/tests/test_with_columns_generator.py

@@ -0,0 +1,23 @@
+import pyarrow as pa
+from datafusion import SessionContext, column
+
+
+def test_with_columns_generator():
+    ctx = SessionContext()
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, 6]), pa.array([8, 5, 8])],
+        names=["a", "b", "c"],
+    )
+    df = ctx.from_arrow(batch)
+
+    def gen():
+        for name in ["d", "e"]:
+            yield (column("a") + column("b")).alias(name)
+
+    df = df.with_columns(gen())
+    result = df.collect()[0]
+
+    assert result.schema.names == ["a", "b", "c", "d", "e"]
+    expected = pa.array([5, 7, 9])
+    assert result.column(3) == expected
+    assert result.column(4) == expected