From 26f45bc9b310d650c6e909597d79ab1b5e86ed1c Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Fri, 31 Jul 2026 22:17:38 +0000 Subject: [PATCH 1/6] [SPARK-58488][PYTHON] Fix type annotations for DataFrameWriter and DataStreamWriter partitionBy/clusterBy ### What changes were proposed in this pull request? `partitionBy` and `clusterBy` on `DataFrameWriter` (batch) and `DataStreamWriter` (streaming), in both classic and Spark Connect, accept either multiple column names as varargs or a single sequence of column names. Their type annotations did not describe this accurately and relied on `# type: ignore` comments. - The single-sequence overload is changed to `__cols: Sequence[str]` and the implementation signature to `*cols: Union[str, Sequence[str]]`. - The runtime unwrap check becomes `not isinstance(cols[0], str) and isinstance(cols[0], Sequence)`, matching the widened annotation. - The streaming writers' implementation was declared `*cols: str` with a `# type: ignore[misc]` suppressing the overload mismatch; it is now honest and the `[misc]` ignores are removed. - The `# type: ignore[assignment]` comments on the unwrap lines are removed. This follows the approach in SPARK-55967, which unified and corrected the column-conversion annotations for the connect DataFrame. ### Why are the changes needed? The annotations were narrower than the runtime contract (declared `List[str]` while the code also accepts a `tuple`, i.e. any `Sequence`), and the streaming implementations did not conform to their `@overload` declarations. Correcting the annotations lets the suppression comments be removed and makes the accepted inputs explicit to users and type checkers. ### Does this PR introduce _any_ user-facing change? No. The annotations are widened to accept more (a tuple / any sequence), which is backward compatible. ### How was this patch tested? Existing tests, full-scope `mypy` over `python/pyspark`, and the typing tests under `python/pyspark/sql/tests/typing`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) --- python/pyspark/sql/connect/readwriter.py | 28 +++++++++----- .../sql/connect/streaming/readwriter.py | 30 +++++++++------ python/pyspark/sql/readwriter.py | 37 ++++++++++++++----- python/pyspark/sql/streaming/readwriter.py | 26 ++++++++----- 4 files changed, 82 insertions(+), 39 deletions(-) diff --git a/python/pyspark/sql/connect/readwriter.py b/python/pyspark/sql/connect/readwriter.py index 5c2c0c80ccdfb..c17a042e4bffb 100644 --- a/python/pyspark/sql/connect/readwriter.py +++ b/python/pyspark/sql/connect/readwriter.py @@ -15,7 +15,7 @@ # limitations under the License. # from typing import Dict -from typing import Optional, Union, List, overload, Tuple, cast, Callable +from typing import Optional, Sequence, Union, List, overload, Tuple, cast, Callable from typing import TYPE_CHECKING from pyspark.sql.connect.plan import ( @@ -620,11 +620,15 @@ def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriter": def partitionBy(self, *cols: str) -> "DataFrameWriter": ... @overload - def partitionBy(self, *cols: List[str]) -> "DataFrameWriter": ... + def partitionBy(self, __cols: Sequence[str]) -> "DataFrameWriter": ... - def partitionBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] # type: ignore[assignment] + def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": + if ( + len(cols) == 1 + and not isinstance(cols[0], str) + and isinstance(cols[0], Sequence) + ): + cols = tuple(cols[0]) self._write.partitioning_cols = cast(List[str], cols) return self @@ -736,11 +740,15 @@ def sortBy( def clusterBy(self, *cols: str) -> "DataFrameWriter": ... @overload - def clusterBy(self, *cols: List[str]) -> "DataFrameWriter": ... - - def clusterBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] # type: ignore[assignment] + def clusterBy(self, __cols: Sequence[str]) -> "DataFrameWriter": ... + + def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": + if ( + len(cols) == 1 + and not isinstance(cols[0], str) + and isinstance(cols[0], Sequence) + ): + cols = tuple(cols[0]) assert len(cols) > 0, "clusterBy needs one or more clustering columns." self._write.clustering_cols = cast(List[str], cols) return self diff --git a/python/pyspark/sql/connect/streaming/readwriter.py b/python/pyspark/sql/connect/streaming/readwriter.py index 130844309ae4c..0820d410a28b0 100644 --- a/python/pyspark/sql/connect/streaming/readwriter.py +++ b/python/pyspark/sql/connect/streaming/readwriter.py @@ -18,7 +18,7 @@ import re import sys import pickle -from typing import cast, overload, Callable, Dict, List, Optional, TYPE_CHECKING, Union +from typing import cast, overload, Callable, Dict, List, Optional, Sequence, TYPE_CHECKING, Union from pyspark.serializers import CloudPickleSerializer from pyspark.sql.connect.plan import ( @@ -508,11 +508,15 @@ def options(self, **options: "OptionalPrimitiveType") -> "DataStreamWriter": def partitionBy(self, *cols: str) -> "DataStreamWriter": ... @overload - def partitionBy(self, __cols: List[str]) -> "DataStreamWriter": ... - - def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] + def partitionBy(self, __cols: Sequence[str]) -> "DataStreamWriter": ... + + def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": + if ( + len(cols) == 1 + and not isinstance(cols[0], str) + and isinstance(cols[0], Sequence) + ): + cols = tuple(cols[0]) # Clear any existing columns (if any). while len(self._write_proto.partitioning_column_names) > 0: self._write_proto.partitioning_column_names.pop() @@ -525,11 +529,15 @@ def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] def clusterBy(self, *cols: str) -> "DataStreamWriter": ... @overload - def clusterBy(self, __cols: List[str]) -> "DataStreamWriter": ... - - def clusterBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] + def clusterBy(self, __cols: Sequence[str]) -> "DataStreamWriter": ... + + def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": + if ( + len(cols) == 1 + and not isinstance(cols[0], str) + and isinstance(cols[0], Sequence) + ): + cols = tuple(cols[0]) # Clear any existing columns (if any). while len(self._write_proto.clustering_column_names) > 0: self._write_proto.clustering_column_names.pop() diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py index afe8000b5c456..0e25c8313b88e 100644 --- a/python/pyspark/sql/readwriter.py +++ b/python/pyspark/sql/readwriter.py @@ -15,7 +15,18 @@ # limitations under the License. # import sys -from typing import cast, overload, Dict, Iterable, List, Optional, Tuple, TYPE_CHECKING, Union +from typing import ( + cast, + overload, + Dict, + Iterable, + List, + Optional, + Sequence, + Tuple, + TYPE_CHECKING, + Union, +) from pyspark.util import is_remote_only from pyspark.sql.types import StructType @@ -1497,9 +1508,9 @@ def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriter": def partitionBy(self, *cols: str) -> "DataFrameWriter": ... @overload - def partitionBy(self, *cols: List[str]) -> "DataFrameWriter": ... + def partitionBy(self, __cols: Sequence[str]) -> "DataFrameWriter": ... - def partitionBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": + def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": """Partitions the output by the given columns on the file system. If specified, the output is laid out on the file system similar @@ -1546,8 +1557,12 @@ def partitionBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": """ from pyspark.sql.classic.column import _to_seq - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - 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]) self._jwrite = self._jwrite.partitionBy( _to_seq(self._spark._sc, cast(Iterable["ColumnOrName"], cols)) ) @@ -1743,9 +1758,9 @@ def sortBy( def clusterBy(self, *cols: str) -> "DataFrameWriter": ... @overload - def clusterBy(self, *cols: List[str]) -> "DataFrameWriter": ... + def clusterBy(self, __cols: Sequence[str]) -> "DataFrameWriter": ... - def clusterBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": + def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": """Clusters the data by the given columns to optimize query performance. .. versionadded:: 4.0.0 @@ -1767,8 +1782,12 @@ def clusterBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": """ from pyspark.sql.classic.column import _to_seq - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - 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]) assert len(cols) > 0, "clusterBy needs one or more clustering columns." self._jwrite = self._jwrite.clusterBy(cols[0], _to_seq(self._spark._sc, cols[1:])) return self diff --git a/python/pyspark/sql/streaming/readwriter.py b/python/pyspark/sql/streaming/readwriter.py index 6b7faa6222076..ec6ca2c9e72cf 100644 --- a/python/pyspark/sql/streaming/readwriter.py +++ b/python/pyspark/sql/streaming/readwriter.py @@ -18,7 +18,7 @@ import re import sys from collections.abc import Iterator -from typing import cast, overload, Any, Callable, List, Optional, TYPE_CHECKING, Union +from typing import cast, overload, Any, Callable, List, Optional, Sequence, TYPE_CHECKING, Union from pyspark.sql.readwriter import OptionUtils, to_str from pyspark.sql.streaming.query import StreamingQuery @@ -1193,9 +1193,9 @@ def options(self, **options: "OptionalPrimitiveType") -> "DataStreamWriter": def partitionBy(self, *cols: str) -> "DataStreamWriter": ... @overload - def partitionBy(self, __cols: List[str]) -> "DataStreamWriter": ... + def partitionBy(self, __cols: Sequence[str]) -> "DataStreamWriter": ... - def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": """Partitions the output by the given columns on the file system. If specified, the output is laid out on the file system similar @@ -1240,8 +1240,12 @@ def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] """ from pyspark.sql.classic.column import _to_seq - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] + if ( + len(cols) == 1 + and not isinstance(cols[0], str) + and isinstance(cols[0], Sequence) + ): + cols = tuple(cols[0]) self._jwrite = self._jwrite.partitionBy(_to_seq(self._spark._sc, cols)) return self @@ -1249,9 +1253,9 @@ def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] def clusterBy(self, *cols: str) -> "DataStreamWriter": ... @overload - def clusterBy(self, __cols: List[str]) -> "DataStreamWriter": ... + def clusterBy(self, __cols: Sequence[str]) -> "DataStreamWriter": ... - def clusterBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": """Clusters the output by the given columns. If specified, the output is laid out such that records with similar values on the clustering @@ -1297,8 +1301,12 @@ def clusterBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] """ from pyspark.sql.classic.column import _to_seq - if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] + if ( + len(cols) == 1 + and not isinstance(cols[0], str) + and isinstance(cols[0], Sequence) + ): + cols = tuple(cols[0]) self._jwrite = self._jwrite.clusterBy(_to_seq(self._spark._sc, cols)) return self From 627984bc84db069b8f8484c5577bf009379fbf95 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Fri, 31 Jul 2026 23:20:23 +0000 Subject: [PATCH 2/6] [SPARK-58488][PYTHON] Fix type annotation for DataFrameNaFunctions.replace ### What changes were proposed in this pull request? `DataFrameNaFunctions.replace` (the `df.na.replace` accessor) had the same annotation issue that SPARK-56731 fixed for `DataFrame.replace`: the `Dict` overload allowed `subset` to be passed positionally after skipping `value`, which the implementation does not support, so the overloads did not conform to the implementation and were suppressed with `# type: ignore[misc]`. This adds `*,` before `subset` in the `Dict` overload of `DataFrameNaFunctions.replace` in both the base `python/pyspark/sql/dataframe.py` and classic `python/pyspark/sql/classic/dataframe.py`, making `subset` keyword-only there, and removes the now-unnecessary `# type: ignore[misc]`. Only the `@overload` annotations change; the implementation signatures are untouched, so there is no runtime behavior change. ### Why are the changes needed? To propagate the SPARK-56731 fix to the `DataFrameNaFunctions.replace` variant, which was not covered, so the overloads honestly describe the accepted call forms and the `# type: ignore[misc]` can be removed. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests, full-scope `mypy` over `python/pyspark`, and the typing tests under `python/pyspark/sql/tests/typing`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) --- python/pyspark/sql/classic/dataframe.py | 3 ++- python/pyspark/sql/dataframe.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index eeed6fef44aa4..5f58cceb4393c 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -2053,6 +2053,7 @@ def replace( def replace( self, to_replace: Dict["LiteralType", "OptionalPrimitiveType"], + *, subset: Optional[List[str]] = ..., ) -> ParentDataFrame: ... @@ -2064,7 +2065,7 @@ def replace( subset: Optional[List[str]] = ..., ) -> ParentDataFrame: ... - def replace( # type: ignore[misc] + def replace( self, to_replace: Union[List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"]], value: Optional[ diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py index 6c4d32ea1797f..91bb4fca8e9aa 100644 --- a/python/pyspark/sql/dataframe.py +++ b/python/pyspark/sql/dataframe.py @@ -7173,6 +7173,7 @@ def replace( def replace( self, to_replace: Dict["LiteralType", "OptionalPrimitiveType"], + *, subset: Optional[List[str]] = ..., ) -> DataFrame: ... @@ -7184,7 +7185,7 @@ def replace( subset: Optional[List[str]] = ..., ) -> DataFrame: ... - @dispatch_df_method # type: ignore[misc] + @dispatch_df_method def replace( self, to_replace: Union[List["LiteralType"], Dict["LiteralType", "OptionalPrimitiveType"]], From 00d2069f834f4f98ab4e7a314f6407653cc9092c Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Fri, 31 Jul 2026 23:31:05 +0000 Subject: [PATCH 3/6] [SPARK-58488][PYTHON] Remove unnecessary type: ignore in DataFrameWriter bucketBy/sortBy ### What changes were proposed in this pull request? `DataFrameWriter.bucketBy` and `DataFrameWriter.sortBy` (both classic and Spark Connect) normalize a single list/tuple of column names into a first column plus the rest with `col, cols = col[0], col[1:]`. The right-hand `col[1:]` is a slice (a list) assigned into the tuple-typed `*cols`, which required a `# type: ignore[assignment]`. Wrapping the slice in `tuple(...)` keeps `cols` a tuple, matching its declared type, so the ignore is no longer needed. The overload signatures already use `TupleOrListOfString = Union[List[str], Tuple[str, ...]]`, so they already describe the accepted inputs; only the unwrap line changes. ### Why are the changes needed? To remove unnecessary `# type: ignore[assignment]` comments by writing the normalization in a way that type-checks, consistent with the fix applied to `partitionBy`/`clusterBy`. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests, full-scope `mypy` over `python/pyspark`, and the typing tests under `python/pyspark/sql/tests/typing`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) --- python/pyspark/sql/connect/readwriter.py | 4 ++-- python/pyspark/sql/readwriter.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/connect/readwriter.py b/python/pyspark/sql/connect/readwriter.py index c17a042e4bffb..0c003873e39ae 100644 --- a/python/pyspark/sql/connect/readwriter.py +++ b/python/pyspark/sql/connect/readwriter.py @@ -663,7 +663,7 @@ def bucketBy( }, ) - col, cols = col[0], col[1:] # type: ignore[assignment] + col, cols = col[0], tuple(col[1:]) for c in cols: if not isinstance(c, str): @@ -709,7 +709,7 @@ def sortBy( }, ) - col, cols = col[0], col[1:] # type: ignore[assignment] + col, cols = col[0], tuple(col[1:]) for c in cols: if not isinstance(c, str): diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py index 0e25c8313b88e..81210a69a00f8 100644 --- a/python/pyspark/sql/readwriter.py +++ b/python/pyspark/sql/readwriter.py @@ -1643,7 +1643,7 @@ def bucketBy( }, ) - col, cols = col[0], col[1:] # type: ignore[assignment] + col, cols = col[0], tuple(col[1:]) for c in cols: if not isinstance(c, str): @@ -1727,7 +1727,7 @@ def sortBy( }, ) - col, cols = col[0], col[1:] # type: ignore[assignment] + col, cols = col[0], tuple(col[1:]) for c in cols: if not isinstance(c, str): From 5161e6dfa6660657888c64cf64326e1a05357cfe Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Fri, 31 Jul 2026 23:53:07 +0000 Subject: [PATCH 4/6] [SPARK-58488][PYTHON] Use Sequence[str] for DataFrameWriter bucketBy/sortBy col ### What changes were proposed in this pull request? `DataFrameWriter.bucketBy` and `DataFrameWriter.sortBy` (both classic and Spark Connect) typed their sequence argument with a local alias `TupleOrListOfString = Union[List[str], Tuple[str, ...]]`. This is replaced with `Sequence[str]`, matching the type used for `partitionBy`/`clusterBy` and the approach SPARK-55967 standardized on (`List` -> `Sequence`). The unused `TupleOrListOfString` alias and `Tuple` import are removed. The runtime `isinstance(col, (list, tuple))` checks are widened to `not isinstance(col, str) and isinstance(col, Sequence)` to match the annotation, consistent with the connect DataFrame `_to_cols` helper. ### Why are the changes needed? To use one consistent type (`Sequence[str]`) across all `DataFrameWriter` column varargs methods instead of two different idioms in the same file, aligned with SPARK-55967. ### Does this PR introduce _any_ user-facing change? No. The annotation and runtime check are widened to accept any sequence of strings (previously only `list`/`tuple`), which is backward compatible. ### How was this patch tested? Existing tests, full-scope `mypy` over `python/pyspark`, and the typing tests under `python/pyspark/sql/tests/typing`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) --- python/pyspark/sql/connect/readwriter.py | 15 +++++++-------- python/pyspark/sql/readwriter.py | 14 ++++++-------- 2 files changed, 13 insertions(+), 16 deletions(-) diff --git a/python/pyspark/sql/connect/readwriter.py b/python/pyspark/sql/connect/readwriter.py index 0c003873e39ae..c521fe622db26 100644 --- a/python/pyspark/sql/connect/readwriter.py +++ b/python/pyspark/sql/connect/readwriter.py @@ -15,7 +15,7 @@ # limitations under the License. # from typing import Dict -from typing import Optional, Sequence, Union, List, overload, Tuple, cast, Callable +from typing import Optional, Sequence, Union, List, overload, cast, Callable from typing import TYPE_CHECKING from pyspark.sql.connect.plan import ( @@ -52,7 +52,6 @@ __all__ = ["DataFrameReader", "DataFrameWriter"] PathOrPaths = Union[str, List[str]] -TupleOrListOfString = Union[List[str], Tuple[str, ...]] class OptionUtils: @@ -639,10 +638,10 @@ def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": def bucketBy(self, numBuckets: int, col: str, *cols: str) -> "DataFrameWriter": ... @overload - def bucketBy(self, numBuckets: int, col: TupleOrListOfString) -> "DataFrameWriter": ... + def bucketBy(self, numBuckets: int, col: Sequence[str]) -> "DataFrameWriter": ... def bucketBy( - self, numBuckets: int, col: Union[str, TupleOrListOfString], *cols: Optional[str] + self, numBuckets: int, col: Union[str, Sequence[str]], *cols: Optional[str] ) -> "DataFrameWriter": if not isinstance(numBuckets, int): raise PySparkTypeError( @@ -654,7 +653,7 @@ def bucketBy( }, ) - if isinstance(col, (list, tuple)): + if not isinstance(col, str) and isinstance(col, Sequence): if cols: raise PySparkValueError( errorClass="CANNOT_SET_TOGETHER", @@ -695,12 +694,12 @@ def bucketBy( def sortBy(self, col: str, *cols: str) -> "DataFrameWriter": ... @overload - def sortBy(self, col: TupleOrListOfString) -> "DataFrameWriter": ... + def sortBy(self, col: Sequence[str]) -> "DataFrameWriter": ... def sortBy( - self, col: Union[str, TupleOrListOfString], *cols: Optional[str] + self, col: Union[str, Sequence[str]], *cols: Optional[str] ) -> "DataFrameWriter": - if isinstance(col, (list, tuple)): + if not isinstance(col, str) and isinstance(col, Sequence): if cols: raise PySparkValueError( errorClass="CANNOT_SET_TOGETHER", diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py index 81210a69a00f8..4346291b036fe 100644 --- a/python/pyspark/sql/readwriter.py +++ b/python/pyspark/sql/readwriter.py @@ -23,7 +23,6 @@ List, Optional, Sequence, - Tuple, TYPE_CHECKING, Union, ) @@ -45,7 +44,6 @@ __all__ = ["DataFrameReader", "DataFrameWriter", "DataFrameWriterV2"] PathOrPaths = Union[str, List[str]] -TupleOrListOfString = Union[List[str], Tuple[str, ...]] class OptionUtils: @@ -1572,10 +1570,10 @@ def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": def bucketBy(self, numBuckets: int, col: str, *cols: str) -> "DataFrameWriter": ... @overload - def bucketBy(self, numBuckets: int, col: TupleOrListOfString) -> "DataFrameWriter": ... + def bucketBy(self, numBuckets: int, col: Sequence[str]) -> "DataFrameWriter": ... def bucketBy( - self, numBuckets: int, col: Union[str, TupleOrListOfString], *cols: Optional[str] + self, numBuckets: int, col: Union[str, Sequence[str]], *cols: Optional[str] ) -> "DataFrameWriter": """Buckets the output by the given columns. If specified, the output is laid out on the file system similar to Hive's bucketing scheme, @@ -1634,7 +1632,7 @@ def bucketBy( }, ) - if isinstance(col, (list, tuple)): + if not isinstance(col, str) and isinstance(col, Sequence): if cols: raise PySparkValueError( errorClass="CANNOT_SET_TOGETHER", @@ -1674,10 +1672,10 @@ def bucketBy( def sortBy(self, col: str, *cols: str) -> "DataFrameWriter": ... @overload - def sortBy(self, col: TupleOrListOfString) -> "DataFrameWriter": ... + def sortBy(self, col: Sequence[str]) -> "DataFrameWriter": ... def sortBy( - self, col: Union[str, TupleOrListOfString], *cols: Optional[str] + self, col: Union[str, Sequence[str]], *cols: Optional[str] ) -> "DataFrameWriter": """Sorts the output in each bucket by the given columns on the file system. @@ -1718,7 +1716,7 @@ def sortBy( """ from pyspark.sql.classic.column import _to_seq - if isinstance(col, (list, tuple)): + if not isinstance(col, str) and isinstance(col, Sequence): if cols: raise PySparkValueError( errorClass="CANNOT_SET_TOGETHER", From 42e8bb3b72cc500ef952f65b0d1f394711f1d139 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Sat, 1 Aug 2026 20:45:47 +0000 Subject: [PATCH 5/6] [SPARK-58488][PYTHON] Apply ruff format to writer varargs changes ### What changes were proposed in this pull request? Run `ruff format` (the Python formatter used by `dev/lint-python`) on the writer `partitionBy`/`clusterBy`/`bucketBy`/`sortBy` changes. The multi-line `if` conditions and `sortBy` signature fit within the line length, so the formatter collapses them to a single line. No logic change. ### Why are the changes needed? To pass the `ruff format` check in the Python linter CI job. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? `ruff format --check` passes; full-scope `mypy` over `python/pyspark` still clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) --- python/pyspark/sql/connect/readwriter.py | 16 +++------------- .../pyspark/sql/connect/streaming/readwriter.py | 12 ++---------- python/pyspark/sql/readwriter.py | 16 +++------------- python/pyspark/sql/streaming/readwriter.py | 12 ++---------- 4 files changed, 10 insertions(+), 46 deletions(-) diff --git a/python/pyspark/sql/connect/readwriter.py b/python/pyspark/sql/connect/readwriter.py index c521fe622db26..4c72a43380deb 100644 --- a/python/pyspark/sql/connect/readwriter.py +++ b/python/pyspark/sql/connect/readwriter.py @@ -622,11 +622,7 @@ def partitionBy(self, *cols: str) -> "DataFrameWriter": ... def partitionBy(self, __cols: Sequence[str]) -> "DataFrameWriter": ... def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) self._write.partitioning_cols = cast(List[str], cols) @@ -696,9 +692,7 @@ def sortBy(self, col: str, *cols: str) -> "DataFrameWriter": ... @overload def sortBy(self, col: Sequence[str]) -> "DataFrameWriter": ... - def sortBy( - self, col: Union[str, Sequence[str]], *cols: Optional[str] - ) -> "DataFrameWriter": + def sortBy(self, col: Union[str, Sequence[str]], *cols: Optional[str]) -> "DataFrameWriter": if not isinstance(col, str) and isinstance(col, Sequence): if cols: raise PySparkValueError( @@ -742,11 +736,7 @@ def clusterBy(self, *cols: str) -> "DataFrameWriter": ... def clusterBy(self, __cols: Sequence[str]) -> "DataFrameWriter": ... def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) assert len(cols) > 0, "clusterBy needs one or more clustering columns." self._write.clustering_cols = cast(List[str], cols) diff --git a/python/pyspark/sql/connect/streaming/readwriter.py b/python/pyspark/sql/connect/streaming/readwriter.py index 0820d410a28b0..df29fa29a9c55 100644 --- a/python/pyspark/sql/connect/streaming/readwriter.py +++ b/python/pyspark/sql/connect/streaming/readwriter.py @@ -511,11 +511,7 @@ def partitionBy(self, *cols: str) -> "DataStreamWriter": ... def partitionBy(self, __cols: Sequence[str]) -> "DataStreamWriter": ... def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) # Clear any existing columns (if any). while len(self._write_proto.partitioning_column_names) > 0: @@ -532,11 +528,7 @@ def clusterBy(self, *cols: str) -> "DataStreamWriter": ... def clusterBy(self, __cols: Sequence[str]) -> "DataStreamWriter": ... def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) # Clear any existing columns (if any). while len(self._write_proto.clustering_column_names) > 0: diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py index 4346291b036fe..028903804f440 100644 --- a/python/pyspark/sql/readwriter.py +++ b/python/pyspark/sql/readwriter.py @@ -1555,11 +1555,7 @@ def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": """ from pyspark.sql.classic.column import _to_seq - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) self._jwrite = self._jwrite.partitionBy( _to_seq(self._spark._sc, cast(Iterable["ColumnOrName"], cols)) @@ -1674,9 +1670,7 @@ def sortBy(self, col: str, *cols: str) -> "DataFrameWriter": ... @overload def sortBy(self, col: Sequence[str]) -> "DataFrameWriter": ... - def sortBy( - self, col: Union[str, Sequence[str]], *cols: Optional[str] - ) -> "DataFrameWriter": + def sortBy(self, col: Union[str, Sequence[str]], *cols: Optional[str]) -> "DataFrameWriter": """Sorts the output in each bucket by the given columns on the file system. .. versionadded:: 2.3.0 @@ -1780,11 +1774,7 @@ def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataFrameWriter": """ from pyspark.sql.classic.column import _to_seq - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) assert len(cols) > 0, "clusterBy needs one or more clustering columns." self._jwrite = self._jwrite.clusterBy(cols[0], _to_seq(self._spark._sc, cols[1:])) diff --git a/python/pyspark/sql/streaming/readwriter.py b/python/pyspark/sql/streaming/readwriter.py index ec6ca2c9e72cf..c73598fee3153 100644 --- a/python/pyspark/sql/streaming/readwriter.py +++ b/python/pyspark/sql/streaming/readwriter.py @@ -1240,11 +1240,7 @@ def partitionBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": """ from pyspark.sql.classic.column import _to_seq - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) self._jwrite = self._jwrite.partitionBy(_to_seq(self._spark._sc, cols)) return self @@ -1301,11 +1297,7 @@ def clusterBy(self, *cols: Union[str, Sequence[str]]) -> "DataStreamWriter": """ from pyspark.sql.classic.column import _to_seq - if ( - len(cols) == 1 - and not isinstance(cols[0], str) - and isinstance(cols[0], Sequence) - ): + if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence): cols = tuple(cols[0]) self._jwrite = self._jwrite.clusterBy(_to_seq(self._spark._sc, cols)) return self From b2b8a1444ea3d25f320851f18a4e0083df88cc84 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Sun, 2 Aug 2026 10:34:30 +0000 Subject: [PATCH 6/6] [SPARK-58488][PYTHON] Use Sequence[str] for save/saveAsTable partitionBy keyword arg ### What changes were proposed in this pull request? The `partitionBy` keyword argument on the writer save methods (`DataFrameWriter.save`, `saveAsTable`, `parquet`, `orc` and the `DataStreamWriter` `start`/`toTable`/`table` variants, in both classic and Spark Connect) was annotated `Optional[Union[str, List[str]]]`. These just forward to `self.partitionBy(partitionBy)`, whose implementation already accepts a list or a tuple at runtime. Widened the annotation to `Optional[Union[str, Sequence[str]]]` so it matches, and removed the now-unused `List` import from `sql/streaming/readwriter.py`. ### Why are the changes needed? To keep these `partitionBy` keyword-arg annotations consistent with the `partitionBy` method they forward to (which accepts any sequence), following SPARK-55967's use of `Sequence`. ### Does this PR introduce _any_ user-facing change? No. The annotation is widened to accept any sequence of strings (a tuple already worked at runtime), which is backward compatible. ### How was this patch tested? Existing tests, full-scope `mypy` over `python/pyspark`, `ruff` format and lint checks, and the typing tests under `python/pyspark/sql/tests/typing`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) --- python/pyspark/sql/connect/readwriter.py | 8 ++++---- python/pyspark/sql/connect/streaming/readwriter.py | 6 +++--- python/pyspark/sql/readwriter.py | 8 ++++---- python/pyspark/sql/streaming/readwriter.py | 6 +++--- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/python/pyspark/sql/connect/readwriter.py b/python/pyspark/sql/connect/readwriter.py index 4c72a43380deb..5d96bd781f616 100644 --- a/python/pyspark/sql/connect/readwriter.py +++ b/python/pyspark/sql/connect/readwriter.py @@ -749,7 +749,7 @@ def save( path: Optional[str] = None, format: Optional[str] = None, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, **options: "OptionalPrimitiveType", ) -> None: self.mode(mode).options(**options) @@ -782,7 +782,7 @@ def saveAsTable( name: str, format: Optional[str] = None, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, **options: "OptionalPrimitiveType", ) -> None: self.mode(mode).options(**options) @@ -827,7 +827,7 @@ def parquet( self, path: str, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, compression: Optional[str] = None, ) -> None: self.mode(mode) @@ -930,7 +930,7 @@ def orc( self, path: str, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, compression: Optional[str] = None, ) -> None: self.mode(mode) diff --git a/python/pyspark/sql/connect/streaming/readwriter.py b/python/pyspark/sql/connect/streaming/readwriter.py index df29fa29a9c55..b307b2d93fb95 100644 --- a/python/pyspark/sql/connect/streaming/readwriter.py +++ b/python/pyspark/sql/connect/streaming/readwriter.py @@ -681,7 +681,7 @@ def _start_internal( tableName: Optional[str] = None, format: Optional[str] = None, outputMode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, queryName: Optional[str] = None, **options: "OptionalPrimitiveType", ) -> StreamingQuery: @@ -727,7 +727,7 @@ def start( path: Optional[str] = None, format: Optional[str] = None, outputMode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, queryName: Optional[str] = None, **options: "OptionalPrimitiveType", ) -> "StreamingQuery": @@ -748,7 +748,7 @@ def toTable( tableName: str, format: Optional[str] = None, outputMode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, queryName: Optional[str] = None, **options: "OptionalPrimitiveType", ) -> "StreamingQuery": diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py index 028903804f440..e0ac0e3aaef1b 100644 --- a/python/pyspark/sql/readwriter.py +++ b/python/pyspark/sql/readwriter.py @@ -1785,7 +1785,7 @@ def save( path: Optional[str] = None, format: Optional[str] = None, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, **options: "OptionalPrimitiveType", ) -> None: """Saves the contents of the :class:`DataFrame` to a data source. @@ -1902,7 +1902,7 @@ def saveAsTable( name: str, format: Optional[str] = None, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, **options: "OptionalPrimitiveType", ) -> None: """Saves the content of the :class:`DataFrame` as the specified table. @@ -2046,7 +2046,7 @@ def parquet( self, path: str, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, compression: Optional[str] = None, ) -> None: """Saves the content of the :class:`DataFrame` in Parquet format at the specified path. @@ -2331,7 +2331,7 @@ def orc( self, path: str, mode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, compression: Optional[str] = None, ) -> None: """Saves the content of the :class:`DataFrame` in ORC format at the specified path. diff --git a/python/pyspark/sql/streaming/readwriter.py b/python/pyspark/sql/streaming/readwriter.py index c73598fee3153..9c66531938ee2 100644 --- a/python/pyspark/sql/streaming/readwriter.py +++ b/python/pyspark/sql/streaming/readwriter.py @@ -18,7 +18,7 @@ import re import sys from collections.abc import Iterator -from typing import cast, overload, Any, Callable, List, Optional, Sequence, TYPE_CHECKING, Union +from typing import cast, overload, Any, Callable, Optional, Sequence, TYPE_CHECKING, Union from pyspark.sql.readwriter import OptionUtils, to_str from pyspark.sql.streaming.query import StreamingQuery @@ -1754,7 +1754,7 @@ def start( path: Optional[str] = None, format: Optional[str] = None, outputMode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, queryName: Optional[str] = None, **options: "OptionalPrimitiveType", ) -> "StreamingQuery": @@ -1842,7 +1842,7 @@ def toTable( tableName: str, format: Optional[str] = None, outputMode: Optional[str] = None, - partitionBy: Optional[Union[str, List[str]]] = None, + partitionBy: Optional[Union[str, Sequence[str]]] = None, queryName: Optional[str] = None, **options: "OptionalPrimitiveType", ) -> "StreamingQuery":