Add type checks for expressions in DataFrame methods and expr_list_to_raw_expr_list for improved error handling

Author: kosiew

python/datafusion/dataframe.py

@@ -424,6 +424,9 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
+            if isinstance(p, str) or not isinstance(p, Expr):
+                error = "Use col() or lit() to construct expressions"
+                raise TypeError(error)
             df = df.filter(expr_list_to_raw_expr_list(p)[0])
         return DataFrame(df)
 
@@ -437,7 +440,10 @@ def with_column(self, name: str, expr: Expr) -> DataFrame:
         Returns:
             DataFrame with the new column.
         """
-        return DataFrame(self.df.with_column(name, expr_list_to_raw_expr_list(expr)[0]))
+        if not isinstance(expr, Expr):
+            error = "Use col() or lit() to construct expressions"
+            raise TypeError(error)
+        return DataFrame(self.df.with_column(name, expr.expr))
 
     def with_columns(
         self, *exprs: Expr | Iterable[Expr], **named_exprs: Expr
@@ -468,13 +474,22 @@ def _simplify_expression(
         ) -> list[expr_internal.Expr]:
             expr_list = []
             for expr in exprs:
-                if isinstance(expr, Iterable) and not isinstance(expr, Expr):
-                    expr_list.extend(expr_list_to_raw_expr_list(inner_expr)[0] for inner_expr in expr)
+                if isinstance(expr, Expr):
+                    expr_list.append(expr.expr)
+                elif isinstance(expr, Iterable) and not isinstance(expr, (str, Expr)):
+                    for inner_expr in expr:
+                        if not isinstance(inner_expr, Expr):
+                            error = "Use col() or lit() to construct expressions"
+                            raise TypeError(error)
+                        expr_list.append(inner_expr.expr)
                 else:
-                    expr_list.append(expr_list_to_raw_expr_list(expr)[0])
+                    error = "Use col() or lit() to construct expressions"
+                    raise TypeError(error)
             if named_exprs:
                 for alias, expr in named_exprs.items():
-                    expr_list_to_raw_expr_list(expr)[0]
+                    if not isinstance(expr, Expr):
+                        error = "Use col() or lit() to construct expressions"
+                        raise TypeError(error)
                     expr_list.append(expr.alias(alias).expr)
             return expr_list
 
@@ -516,10 +531,24 @@ def aggregate(
         group_by_list = group_by if isinstance(group_by, list) else [group_by]
         aggs_list = aggs if isinstance(aggs, list) else [aggs]
 
-        group_by_exprs = [
-            Expr.column(e).expr if isinstance(e, str) else e.expr for e in group_by_list
-        ]
-        aggs_exprs = [expr_list_to_raw_expr_list(agg)[0] for agg in aggs_list]
+        group_by_exprs = []
+        for e in group_by_list:
+            if isinstance(e, str):
+                group_by_exprs.append(Expr.column(e).expr)
+            elif isinstance(e, Expr):
+                group_by_exprs.append(e.expr)
+            else:
+                error = (
+                    "Expected Expr or column name, found:"
+                    f" {type(e).__name__}. Use col() or lit() to construct expressions"
+                )
+                raise TypeError(error)
+        aggs_exprs = []
+        for agg in aggs_list:
+            if not isinstance(agg, Expr):
+                error = "Use col() or lit() to construct expressions"
+                raise TypeError(error)
+            aggs_exprs.append(agg.expr)
         return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
     def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
@@ -765,7 +794,12 @@ def join_on(
         Returns:
             DataFrame after join.
         """
-        exprs = [expr_list_to_raw_expr_list(expr)[0] for expr in on_exprs]
+        exprs = []
+        for expr in on_exprs:
+            if not isinstance(expr, Expr):
+                error = "Use col() or lit() to construct expressions"
+                raise TypeError(error)
+            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/datafusion/expr.py

@@ -223,7 +223,19 @@ def expr_list_to_raw_expr_list(
         expr_list = [expr_list]
     if expr_list is None:
         return None
-    return [Expr.column(e).expr if isinstance(e, str) else e.expr for e in expr_list]
+    raw_exprs: list[expr_internal.Expr] = []
+    for e in expr_list:
+        if isinstance(e, str):
+            raw_exprs.append(Expr.column(e).expr)
+        elif isinstance(e, Expr):
+            raw_exprs.append(e.expr)
+        else:
+            error = (
+                "Expected Expr or column name, found:"
+                f" {type(e).__name__}. Use col() or lit() to construct expressions."
+            )
+            raise TypeError(error)
+    return raw_exprs
 
 
 def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr: