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

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Any, Protocol, Sequence
+from typing import TYPE_CHECKING, Any, Protocol
 
 import pyarrow as pa
 
@@ -31,7 +31,7 @@
 
 from datafusion.catalog import Catalog, CatalogProvider, Table
 from datafusion.dataframe import DataFrame
-from datafusion.expr import SortKey, sort_list_to_raw_sort_list
+from datafusion.expr import Expr, SortExpr, sort_list_to_raw_sort_list
 from datafusion.record_batch import RecordBatchStream
 from datafusion.user_defined import AggregateUDF, ScalarUDF, TableFunction, WindowUDF
 
@@ -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: Sequence[Sequence[SortKey]] | None = None,
+        file_sort_order: list[list[Expr | SortExpr]] | None = None,
     ) -> None:
         """Register multiple files as a single table.
 
@@ -567,20 +567,23 @@ def register_listing_table(
             table_partition_cols: Partition columns.
             file_extension: File extension of the provided table.
             schema: The data source schema.
-            file_sort_order: Sort order for the file. Each sort key can be
-                specified as a column name (``str``), an expression
-                (``Expr``), or a ``SortExpr``.
+            file_sort_order: Sort order for the file.
         """
         if table_partition_cols is None:
             table_partition_cols = []
         table_partition_cols = self._convert_table_partition_cols(table_partition_cols)
+        file_sort_order_raw = (
+            [sort_list_to_raw_sort_list(f) for f in file_sort_order]
+            if file_sort_order is not None
+            else None
+        )
         self.ctx.register_listing_table(
             name,
             str(path),
             table_partition_cols,
             file_extension,
             schema,
-            self._convert_file_sort_order(file_sort_order),
+            file_sort_order_raw,
         )
 
     def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame:
@@ -805,7 +808,7 @@ def register_parquet(
         file_extension: str = ".parquet",
         skip_metadata: bool = True,
         schema: pa.Schema | None = None,
-        file_sort_order: Sequence[Sequence[SortKey]] | None = None,
+        file_sort_order: list[list[SortExpr]] | None = None,
     ) -> None:
         """Register a Parquet file as a table.
 
@@ -824,9 +827,7 @@ def register_parquet(
                 that may be in the file schema. This can help avoid schema
                 conflicts due to metadata.
             schema: The data source schema.
-            file_sort_order: Sort order for the file. Each sort key can be
-                specified as a column name (``str``), an expression
-                (``Expr``), or a ``SortExpr``.
+            file_sort_order: Sort order for the file.
         """
         if table_partition_cols is None:
             table_partition_cols = []
@@ -839,7 +840,9 @@ def register_parquet(
             file_extension,
             skip_metadata,
             schema,
-            self._convert_file_sort_order(file_sort_order),
+            [sort_list_to_raw_sort_list(exprs) for exprs in file_sort_order]
+            if file_sort_order is not None
+            else None,
         )
 
     def register_csv(
@@ -1096,7 +1099,7 @@ def read_parquet(
         file_extension: str = ".parquet",
         skip_metadata: bool = True,
         schema: pa.Schema | None = None,
-        file_sort_order: Sequence[Sequence[SortKey]] | None = None,
+        file_sort_order: list[list[Expr | SortExpr]] | None = None,
     ) -> DataFrame:
         """Read a Parquet source into a :py:class:`~datafusion.dataframe.Dataframe`.
 
@@ -1113,17 +1116,19 @@ def read_parquet(
             schema: An optional schema representing the parquet files. If None,
                 the parquet reader will try to infer it based on data in the
                 file.
-            file_sort_order: Sort order for the file. Each sort key can be
-                specified as a column name (``str``), an expression
-                (``Expr``), or a ``SortExpr``.
+            file_sort_order: Sort order for the file.
 
         Returns:
             DataFrame representation of the read Parquet files
         """
         if table_partition_cols is None:
             table_partition_cols = []
         table_partition_cols = self._convert_table_partition_cols(table_partition_cols)
-        file_sort_order = self._convert_file_sort_order(file_sort_order)
+        file_sort_order = (
+            [sort_list_to_raw_sort_list(f) for f in file_sort_order]
+            if file_sort_order is not None
+            else None
+        )
         return DataFrame(
             self.ctx.read_parquet(
                 str(path),
@@ -1174,24 +1179,6 @@ def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream:
         """Execute the ``plan`` and return the results."""
         return RecordBatchStream(self.ctx.execute(plan._raw_plan, partitions))
 
-    @staticmethod
-    def _convert_file_sort_order(
-        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
-            else None
-        )
-
     @staticmethod
     def _convert_table_partition_cols(
         table_partition_cols: list[tuple[str, str | pa.DataType]],