Enhance type annotations for file_sort_order and order_by parameters to support string inputs

Author: kosiew

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]] | None = None,
+        file_sort_order: list[list[Expr | SortExpr | str]] | 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[SortExpr]] | None = None,
+        file_sort_order: list[list[Expr | SortExpr | str]] | 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]] | None = None,
+        file_sort_order: list[list[Expr | SortExpr | str]] | None = None,
     ) -> DataFrame:
         """Read a Parquet source into a :py:class:`~datafusion.dataframe.Dataframe`.
 

python/datafusion/dataframe.py

@@ -45,7 +45,7 @@
     Expr,
     SortExpr,
     expr_list_to_raw_expr_list,
-    sort_or_default,
+    sort_list_to_raw_sort_list,
 )
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
@@ -551,20 +551,7 @@ def sort(self, *exprs: Expr | SortExpr | str) -> DataFrame:
         Returns:
             DataFrame after sorting.
         """
-        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__}. {_EXPR_TYPE_ERROR}."
-                )
-                raise TypeError(error)
+        exprs_raw = sort_list_to_raw_sort_list(list(exprs))
         return DataFrame(self.df.sort(*exprs_raw))
 
     def cast(self, mapping: dict[str, pa.DataType[Any]]) -> DataFrame:

python/datafusion/expr.py

@@ -250,12 +250,27 @@ def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
 
 
 def sort_list_to_raw_sort_list(
-    sort_list: Optional[list[Expr | SortExpr] | Expr | SortExpr],
+    sort_list: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str],
 ) -> Optional[list[expr_internal.SortExpr]]:
     """Helper function to return an optional sort list to raw variant."""
-    if isinstance(sort_list, (Expr, SortExpr)):
+    if isinstance(sort_list, (Expr, SortExpr, str)):
         sort_list = [sort_list]
-    return [sort_or_default(e) for e in sort_list] if sort_list is not None else None
+    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
 
 
 class Expr:

python/datafusion/functions.py

@@ -429,7 +429,7 @@ def window(
     name: str,
     args: list[Expr],
     partition_by: list[Expr] | Expr | None = None,
-    order_by: list[Expr | SortExpr] | Expr | SortExpr | None = None,
+    order_by: list[Expr | SortExpr | str] | Expr | SortExpr | str | None = None,
     window_frame: WindowFrame | None = None,
     ctx: SessionContext | None = None,
 ) -> Expr:
@@ -1723,7 +1723,7 @@ def array_agg(
     expression: Expr,
     distinct: bool = False,
     filter: Optional[Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Aggregate values into an array.
 
@@ -2222,7 +2222,7 @@ def regr_syy(
 def first_value(
     expression: Expr,
     filter: Optional[Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
     null_treatment: NullTreatment = NullTreatment.RESPECT_NULLS,
 ) -> Expr:
     """Returns the first value in a group of values.
@@ -2254,7 +2254,7 @@ def first_value(
 def last_value(
     expression: Expr,
     filter: Optional[Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
     null_treatment: NullTreatment = NullTreatment.RESPECT_NULLS,
 ) -> Expr:
     """Returns the last value in a group of values.
@@ -2287,7 +2287,7 @@ def nth_value(
     expression: Expr,
     n: int,
     filter: Optional[Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
     null_treatment: NullTreatment = NullTreatment.RESPECT_NULLS,
 ) -> Expr:
     """Returns the n-th value in a group of values.
@@ -2408,7 +2408,7 @@ def lead(
     shift_offset: int = 1,
     default_value: Optional[Any] = None,
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a lead window function.
 
@@ -2461,7 +2461,7 @@ def lag(
     shift_offset: int = 1,
     default_value: Optional[Any] = None,
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a lag window function.
 
@@ -2508,7 +2508,7 @@ def lag(
 
 def row_number(
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a row number window function.
 
@@ -2542,7 +2542,7 @@ def row_number(
 
 def rank(
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a rank window function.
 
@@ -2581,7 +2581,7 @@ def rank(
 
 def dense_rank(
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a dense_rank window function.
 
@@ -2615,7 +2615,7 @@ def dense_rank(
 
 def percent_rank(
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a percent_rank window function.
 
@@ -2650,7 +2650,7 @@ def percent_rank(
 
 def cume_dist(
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a cumulative distribution window function.
 
@@ -2686,7 +2686,7 @@ def cume_dist(
 def ntile(
     groups: int,
     partition_by: Optional[list[Expr] | Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Create a n-tile window function.
 
@@ -2727,7 +2727,7 @@ def string_agg(
     expression: Expr,
     delimiter: str,
     filter: Optional[Expr] = None,
-    order_by: Optional[list[Expr | SortExpr] | Expr | SortExpr] = None,
+    order_by: Optional[list[Expr | SortExpr | str] | Expr | SortExpr | str] = None,
 ) -> Expr:
     """Concatenates the input strings.