feat: add ensure_expr helper for validation; refine expression handling, sorting, and docs

- Introduce `ensure_expr` helper and improve internal expression
  validation
- Update error messages and tests to consistently use `EXPR_TYPE_ERROR`
- Refactor expression handling with `_to_raw_expr`, `_ensure_expr`, and
  `SortKey`
- Improve type safety and consistency in sort key definitions and file
  sort order
- Add parameterized parquet sorting tests
- Enhance DataFrame docstrings with clearer guidance and usage examples
- Fix minor typos and error message clarity

Author: kosiew

python/datafusion/context.py

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, Sequence
 
 import pyarrow as pa
 
@@ -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[SortKey]] | None = None,
+        file_sort_order: Sequence[Sequence[SortKey]] | None = None,
     ) -> None:
         """Register multiple files as a single table.
 
@@ -805,7 +805,7 @@ def register_parquet(
         file_extension: str = ".parquet",
         skip_metadata: bool = True,
         schema: pa.Schema | None = None,
-        file_sort_order: list[list[SortKey]] | None = None,
+        file_sort_order: Sequence[Sequence[SortKey]] | None = None,
     ) -> None:
         """Register a Parquet file as a table.
 
@@ -1096,7 +1096,7 @@ def read_parquet(
         file_extension: str = ".parquet",
         skip_metadata: bool = True,
         schema: pa.Schema | None = None,
-        file_sort_order: list[list[SortKey]] | None = None,
+        file_sort_order: Sequence[Sequence[SortKey]] | None = None,
     ) -> DataFrame:
         """Read a Parquet source into a :py:class:`~datafusion.dataframe.Dataframe`.
 
@@ -1176,8 +1176,16 @@ def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream:
 
     @staticmethod
     def _convert_file_sort_order(
-        file_sort_order: list[list[Expr | SortExpr | str]] | None,
+        file_sort_order: Sequence[Sequence[SortKey]] | None,
     ) -> list[list[Any]] | None:
+        """Convert nested ``SortKey`` sequences into raw sort representations.
+
+        Each ``SortKey`` can be a column name string, an ``Expr``, or a
+        ``SortExpr`` and will be converted using
+        :func:`datafusion.expr.sort_list_to_raw_sort_list`.
+        """
+        # Convert each ``SortKey`` in the provided sort order to the low-level
+        # representation expected by the Rust bindings.
         return (
             [sort_list_to_raw_sort_list(f) for f in file_sort_order]
             if file_sort_order is not None

python/datafusion/dataframe.py

@@ -41,9 +41,9 @@
 from datafusion._internal import ParquetColumnOptions as ParquetColumnOptionsInternal
 from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal
 from datafusion.expr import (
-    EXPR_TYPE_ERROR,
     Expr,
     SortKey,
+    ensure_expr,
     expr_list_to_raw_expr_list,
     sort_list_to_raw_sort_list,
 )
@@ -58,8 +58,6 @@
     import polars as pl
     import pyarrow as pa
 
-    from datafusion._internal import expr as expr_internal
-
 from enum import Enum
 
 
@@ -418,9 +416,17 @@ def filter(self, *predicates: Expr) -> DataFrame:
         """Return a DataFrame for which ``predicate`` evaluates to ``True``.
 
         Rows for which ``predicate`` evaluates to ``False`` or ``None`` are filtered
-        out.  If more than one predicate is provided, these predicates will be
-        combined as a logical AND. If more complex logic is required, see the
-        logical operations in :py:mod:`~datafusion.functions`.
+        out. If more than one predicate is provided, these predicates will be
+        combined as a logical AND. Each ``predicate`` must be an
+        :class:`~datafusion.expr.Expr` created using helper functions such as
+        :func:`datafusion.col` or :func:`datafusion.lit`; plain strings are not
+        accepted. If more complex logic is required, see the logical operations in
+        :py:mod:`~datafusion.functions`.
+
+        Example::
+
+            from datafusion import col, lit
+            df.filter(col("a") > lit(1))
 
         Args:
             predicates: Predicate expression(s) to filter the DataFrame.
@@ -430,40 +436,49 @@ 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)
+            df = df.filter(ensure_expr(p))
         return DataFrame(df)
 
     def with_column(self, name: str, expr: Expr) -> DataFrame:
         """Add an additional column to the DataFrame.
 
+        The ``expr`` must be an :class:`~datafusion.expr.Expr` constructed with
+        :func:`datafusion.col` or :func:`datafusion.lit`; plain strings are not
+        accepted.
+
+        Example::
+
+            from datafusion import col, lit
+            df.with_column("b", col("a") + lit(1))
+
         Args:
             name: Name of the column to add.
             expr: Expression to compute the column.
 
         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))
+        return DataFrame(self.df.with_column(name, ensure_expr(expr)))
 
     def with_columns(
         self, *exprs: Expr | Iterable[Expr], **named_exprs: Expr
     ) -> DataFrame:
         """Add columns to the DataFrame.
 
-        By passing expressions, iteratables of expressions, or named expressions. To
-        pass named expressions use the form name=Expr.
+        By passing expressions, iterables of expressions, or named expressions.
+        All expressions must be :class:`~datafusion.expr.Expr` objects created via
+        :func:`datafusion.col` or :func:`datafusion.lit`; plain strings are not
+        accepted. To pass named expressions use the form ``name=Expr``.
 
-        Example usage: The following will add 4 columns labeled a, b, c, and d::
+        Example usage: The following will add 4 columns labeled ``a``, ``b``, ``c``,
+        and ``d``::
 
+            from datafusion import col, lit
             df = df.with_columns(
-                lit(0).alias('a'),
-                [lit(1).alias('b'), lit(2).alias('c')],
+                col("x").alias("a"),
+                [lit(1).alias("b"), col("y").alias("c")],
                 d=lit(3)
-                )
+            )
 
         Args:
             exprs: Either a single expression or an iterable of expressions to add.
@@ -473,30 +488,19 @@ def with_columns(
             DataFrame with the new columns added.
         """
 
-        def _simplify_expression(
-            *exprs: Expr | Iterable[Expr], **named_exprs: Expr
-        ) -> list[expr_internal.Expr]:
-            expr_list: list[expr_internal.Expr] = []
-            for expr in exprs:
+        def _iter_exprs(items: Iterable[Expr | Iterable[Expr]]) -> Iterable[Expr | str]:
+            for expr in items:
                 if isinstance(expr, str):
-                    raise TypeError(EXPR_TYPE_ERROR)
-                if isinstance(expr, Iterable) and not isinstance(expr, Expr):
-                    expr_value = list(expr)
-                    if any(isinstance(inner, str) for inner in expr_value):
-                        raise TypeError(EXPR_TYPE_ERROR)
+                    yield expr
+                elif isinstance(expr, Iterable) and not isinstance(expr, Expr):
+                    yield from _iter_exprs(expr)
                 else:
-                    expr_value = expr
-                try:
-                    expr_list.extend(expr_list_to_raw_expr_list(expr_value))
-                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)
-            return expr_list
-
-        expressions = _simplify_expression(*exprs, **named_exprs)
+                    yield expr
+
+        expressions = [ensure_expr(e) for e in _iter_exprs(exprs)]
+        for alias, expr in named_exprs.items():
+            ensure_expr(expr)
+            expressions.append(expr.alias(alias).expr)
 
         return DataFrame(self.df.with_columns(expressions))
 
@@ -535,11 +539,7 @@ def aggregate(
         aggs_list = 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)
+        aggs_exprs = [ensure_expr(agg) for agg in aggs_list]
         return DataFrame(self.df.aggregate(group_by_exprs, aggs_exprs))
 
     def sort(self, *exprs: SortKey) -> DataFrame:
@@ -554,7 +554,7 @@ def sort(self, *exprs: SortKey) -> DataFrame:
         Returns:
             DataFrame after sorting.
         """
-        exprs_raw = sort_list_to_raw_sort_list(list(exprs))
+        exprs_raw = sort_list_to_raw_sort_list(exprs)
         return DataFrame(self.df.sort(*exprs_raw))
 
     def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame:
@@ -766,8 +766,15 @@ def join_on(
     ) -> DataFrame:
         """Join two :py:class:`DataFrame` using the specified expressions.
 
-        On expressions are used to support in-equality predicates. Equality
-        predicates are correctly optimized
+        Join predicates must be :class:`~datafusion.expr.Expr` objects, typically
+        built with :func:`datafusion.col`; plain strings are not accepted. On
+        expressions are used to support in-equality predicates. Equality predicates
+        are correctly optimized.
+
+        Example::
+
+            from datafusion import col
+            df.join_on(other_df, col("id") == col("other_id"))
 
         Args:
             right: Other DataFrame to join with.
@@ -778,11 +785,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 = [ensure_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, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence
 
 import pyarrow as pa
 
@@ -41,9 +41,8 @@
 
 
 # Standard error message for invalid expression types
-EXPR_TYPE_ERROR = "Use col() or lit() to construct expressions"
-
-SortKey = Union["Expr", "SortExpr", str]
+# Mention both alias forms of column and literal helpers
+EXPR_TYPE_ERROR = "Use col()/column() or lit()/literal() 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
@@ -219,9 +218,54 @@
     "WindowExpr",
     "WindowFrame",
     "WindowFrameBound",
+    "ensure_expr",
 ]
 
 
+def ensure_expr(value: Expr | Any) -> expr_internal.Expr:
+    """Return the internal expression from ``Expr`` or raise ``TypeError``.
+
+    This helper rejects plain strings and other non-:class:`Expr` values so
+    higher level APIs consistently require explicit :func:`~datafusion.col` or
+    :func:`~datafusion.lit` expressions.
+
+    Args:
+        value: Candidate expression or other object.
+
+    Returns:
+        The internal expression representation.
+
+    Raises:
+        TypeError: If ``value`` is not an instance of :class:`Expr`.
+    """
+    if not isinstance(value, Expr):
+        raise TypeError(EXPR_TYPE_ERROR)
+    return value.expr
+
+
+def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
+    """Convert a Python expression or column name to its raw variant.
+
+    Args:
+        value: Candidate expression or column name.
+
+    Returns:
+        The internal :class:`~datafusion._internal.expr.Expr` representation.
+
+    Raises:
+        TypeError: If ``value`` is neither an :class:`Expr` nor ``str``.
+    """
+    if isinstance(value, str):
+        return Expr.column(value).expr
+    if isinstance(value, Expr):
+        return value.expr
+    error = (
+        "Expected Expr or column name, found:"
+        f" {type(value).__name__}. {EXPR_TYPE_ERROR}."
+    )
+    raise TypeError(error)
+
+
 def expr_list_to_raw_expr_list(
     expr_list: Optional[Sequence[Expr | str] | Expr | str],
 ) -> Optional[list[expr_internal.Expr]]:
@@ -230,30 +274,18 @@ def expr_list_to_raw_expr_list(
         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 [_to_raw_expr(e) for e in expr_list]
 
 
 def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
-    """Helper function to return a default Sort if an Expr is provided."""
+    """Return a :class:`SortExpr`, defaulting attributes when necessary."""
     if isinstance(e, SortExpr):
         return e.raw_sort
     return SortExpr(e, ascending=True, nulls_first=True).raw_sort
 
 
 def sort_list_to_raw_sort_list(
-    sort_list: Optional[list[SortKey] | SortKey],
+    sort_list: Optional[Sequence[SortKey] | SortKey],
 ) -> Optional[list[expr_internal.SortExpr]]:
     """Helper function to return an optional sort list to raw variant."""
     if isinstance(sort_list, (Expr, SortExpr, str)):
@@ -262,17 +294,11 @@ def sort_list_to_raw_sort_list(
         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
+        if isinstance(item, SortExpr):
+            raw_sort_list.append(sort_or_default(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))
+            raw_expr = _to_raw_expr(item)  # may raise ``TypeError``
+            raw_sort_list.append(sort_or_default(Expr(raw_expr)))
     return raw_sort_list
 
 
@@ -1335,3 +1361,6 @@ def nulls_first(self) -> bool:
     def __repr__(self) -> str:
         """Generate a string representation of this expression."""
         return self.raw_sort.__repr__()
+
+
+SortKey = Expr | SortExpr | str