[SPARK-58522][PYTHON] Accept a single tuple of columns in Window and TableArg partitionBy/orderBy

python/pyspark/sql/classic/table_arg.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 #
 
-from typing import TYPE_CHECKING
+from typing import cast, Iterable, overload, Sequence, TYPE_CHECKING, Union
 
 from pyspark.sql.classic.column import _to_java_column, _to_seq
 from pyspark.sql.table_arg import TableArg as ParentTableArg
@@ -30,19 +30,31 @@ class TableArg(ParentTableArg):
     def __init__(self, j_table_arg: "JavaObject"):
         self._j_table_arg = j_table_arg
 
-    def partitionBy(self, *cols: "ColumnOrName") -> "TableArg":
+    @overload
+    def partitionBy(self, *cols: "ColumnOrName") -> "TableArg": ...
+
+    @overload
+    def partitionBy(self, __cols: Sequence["ColumnOrName"]) -> "TableArg": ...
+
+    def partitionBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "TableArg":
         sc = get_active_spark_context()
-        if len(cols) == 1 and isinstance(cols[0], list):
-            cols = cols[0]
-        j_cols = _to_seq(sc, cols, _to_java_column)
+        if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence):
+            cols = tuple(cols[0])
+        j_cols = _to_seq(sc, cast(Iterable["ColumnOrName"], cols), _to_java_column)
         new_j_table_arg = self._j_table_arg.partitionBy(j_cols)
         return TableArg(new_j_table_arg)
 
-    def orderBy(self, *cols: "ColumnOrName") -> "TableArg":
+    @overload
+    def orderBy(self, *cols: "ColumnOrName") -> "TableArg": ...
+
+    @overload
+    def orderBy(self, __cols: Sequence["ColumnOrName"]) -> "TableArg": ...
+
+    def orderBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "TableArg":
         sc = get_active_spark_context()
-        if len(cols) == 1 and isinstance(cols[0], list):
-            cols = cols[0]
-        j_cols = _to_seq(sc, cols, _to_java_column)
+        if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence):
+            cols = tuple(cols[0])
+        j_cols = _to_seq(sc, cast(Iterable["ColumnOrName"], cols), _to_java_column)
         new_j_table_arg = self._j_table_arg.orderBy(j_cols)
         return TableArg(new_j_table_arg)
 

python/pyspark/sql/classic/window.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 #
 import sys
-from typing import cast, Iterable, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import cast, Iterable, overload, Sequence, Tuple, TYPE_CHECKING, Union
 
 from pyspark.sql.window import (
     Window as ParentWindow,
@@ -36,13 +36,21 @@ def _to_java_cols(
 ) -> "JavaObject":
     from pyspark.sql.classic.column import _to_seq, _to_java_column
 
-    if len(cols) == 1 and isinstance(cols[0], list):
-        cols = cols[0]  # type: ignore[assignment]
+    if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence):
+        cols = tuple(cols[0])
     sc = get_active_spark_context()
     return _to_seq(sc, cast(Iterable["ColumnOrName"], cols), _to_java_column)
 
 
 class Window(ParentWindow):
+    @overload
+    @staticmethod
+    def partitionBy(*cols: "ColumnOrName") -> ParentWindowSpec: ...
+
+    @overload
+    @staticmethod
+    def partitionBy(__cols: Sequence["ColumnOrName"]) -> ParentWindowSpec: ...
+
     @staticmethod
     def partitionBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> ParentWindowSpec:
         from py4j.java_gateway import JVMView
@@ -53,6 +61,14 @@ def partitionBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> Paren
         ).partitionBy(_to_java_cols(cols))
         return WindowSpec(jspec)
 
+    @overload
+    @staticmethod
+    def orderBy(*cols: "ColumnOrName") -> ParentWindowSpec: ...
+
+    @overload
+    @staticmethod
+    def orderBy(__cols: Sequence["ColumnOrName"]) -> ParentWindowSpec: ...
+
     @staticmethod
     def orderBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> ParentWindowSpec:
         from py4j.java_gateway import JVMView
@@ -100,11 +116,23 @@ def __new__(cls, jspec: "JavaObject") -> "WindowSpec":
     def __init__(self, jspec: "JavaObject") -> None:
         self._jspec = jspec
 
+    @overload
+    def partitionBy(self, *cols: "ColumnOrName") -> ParentWindowSpec: ...
+
+    @overload
+    def partitionBy(self, __cols: Sequence["ColumnOrName"]) -> ParentWindowSpec: ...
+
     def partitionBy(
         self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]
     ) -> ParentWindowSpec:
         return WindowSpec(self._jspec.partitionBy(_to_java_cols(cols)))
 
+    @overload
+    def orderBy(self, *cols: "ColumnOrName") -> ParentWindowSpec: ...
+
+    @overload
+    def orderBy(self, __cols: Sequence["ColumnOrName"]) -> ParentWindowSpec: ...
+
     def orderBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> ParentWindowSpec:
         return WindowSpec(self._jspec.orderBy(_to_java_cols(cols)))
 

python/pyspark/sql/connect/table_arg.py

@@ -17,6 +17,7 @@
 
 from typing import (
     Iterable,
+    overload,
     TYPE_CHECKING,
     Union,
     Sequence,
@@ -39,8 +40,8 @@
 
 
 def _to_cols(cols: Tuple[Union["ColumnOrName", Sequence["ColumnOrName"]], ...]) -> List[Column]:
-    if len(cols) == 1 and isinstance(cols[0], list):
-        cols = cols[0]  # type: ignore[assignment]
+    if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence):
+        cols = tuple(cols[0])
     return [F._to_col(c) for c in cast(Iterable["ColumnOrName"], cols)]
 
 
@@ -54,7 +55,13 @@ def _is_partitioned(self) -> bool:
             self._subquery_expr._with_single_partition
         )
 
-    def partitionBy(self, *cols: "ColumnOrName") -> "TableArg":
+    @overload
+    def partitionBy(self, *cols: "ColumnOrName") -> "TableArg": ...
+
+    @overload
+    def partitionBy(self, __cols: Sequence["ColumnOrName"]) -> "TableArg": ...
+
+    def partitionBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "TableArg":
         if self._is_partitioned():
             raise IllegalArgumentException(
                 "Cannot call partitionBy() after partitionBy() or "
@@ -72,7 +79,13 @@ def partitionBy(self, *cols: "ColumnOrName") -> "TableArg":
         )
         return TableArg(new_expr)
 
-    def orderBy(self, *cols: "ColumnOrName") -> "TableArg":
+    @overload
+    def orderBy(self, *cols: "ColumnOrName") -> "TableArg": ...
+
+    @overload
+    def orderBy(self, __cols: Sequence["ColumnOrName"]) -> "TableArg": ...
+
+    def orderBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "TableArg":
         if not self._is_partitioned():
             raise IllegalArgumentException(
                 "Please call partitionBy() or withSinglePartition() before orderBy()."

python/pyspark/sql/connect/window.py

@@ -14,7 +14,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import TYPE_CHECKING, Any, Union, Sequence, List, Optional, Tuple, cast, Iterable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Union,
+    Sequence,
+    List,
+    Optional,
+    Tuple,
+    cast,
+    Iterable,
+    overload,
+)
 
 from pyspark.sql.column import Column
 from pyspark.sql.window import (
@@ -31,8 +42,8 @@
 
 
 def _to_cols(cols: Tuple[Union["ColumnOrName", Sequence["ColumnOrName"]], ...]) -> List[Column]:
-    if len(cols) == 1 and isinstance(cols[0], list):
-        cols = cols[0]  # type: ignore[assignment]
+    if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence):
+        cols = tuple(cols[0])
     return [F._to_col(c) for c in cast(Iterable["ColumnOrName"], cols)]
 
 
@@ -82,13 +93,25 @@ def __init__(
         self._orderSpec = orderSpec
         self._frame = frame
 
+    @overload
+    def partitionBy(self, *cols: "ColumnOrName") -> "WindowSpec": ...
+
+    @overload
+    def partitionBy(self, __cols: Sequence["ColumnOrName"]) -> "WindowSpec": ...
+
     def partitionBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec":
         return WindowSpec(
             partitionSpec=[c._expr for c in _to_cols(cols)],  # type: ignore[misc]
             orderSpec=self._orderSpec,
             frame=self._frame,
         )
 
+    @overload
+    def orderBy(self, *cols: "ColumnOrName") -> "WindowSpec": ...
+
+    @overload
+    def orderBy(self, __cols: Sequence["ColumnOrName"]) -> "WindowSpec": ...
+
     def orderBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec":
         return WindowSpec(
             partitionSpec=self._partitionSpec,
@@ -136,13 +159,29 @@ def __repr__(self) -> str:
 class Window(ParentWindow):
     _spec = WindowSpec(partitionSpec=[], orderSpec=[], frame=None)
 
+    @overload
+    @staticmethod
+    def partitionBy(*cols: "ColumnOrName") -> "WindowSpec": ...
+
+    @overload
+    @staticmethod
+    def partitionBy(__cols: Sequence["ColumnOrName"]) -> "WindowSpec": ...
+
     @staticmethod
     def partitionBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec":
-        return Window._spec.partitionBy(*cols)
+        return Window._spec.partitionBy(*cols)  # type: ignore[arg-type]
+
+    @overload
+    @staticmethod
+    def orderBy(*cols: "ColumnOrName") -> "WindowSpec": ...
+
+    @overload
+    @staticmethod
+    def orderBy(__cols: Sequence["ColumnOrName"]) -> "WindowSpec": ...
 
     @staticmethod
     def orderBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec":
-        return Window._spec.orderBy(*cols)
+        return Window._spec.orderBy(*cols)  # type: ignore[arg-type]
 
     @staticmethod
     def rowsBetween(start: int, end: int) -> "WindowSpec":

python/pyspark/sql/table_arg.py

@@ -17,7 +17,7 @@
 
 # mypy: disable-error-code="empty-body"
 
-from typing import TYPE_CHECKING
+from typing import overload, Sequence, TYPE_CHECKING, Union
 
 from pyspark.sql.tvf_argument import TableValuedFunctionArgument
 from pyspark.sql.utils import dispatch_table_arg_method
@@ -35,8 +35,14 @@ class TableArg(TableValuedFunctionArgument):
     to TVF(Table-Valued Function)s including UDTF(User-Defined Table Function)s.
     """
 
+    @overload
+    def partitionBy(self, *cols: "ColumnOrName") -> "TableArg": ...
+
+    @overload
+    def partitionBy(self, __cols: Sequence["ColumnOrName"]) -> "TableArg": ...
+
     @dispatch_table_arg_method
-    def partitionBy(self, *cols: "ColumnOrName") -> "TableArg":
+    def partitionBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "TableArg":
         """
         Partitions the data based on the specified columns.
 
@@ -95,8 +101,14 @@ def partitionBy(self, *cols: "ColumnOrName") -> "TableArg":
         """
         ...
 
+    @overload
+    def orderBy(self, *cols: "ColumnOrName") -> "TableArg": ...
+
+    @overload
+    def orderBy(self, __cols: Sequence["ColumnOrName"]) -> "TableArg": ...
+
     @dispatch_table_arg_method
-    def orderBy(self, *cols: "ColumnOrName") -> "TableArg":
+    def orderBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "TableArg":
         """
         Orders the data within each partition by the specified columns.
 

python/pyspark/sql/tests/test_functions.py

@@ -2059,6 +2059,24 @@ def test_window_functions_without_partitionBy(self):
         for r, ex in zip(rs, expected):
             self.assertEqual(tuple(r), ex[: len(r)])
 
+    def test_window_partitionBy_orderBy_with_sequence(self):
+        # partitionBy/orderBy accept the columns either spread out as varargs or
+        # passed as a single list/tuple; all forms should be equivalent.
+        df = self.spark.createDataFrame(
+            [(1, "a", 3), (1, "b", 3), (2, "c", 4), (2, "d", 4)], ["key", "value", "number"]
+        )
+
+        def row_numbers(w):
+            return [r[0] for r in df.select(F.row_number().over(w)).orderBy("value").collect()]
+
+        # Window.partitionBy is the static method; the chained .orderBy exercises
+        # WindowSpec.orderBy. Both accept a single list or tuple of columns.
+        varargs = Window.partitionBy("key", "number").orderBy("value", "key")
+        as_list = Window.partitionBy(["key", "number"]).orderBy(["value", "key"])
+        as_tuple = Window.partitionBy(("key", "number")).orderBy(("value", "key"))
+        self.assertEqual(row_numbers(varargs), row_numbers(as_list))
+        self.assertEqual(row_numbers(varargs), row_numbers(as_tuple))
+
     def test_window_functions_cumulative_sum(self):
         df = self.spark.createDataFrame([("one", 1), ("two", 2)], ["key", "value"])
 

python/pyspark/sql/tests/test_udtf.py

@@ -1310,6 +1310,16 @@ def eval(self, row: Row):
             ],
             checkRowOrder=True,
         )
+        assertDataFrameEqual(
+            func(df.asTable().partitionBy(("key", "number")).orderBy(df.value)),
+            [
+                Row(key=1, value="a"),
+                Row(key=1, value="b"),
+                Row(key=2, value="c"),
+                Row(key=2, value="d"),
+            ],
+            checkRowOrder=True,
+        )
         assertDataFrameEqual(
             func(df.asTable().partitionBy("key").orderBy(df.value.desc())),
             [
@@ -1330,6 +1340,16 @@ def eval(self, row: Row):
             ],
             checkRowOrder=True,
         )
+        assertDataFrameEqual(
+            func(df.asTable().partitionBy("key").orderBy(("number", "value"))),
+            [
+                Row(key=1, value="a"),
+                Row(key=1, value="b"),
+                Row(key=2, value="c"),
+                Row(key=2, value="d"),
+            ],
+            checkRowOrder=True,
+        )
         assertDataFrameEqual(
             func(df.asTable().withSinglePartition()),
             [