Add type checks for expressions in DataFrame methods to enforce correct usage

Author: kosiew

python/datafusion/dataframe.py

@@ -445,6 +445,12 @@ def with_column(self, name: str, expr: Expr) -> DataFrame:
         Returns:
             DataFrame with the new column.
         """
+        if not isinstance(expr, Expr):
+            msg = (
+                f"Expected Expr, got {type(expr).__name__}. "
+                "Use col() or lit() to construct expressions."
+            )
+            raise TypeError(msg)
         return DataFrame(self.df.with_column(name, expr.expr))
 
     def with_columns(
@@ -479,11 +485,28 @@ def _simplify_expression(
                 if isinstance(expr, Expr):
                     expr_list.append(expr.expr)
                 elif isinstance(expr, Iterable):
-                    expr_list.extend(inner_expr.expr for inner_expr in expr)
+                    for inner_expr in expr:
+                        if not isinstance(inner_expr, Expr):
+                            msg = (
+                                f"Expected Expr, got {type(inner_expr).__name__}. "
+                                "Use col() or lit() to construct expressions."
+                            )
+                            raise TypeError(msg)
+                        expr_list.append(inner_expr.expr)
                 else:
-                    raise NotImplementedError
+                    msg = (
+                        f"Expected Expr, got {type(expr).__name__}. "
+                        "Use col() or lit() to construct expressions."
+                    )
+                    raise TypeError(msg)
             if named_exprs:
                 for alias, expr in named_exprs.items():
+                    if not isinstance(expr, Expr):
+                        msg = (
+                            f"Expected Expr, got {type(expr).__name__}. "
+                            "Use col() or lit() to construct expressions."
+                        )
+                        raise TypeError(msg)
                     expr_list.append(expr.alias(alias).expr)
             return expr_list
 
@@ -528,7 +551,16 @@ def aggregate(
         group_by_exprs = [
             Expr.column(e).expr if isinstance(e, str) else e.expr for e in group_by_list
         ]
-        aggs_exprs = [e.expr for e in aggs_list]
+
+        aggs_exprs: list[expr_internal.Expr] = []
+        for agg in aggs_list:
+            if not isinstance(agg, Expr):
+                msg = (
+                    f"Expected Expr, got {type(agg).__name__}. "
+                    "Use col() or lit() to construct expressions."
+                )
+                raise TypeError(msg)
+            aggs_exprs.append(agg.expr)
         return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
     def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
@@ -770,7 +802,15 @@ def join_on(
         Returns:
             DataFrame after join.
         """
-        exprs = [expr.expr for expr in on_exprs]
+        exprs = []
+        for expr in on_exprs:
+            if not isinstance(expr, Expr):
+                msg = (
+                    f"Expected Expr, got {type(expr).__name__}. "
+                    "Use col() or lit() to construct expressions."
+                )
+                raise TypeError(msg)
+            exprs.append(expr.expr)
         return DataFrame(self.df.join_on(right.df, exprs, how))
 
     def explain(self, verbose: bool = False, analyze: bool = False) -> None:

python/tests/test_dataframe.py

@@ -350,6 +350,11 @@ def test_with_column(df):
     assert result.column(2) == pa.array([5, 7, 9])
 
 
+def test_with_column_invalid_expr(df):
+    with pytest.raises(TypeError, match=r"Use col\(\) or lit\(\)"):
+        df.with_column("c", "a")
+
+
 def test_with_columns(df):
     df = df.with_columns(
         (column("a") + column("b")).alias("c"),
@@ -381,6 +386,13 @@ def test_with_columns(df):
     assert result.column(6) == pa.array([5, 7, 9])
 
 
+def test_with_columns_invalid_expr(df):
+    with pytest.raises(TypeError, match=r"Use col\(\) or lit\(\)"):
+        df.with_columns("a")
+    with pytest.raises(TypeError, match=r"Use col\(\) or lit\(\)"):
+        df.with_columns(c="a")
+
+
 def test_cast(df):
     df = df.cast({"a": pa.float16(), "b": pa.list_(pa.uint32())})
     expected = pa.schema(
@@ -539,6 +551,25 @@ def test_join_on():
     assert table.to_pydict() == expected
 
 
+def test_join_on_invalid_expr():
+    ctx = SessionContext()
+
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2]), pa.array([4, 5])],
+        names=["a", "b"],
+    )
+    df = ctx.create_dataframe([[batch]], "l")
+    df1 = ctx.create_dataframe([[batch]], "r")
+
+    with pytest.raises(TypeError, match=r"Use col\(\) or lit\(\)"):
+        df.join_on(df1, "a")
+
+
+def test_aggregate_invalid_aggs(df):
+    with pytest.raises(TypeError, match=r"Use col\(\) or lit\(\)"):
+        df.aggregate([], "a")
+
+
 def test_distinct():
     ctx = SessionContext()