From a8013665fd063468e1845279a667b527a4f13270 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Mon, 3 Aug 2026 06:53:56 +0000 Subject: [PATCH 1/2] [SPARK-58522][PYTHON] Remove type: ignore[assignment] in Window/TableArg column helpers Use tuple(cols[0]) instead of cols = cols[0] in the single-sequence unwrap of _to_cols (connect Window, connect TableArg) and _to_java_cols (classic Window). tuple() keeps cols tuple-typed, so the reused variable no longer needs # type: ignore[assignment]. No behavior change. Co-authored-by: Isaac --- python/pyspark/sql/classic/window.py | 2 +- python/pyspark/sql/connect/table_arg.py | 2 +- python/pyspark/sql/connect/window.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/pyspark/sql/classic/window.py b/python/pyspark/sql/classic/window.py index 0a9e0b3c96c64..ce56f113d738a 100644 --- a/python/pyspark/sql/classic/window.py +++ b/python/pyspark/sql/classic/window.py @@ -37,7 +37,7 @@ def _to_java_cols( 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] + cols = tuple(cols[0]) sc = get_active_spark_context() return _to_seq(sc, cast(Iterable["ColumnOrName"], cols), _to_java_column) diff --git a/python/pyspark/sql/connect/table_arg.py b/python/pyspark/sql/connect/table_arg.py index 789f241de509b..295c158689cae 100644 --- a/python/pyspark/sql/connect/table_arg.py +++ b/python/pyspark/sql/connect/table_arg.py @@ -40,7 +40,7 @@ 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] + cols = tuple(cols[0]) return [F._to_col(c) for c in cast(Iterable["ColumnOrName"], cols)] diff --git a/python/pyspark/sql/connect/window.py b/python/pyspark/sql/connect/window.py index 55a43c2158a3b..36d31642cef06 100644 --- a/python/pyspark/sql/connect/window.py +++ b/python/pyspark/sql/connect/window.py @@ -32,7 +32,7 @@ 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] + cols = tuple(cols[0]) return [F._to_col(c) for c in cast(Iterable["ColumnOrName"], cols)] From 2f4a30386cd6bdb535ac50f2e67c52616ff6a5fc Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Mon, 3 Aug 2026 18:10:35 +0000 Subject: [PATCH 2/2] [SPARK-58522][PYTHON] Accept a single tuple of columns in Window and TableArg partitionBy/orderBy Window.partitionBy/orderBy (on Window and WindowSpec) and TableArg.partitionBy/orderBy unwrapped a single collection of columns only when it was a list, so a tuple was treated as one invalid column argument. Widen the runtime check to any non-str Sequence, widen the implementation signatures to Union[ColumnOrName, Sequence[ColumnOrName]] across the base, classic and connect layers, and add the two-overload public signature used by select/groupBy/describe. Window's annotation already claimed Sequence support the runtime did not provide; TableArg's admitted no sequence at all even though the runtime and docstrings did. The two new # type: ignore[arg-type] in connect Window follow the existing DataFrame.agg -> groupBy().agg(*exprs) precedent: the wide varargs are forwarded into a now-overloaded method, which the overloads cannot express. Co-authored-by: Isaac --- python/pyspark/sql/classic/table_arg.py | 30 +++++++++----- python/pyspark/sql/classic/window.py | 32 ++++++++++++++- python/pyspark/sql/connect/table_arg.py | 19 +++++++-- python/pyspark/sql/connect/window.py | 47 ++++++++++++++++++++-- python/pyspark/sql/table_arg.py | 18 +++++++-- python/pyspark/sql/tests/test_functions.py | 18 +++++++++ python/pyspark/sql/tests/test_udtf.py | 20 +++++++++ python/pyspark/sql/window.py | 30 +++++++++++++- 8 files changed, 192 insertions(+), 22 deletions(-) diff --git a/python/pyspark/sql/classic/table_arg.py b/python/pyspark/sql/classic/table_arg.py index 97f6fa79e1894..79f46beedd52d 100644 --- a/python/pyspark/sql/classic/table_arg.py +++ b/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) diff --git a/python/pyspark/sql/classic/window.py b/python/pyspark/sql/classic/window.py index ce56f113d738a..c6a35d3e0e02d 100644 --- a/python/pyspark/sql/classic/window.py +++ b/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): + 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))) diff --git a/python/pyspark/sql/connect/table_arg.py b/python/pyspark/sql/connect/table_arg.py index 295c158689cae..4e4c08c3f53a8 100644 --- a/python/pyspark/sql/connect/table_arg.py +++ b/python/pyspark/sql/connect/table_arg.py @@ -17,6 +17,7 @@ from typing import ( Iterable, + overload, TYPE_CHECKING, Union, Sequence, @@ -39,7 +40,7 @@ def _to_cols(cols: Tuple[Union["ColumnOrName", Sequence["ColumnOrName"]], ...]) -> List[Column]: - if len(cols) == 1 and isinstance(cols[0], list): + 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()." diff --git a/python/pyspark/sql/connect/window.py b/python/pyspark/sql/connect/window.py index 36d31642cef06..69fa0bb2fb8e2 100644 --- a/python/pyspark/sql/connect/window.py +++ b/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,7 +42,7 @@ def _to_cols(cols: Tuple[Union["ColumnOrName", Sequence["ColumnOrName"]], ...]) -> List[Column]: - if len(cols) == 1 and isinstance(cols[0], list): + 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,6 +93,12 @@ 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] @@ -89,6 +106,12 @@ def partitionBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> 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": diff --git a/python/pyspark/sql/table_arg.py b/python/pyspark/sql/table_arg.py index e7338f2023a06..4fc0f7b8afd9c 100644 --- a/python/pyspark/sql/table_arg.py +++ b/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. diff --git a/python/pyspark/sql/tests/test_functions.py b/python/pyspark/sql/tests/test_functions.py index 41764083c3c84..0cb05d9faa6a1 100644 --- a/python/pyspark/sql/tests/test_functions.py +++ b/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"]) diff --git a/python/pyspark/sql/tests/test_udtf.py b/python/pyspark/sql/tests/test_udtf.py index 2b30e41b30db2..9bb05eb4810db 100644 --- a/python/pyspark/sql/tests/test_udtf.py +++ b/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()), [ diff --git a/python/pyspark/sql/window.py b/python/pyspark/sql/window.py index 0b2687a038866..5a5d1e06c4ef4 100644 --- a/python/pyspark/sql/window.py +++ b/python/pyspark/sql/window.py @@ -18,7 +18,7 @@ # mypy: disable-error-code="empty-body" import sys -from typing import Sequence, TYPE_CHECKING, Union +from typing import overload, Sequence, TYPE_CHECKING, Union from pyspark.sql.utils import dispatch_window_method from pyspark.util import ( @@ -66,6 +66,14 @@ class Window: currentRow: int = 0 + @overload + @staticmethod + def partitionBy(*cols: "ColumnOrName") -> "WindowSpec": ... + + @overload + @staticmethod + def partitionBy(__cols: Sequence["ColumnOrName"]) -> "WindowSpec": ... + @staticmethod @dispatch_window_method def partitionBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec": @@ -118,6 +126,14 @@ def partitionBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "Wind """ ... + @overload + @staticmethod + def orderBy(*cols: "ColumnOrName") -> "WindowSpec": ... + + @overload + @staticmethod + def orderBy(__cols: Sequence["ColumnOrName"]) -> "WindowSpec": ... + @staticmethod @dispatch_window_method def orderBy(*cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec": @@ -344,6 +360,12 @@ def __new__(cls, jspec: "JavaObject") -> "WindowSpec": return WindowSpec.__new__(WindowSpec, jspec) + @overload + def partitionBy(self, *cols: "ColumnOrName") -> "WindowSpec": ... + + @overload + def partitionBy(self, __cols: Sequence["ColumnOrName"]) -> "WindowSpec": ... + def partitionBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec": """ Defines the partitioning columns in a :class:`WindowSpec`. @@ -357,6 +379,12 @@ def partitionBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> """ ... + @overload + def orderBy(self, *cols: "ColumnOrName") -> "WindowSpec": ... + + @overload + def orderBy(self, __cols: Sequence["ColumnOrName"]) -> "WindowSpec": ... + def orderBy(self, *cols: Union["ColumnOrName", Sequence["ColumnOrName"]]) -> "WindowSpec": """ Defines the ordering columns in a :class:`WindowSpec`.