UNPICK

Author: kosiew

docs/source/user-guide/dataframe/index.rst

@@ -126,51 +126,6 @@ DataFusion's DataFrame API offers a wide range of operations:
     # Drop columns
     df = df.drop("temporary_column")
 
-String Columns and Expressions
-------------------------------
-
-Some ``DataFrame`` methods accept plain strings when an argument refers to an
-existing column. These include:
-
-* :py:meth:`~datafusion.DataFrame.select`
-* :py:meth:`~datafusion.DataFrame.sort`
-* :py:meth:`~datafusion.DataFrame.drop`
-* :py:meth:`~datafusion.DataFrame.join` (``on`` argument)
-* :py:meth:`~datafusion.DataFrame.aggregate` (grouping columns)
-
-For such methods, you can pass column names directly:
-
-.. code-block:: python
-
-    from datafusion import col, functions as f
-
-    df.sort('id')
-    df.aggregate('id', [f.count(col('value'))])
-
-The same operation can also be written with explicit column expressions, using either ``col()`` or ``column()``:
-
-.. code-block:: python
-
-    from datafusion import col, column, functions as f
-
-    df.sort(col('id'))
-    df.aggregate(column('id'), [f.count(col('value'))])
-
-Note that ``column()`` is an alias of ``col()``, so you can use either name; the example above shows both in action.
-
-Whenever an argument represents an expression—such as in
-:py:meth:`~datafusion.DataFrame.filter` or
-:py:meth:`~datafusion.DataFrame.with_column`—use ``col()`` to reference columns
-and wrap constant values with ``lit()`` (also available as ``literal()``):
-
-.. code-block:: python
-
-    from datafusion import col, lit
-    df.filter(col('age') > lit(21))
-
-Without ``lit()`` DataFusion would treat ``21`` as a column name rather than a
-constant value.
-
 Terminal Operations
 -------------------
 

python/datafusion/context.py

@@ -553,7 +553,7 @@ def register_listing_table(
         table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None,
         file_extension: str = ".parquet",
         schema: pa.Schema | None = None,
-        file_sort_order: list[list[Expr | SortExpr | str]] | None = None,
+        file_sort_order: list[list[Expr | SortExpr]] | None = None,
     ) -> None:
         """Register multiple files as a single table.
 
@@ -808,7 +808,7 @@ def register_parquet(
         file_extension: str = ".parquet",
         skip_metadata: bool = True,
         schema: pa.Schema | None = None,
-        file_sort_order: list[list[Expr | SortExpr | str]] | None = None,
+        file_sort_order: list[list[SortExpr]] | None = None,
     ) -> None:
         """Register a Parquet file as a table.
 
@@ -1099,7 +1099,7 @@ def read_parquet(
         file_extension: str = ".parquet",
         skip_metadata: bool = True,
         schema: pa.Schema | None = None,
-        file_sort_order: list[list[Expr | SortExpr | str]] | None = None,
+        file_sort_order: list[list[Expr | SortExpr]] | None = None,
     ) -> DataFrame:
         """Read a Parquet source into a :py:class:`~datafusion.dataframe.Dataframe`.
 

python/datafusion/dataframe.py

@@ -40,13 +40,7 @@
 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_TYPE_ERROR,
-    Expr,
-    SortExpr,
-    expr_list_to_raw_expr_list,
-    sort_list_to_raw_sort_list,
-)
+from datafusion.expr import Expr, SortExpr, sort_or_default
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
 
@@ -400,7 +394,9 @@ def select(self, *exprs: Expr | str) -> DataFrame:
             df = df.select("a", col("b"), col("a").alias("alternate_a"))
 
         """
-        exprs_internal = expr_list_to_raw_expr_list(exprs)
+        exprs_internal = [
+            Expr.column(arg).expr if isinstance(arg, str) else arg.expr for arg in exprs
+        ]
         return DataFrame(self.df.select(*exprs_internal))
 
     def drop(self, *columns: str) -> DataFrame:
@@ -430,8 +426,6 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
-            if not isinstance(p, Expr):
-                raise TypeError(EXPR_TYPE_ERROR)
             df = df.filter(p.expr)
         return DataFrame(df)
 
@@ -445,8 +439,6 @@ def with_column(self, name: str, expr: Expr) -> DataFrame:
         Returns:
             DataFrame with the new column.
         """
-        if not isinstance(expr, Expr):
-            raise TypeError(EXPR_TYPE_ERROR)
         return DataFrame(self.df.with_column(name, expr.expr))
 
     def with_columns(
@@ -476,22 +468,17 @@ def with_columns(
         def _simplify_expression(
             *exprs: Expr | Iterable[Expr], **named_exprs: Expr
         ) -> list[expr_internal.Expr]:
-            expr_list: list[expr_internal.Expr] = []
+            expr_list = []
             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)
-                ):
-                    raise TypeError(EXPR_TYPE_ERROR)
-                try:
-                    expr_list.extend(expr_list_to_raw_expr_list(expr))
-                except TypeError as err:
-                    raise TypeError(EXPR_TYPE_ERROR) from err
-            for alias, expr in named_exprs.items():
-                if not isinstance(expr, Expr):
-                    raise TypeError(EXPR_TYPE_ERROR)
-                expr_list.append(expr.alias(alias).expr)
+                if isinstance(expr, Expr):
+                    expr_list.append(expr.expr)
+                elif isinstance(expr, Iterable):
+                    expr_list.extend(inner_expr.expr for inner_expr in expr)
+                else:
+                    raise NotImplementedError
+            if named_exprs:
+                for alias, expr in named_exprs.items():
+                    expr_list.append(expr.alias(alias).expr)
             return expr_list
 
         expressions = _simplify_expression(*exprs, **named_exprs)
@@ -516,43 +503,37 @@ def with_column_renamed(self, old_name: str, new_name: str) -> DataFrame:
         return DataFrame(self.df.with_column_renamed(old_name, new_name))
 
     def aggregate(
-        self,
-        group_by: list[Expr | str] | Expr | str,
-        aggs: list[Expr] | Expr,
+        self, group_by: list[Expr] | Expr, aggs: list[Expr] | Expr
     ) -> DataFrame:
         """Aggregates the rows of the current DataFrame.
 
         Args:
-            group_by: List of expressions or column names to group by.
+            group_by: List of expressions to group by.
             aggs: List of expressions to aggregate.
 
         Returns:
             DataFrame after aggregation.
         """
-        group_by_list = group_by if isinstance(group_by, list) else [group_by]
-        aggs_list = aggs if isinstance(aggs, list) else [aggs]
+        group_by = group_by if isinstance(group_by, list) else [group_by]
+        aggs = aggs if isinstance(aggs, list) else [aggs]
 
-        group_by_exprs = expr_list_to_raw_expr_list(group_by_list)
-        aggs_exprs = []
-        for agg in aggs_list:
-            if not isinstance(agg, Expr):
-                raise TypeError(EXPR_TYPE_ERROR)
-            aggs_exprs.append(agg.expr)
-        return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
+        group_by = [e.expr for e in group_by]
+        aggs = [e.expr for e in aggs]
+        return DataFrame(self.df.aggregate(group_by, aggs))
 
-    def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
-        """Sort the DataFrame by the specified sorting expressions or column names.
+    def sort(self, *exprs: Expr | SortExpr) -> DataFrame:
+        """Sort the DataFrame by the specified sorting expressions.
 
         Note that any expression can be turned into a sort expression by
-        calling its ``sort`` method.
+        calling its` ``sort`` method.
 
         Args:
-            exprs: Sort expressions or column names, applied in order.
+            exprs: Sort expressions, applied in order.
 
         Returns:
             DataFrame after sorting.
         """
-        exprs_raw = sort_list_to_raw_sort_list(list(exprs))
+        exprs_raw = [sort_or_default(expr) for expr in exprs]
         return DataFrame(self.df.sort(*exprs_raw))
 
     def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame:
@@ -776,11 +757,7 @@ def join_on(
         Returns:
             DataFrame after join.
         """
-        exprs = []
-        for expr in on_exprs:
-            if not isinstance(expr, Expr):
-                raise TypeError(EXPR_TYPE_ERROR)
-            exprs.append(expr.expr)
+        exprs = [expr.expr for expr in on_exprs]
         return DataFrame(self.df.join_on(right.df, exprs, how))
 
     def explain(self, verbose: bool = False, analyze: bool = False) -> None:

python/datafusion/expr.py

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar, Optional
 
 import pyarrow as pa
 
@@ -39,10 +39,6 @@
 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.
@@ -220,26 +216,12 @@
 
 
 def expr_list_to_raw_expr_list(
-    expr_list: Optional[Sequence[Expr | str] | Expr | str],
+    expr_list: Optional[list[Expr] | Expr],
 ) -> Optional[list[expr_internal.Expr]]:
-    """Convert a sequence of expressions or column names to raw expressions."""
-    if isinstance(expr_list, (Expr, str)):
+    """Helper function to convert an optional list to raw expressions."""
+    if isinstance(expr_list, Expr):
         expr_list = [expr_list]
-    if expr_list is None:
-        return None
-    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__}. {EXPR_TYPE_ERROR}."
-            )
-            raise TypeError(error)
-    return raw_exprs
+    return [e.expr for e in expr_list] if expr_list is not None else None
 
 
 def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
@@ -250,27 +232,12 @@ def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
 
 
 def sort_list_to_raw_sort_list(
-    sort_list: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str],
+    sort_list: Optional[list[Expr | SortExpr] | Expr | SortExpr],
 ) -> Optional[list[expr_internal.SortExpr]]:
     """Helper function to return an optional sort list to raw variant."""
-    if isinstance(sort_list, (Expr, SortExpr, str)):
+    if isinstance(sort_list, (Expr, SortExpr)):
         sort_list = [sort_list]
-    if sort_list is None:
-        return None
-    raw_sort_list = []
-    for item in sort_list:
-        if isinstance(item, str):
-            expr_obj = Expr.column(item)
-        elif isinstance(item, (Expr, SortExpr)):
-            expr_obj = item
-        else:
-            error = (
-                "Expected Expr or column name, found:"
-                f" {type(item).__name__}. {EXPR_TYPE_ERROR}."
-            )
-            raise TypeError(error)
-        raw_sort_list.append(sort_or_default(expr_obj))
-    return raw_sort_list
+    return [sort_or_default(e) for e in sort_list] if sort_list is not None else None
 
 
 class Expr: