Revert "UNPICK"

This reverts commit ca23c74.

Author: kosiew

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

@@ -126,6 +126,49 @@ 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 an explicit column expression:
+
+.. code-block:: python
+
+    from datafusion import col, functions as f
+
+    df.sort(col('id'))
+    df.aggregate(col('id'), [f.count(col('value'))])
+
+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/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, 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
 
@@ -394,9 +400,7 @@ def select(self, *exprs: Expr | str) -> DataFrame:
             df = df.select("a", col("b"), col("a").alias("alternate_a"))
 
         """
-        exprs_internal = [
-            Expr.column(arg).expr if isinstance(arg, str) else arg.expr for arg in exprs
-        ]
+        exprs_internal = expr_list_to_raw_expr_list(exprs)
         return DataFrame(self.df.select(*exprs_internal))
 
     def drop(self, *columns: str) -> DataFrame:
@@ -426,7 +430,9 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
-            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:
@@ -439,6 +445,8 @@ 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(
@@ -472,12 +480,17 @@ def _simplify_expression(
             for expr in exprs:
                 if isinstance(expr, Expr):
                     expr_list.append(expr.expr)
-                elif isinstance(expr, Iterable):
-                    expr_list.extend(inner_expr.expr for inner_expr in expr)
+                elif isinstance(expr, Iterable) and not isinstance(expr, (str, Expr)):
+                    for inner_expr in expr:
+                        if not isinstance(inner_expr, Expr):
+                            raise TypeError(_EXPR_TYPE_ERROR)
+                        expr_list.append(inner_expr.expr)
                 else:
-                    raise NotImplementedError
+                    raise TypeError(_EXPR_TYPE_ERROR)
             if named_exprs:
                 for alias, expr in named_exprs.items():
+                    if not isinstance(expr, Expr):
+                        raise TypeError(_EXPR_TYPE_ERROR)
                     expr_list.append(expr.alias(alias).expr)
             return expr_list
 
@@ -503,37 +516,56 @@ 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] | Expr, aggs: list[Expr] | Expr
+        self,
+        group_by: list[Expr | str] | Expr | str,
+        aggs: list[Expr] | Expr,
     ) -> DataFrame:
         """Aggregates the rows of the current DataFrame.
 
         Args:
-            group_by: List of expressions to group by.
+            group_by: List of expressions or column names to group by.
             aggs: List of expressions to aggregate.
 
         Returns:
             DataFrame after aggregation.
         """
-        group_by = group_by if isinstance(group_by, list) else [group_by]
-        aggs = aggs if isinstance(aggs, list) else [aggs]
+        group_by_list = group_by if isinstance(group_by, list) else [group_by]
+        aggs_list = aggs if isinstance(aggs, list) else [aggs]
 
-        group_by = [e.expr for e in group_by]
-        aggs = [e.expr for e in aggs]
-        return DataFrame(self.df.aggregate(group_by, 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))
 
-    def sort(self, *exprs: Expr | SortExpr) -> DataFrame:
-        """Sort the DataFrame by the specified sorting expressions.
+    def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
+        """Sort the DataFrame by the specified sorting expressions or column names.
 
         Note that any expression can be turned into a sort expression by
-        calling its` ``sort`` method.
+        calling its ``sort`` method.
 
         Args:
-            exprs: Sort expressions, applied in order.
+            exprs: Sort expressions or column names, applied in order.
 
         Returns:
             DataFrame after sorting.
         """
-        exprs_raw = [sort_or_default(expr) for expr in exprs]
+        exprs_raw = []
+        for e in exprs:
+            if isinstance(e, SortExpr):
+                exprs_raw.append(sort_or_default(e))
+            elif isinstance(e, str):
+                exprs_raw.append(sort_or_default(Expr.column(e)))
+            elif isinstance(e, Expr):
+                exprs_raw.append(sort_or_default(e))
+            else:
+                error = (
+                    "Expected Expr or column name, found:"
+                    f" {type(e).__name__}. Use col() or lit() to construct expressions."
+                )
+                raise TypeError(error)
         return DataFrame(self.df.sort(*exprs_raw))
 
     def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame:
@@ -757,7 +789,11 @@ 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):
+                raise TypeError(_EXPR_TYPE_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

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence
 
 import pyarrow as pa
 
@@ -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.
@@ -216,12 +220,26 @@
 
 
 def expr_list_to_raw_expr_list(
-    expr_list: Optional[list[Expr] | Expr],
+    expr_list: Optional[Sequence[Expr | str] | Expr | str],
 ) -> Optional[list[expr_internal.Expr]]:
-    """Helper function to convert an optional list to raw expressions."""
-    if isinstance(expr_list, Expr):
+    """Convert a sequence of expressions or column names to raw expressions."""
+    if isinstance(expr_list, (Expr, str)):
         expr_list = [expr_list]
-    return [e.expr for e in expr_list] if expr_list is not None else None
+    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
 
 
 def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:

python/tests/test_dataframe.py

@@ -33,6 +33,7 @@
     WindowFrame,
     column,
     literal,
+    col,
 )
 from datafusion import (
     functions as f,
@@ -227,6 +228,13 @@ def test_select_mixed_expr_string(df):
     assert result.column(1) == pa.array([1, 2, 3])
 
 
+def test_select_unsupported(df):
+    with pytest.raises(
+        TypeError, match=r"Expected Expr or column name.*col\(\) or lit\(\)"
+    ):
+        df.select(1)
+
+
 def test_filter(df):
     df1 = df.filter(column("a") > literal(2)).select(
         column("a") + column("b"),
@@ -268,6 +276,32 @@ def test_sort(df):
     assert table.to_pydict() == expected
 
 
+def test_sort_string_and_expression_equivalent(df):
+    from datafusion import col
+
+    result_str = df.sort("a").to_pydict()
+    result_expr = df.sort(col("a")).to_pydict()
+    assert result_str == result_expr
+
+
+def test_sort_unsupported(df):
+    with pytest.raises(
+        TypeError, match=r"Expected Expr or column name.*col\(\) or lit\(\)"
+    ):
+        df.sort(1)
+
+
+def test_aggregate_string_and_expression_equivalent(df):
+    result_str = df.aggregate("a", [f.count()]).to_pydict()
+    result_expr = df.aggregate(col("a"), [f.count()]).to_pydict()
+    assert result_str == result_expr
+
+
+def test_filter_string_unsupported(df):
+    with pytest.raises(TypeError, match=r"col\(\) or lit\(\)"):
+        df.filter("a > 1")
+
+
 def test_drop(df):
     df = df.drop("c")
 
@@ -337,6 +371,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"),
@@ -368,6 +407,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(
@@ -526,6 +572,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()