Revert "revert branch UNPICK"

This reverts commit c720298.

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

@@ -426,6 +426,12 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """
         df = self.df
         for p in predicates:
+            if not isinstance(p, Expr):
+                msg = (
+                    f"Expected Expr, got {type(p).__name__}. "
+                    "Use col() or lit() to construct expressions."
+                )
+                raise TypeError(msg)
             df = df.filter(p.expr)
         return DataFrame(df)
 
@@ -503,37 +509,44 @@ 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.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]
+        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 = [
+            sort_or_default(Expr.column(expr) if isinstance(expr, str) else expr)
+            for expr in exprs
+        ]
         return DataFrame(self.df.sort(*exprs_raw))
 
     def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame:

python/tests/test_dataframe.py

@@ -268,6 +268,19 @@ 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_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")