Refactor expression type error handling in DataFrame and expr modules for improved clarity and consistency

Author: kosiew

python/datafusion/dataframe.py

@@ -40,7 +40,13 @@
 from datafusion._internal import DataFrame as DataFrameInternal
 from datafusion._internal import ParquetColumnOptions as ParquetColumnOptionsInternal
 from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal
-from datafusion.expr import Expr, SortExpr, expr_list_to_raw_expr_list, sort_or_default
+from datafusion.expr import (
+    _EXPR_TYPE_ERROR,
+    Expr,
+    SortExpr,
+    expr_list_to_raw_expr_list,
+    sort_or_default,
+)
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
 
@@ -424,10 +430,9 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
-            if not isinstance(p, Expr):
-                error = "Use col() or lit() to construct expressions"
-                raise TypeError(error)
-            df = df.filter(p.expr)
+            if isinstance(p, str) or not isinstance(p, Expr):
+                raise TypeError(_EXPR_TYPE_ERROR)
+            df = df.filter(expr_list_to_raw_expr_list(p)[0])
         return DataFrame(df)
 
     def with_column(self, name: str, expr: Expr) -> DataFrame:
@@ -441,8 +446,7 @@ def with_column(self, name: str, expr: Expr) -> DataFrame:
             DataFrame with the new column.
         """
         if not isinstance(expr, Expr):
-            error = "Use col() or lit() to construct expressions"
-            raise TypeError(error)
+            raise TypeError(_EXPR_TYPE_ERROR)
         return DataFrame(self.df.with_column(name, expr.expr))
 
     def with_columns(
@@ -479,17 +483,14 @@ def _simplify_expression(
                 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)
+                            raise TypeError(_EXPR_TYPE_ERROR)
                         expr_list.append(inner_expr.expr)
                 else:
-                    error = "Use col() or lit() to construct expressions"
-                    raise TypeError(error)
+                    raise TypeError(_EXPR_TYPE_ERROR)
             if named_exprs:
                 for alias, expr in named_exprs.items():
                     if not isinstance(expr, Expr):
-                        error = "Use col() or lit() to construct expressions"
-                        raise TypeError(error)
+                        raise TypeError(_EXPR_TYPE_ERROR)
                     expr_list.append(expr.alias(alias).expr)
             return expr_list
 
@@ -535,8 +536,7 @@ def aggregate(
         aggs_exprs = []
         for agg in aggs_list:
             if not isinstance(agg, Expr):
-                error = "Use col() or lit() to construct expressions"
-                raise TypeError(error)
+                raise TypeError(_EXPR_TYPE_ERROR)
             aggs_exprs.append(agg.expr)
         return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
@@ -786,8 +786,7 @@ def join_on(
         exprs = []
         for expr in on_exprs:
             if not isinstance(expr, Expr):
-                error = "Use col() or lit() to construct expressions"
-                raise TypeError(error)
+                raise TypeError(_EXPR_TYPE_ERROR)
             exprs.append(expr.expr)
         return DataFrame(self.df.join_on(right.df, exprs, how))
 

python/datafusion/expr.py

@@ -39,6 +39,10 @@
 if TYPE_CHECKING:
     from datafusion.plan import LogicalPlan
 
+
+# Standard error message for invalid expression types
+_EXPR_TYPE_ERROR = "Use col() or lit() to construct expressions"
+
 # The following are imported from the internal representation. We may choose to
 # give these all proper wrappers, or to simply leave as is. These were added
 # in order to support passing the `test_imports` unit test.
@@ -232,7 +236,7 @@ def expr_list_to_raw_expr_list(
         else:
             error = (
                 "Expected Expr or column name, found:"
-                f" {type(e).__name__}. Use col() or lit() to construct expressions."
+                f" {type(e).__name__}. {_EXPR_TYPE_ERROR}."
             )
             raise TypeError(error)
     return raw_exprs