From 05d596aa0f98be2fc4827abec321eac8c9d0514e Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Tue, 28 Jul 2026 18:56:13 +0000 Subject: [PATCH] [SPARK-58410][PYTHON] Remove unnecessary type: ignore comments in the DataFrame read/write API ### What changes were proposed in this pull request? Removes 7 `# type: ignore` comments in the PySpark DataFrame read/write API (DataFrameReader/DataFrameWriter and their streaming and Spark Connect variants) by improving the underlying annotations: 1. `OptionUtils._set_opts` (classic + Connect): declare a `SupportsOption` Protocol and annotate `self` with it, so `self.option(...)` type-checks without an `[attr-defined]` ignore. 2. `DataFrameReader.load` (Connect): declare `paths: Optional[List[str]]` up front instead of reusing the widely-typed `path` local, removing an `[arg-type]` ignore. 3. `DataStreamWriter.partitionBy` / `clusterBy` (classic + Connect): widen the overload implementation signatures to `*cols: Union[str, List[str]]` to match their overloads (and the non-streaming `DataFrameWriter`), and normalize into a `Sequence[str]` local, removing four `[misc]` ignores without adding new ones. ### Why are the changes needed? These ignores sit on fixable annotation gaps rather than genuine type deviations. Removing them documents the intended contracts and lets mypy type-check the affected code paths (e.g. the streaming `partitionBy` list branch was previously dead code to the type checker). ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests. Verified with mypy over the full `python/pyspark` scope. Generated-by: Claude Code (Opus 4.8) --- python/pyspark/sql/_typing.pyi | 3 +++ python/pyspark/sql/connect/_typing.py | 4 ++++ python/pyspark/sql/connect/readwriter.py | 14 +++++++------- .../sql/connect/streaming/readwriter.py | 18 +++++++++++------- python/pyspark/sql/readwriter.py | 6 +++--- python/pyspark/sql/streaming/readwriter.py | 18 +++++++++++------- 6 files changed, 39 insertions(+), 24 deletions(-) diff --git a/python/pyspark/sql/_typing.pyi b/python/pyspark/sql/_typing.pyi index 94e3ccf770939..d6461e5770a7a 100644 --- a/python/pyspark/sql/_typing.pyi +++ b/python/pyspark/sql/_typing.pyi @@ -77,6 +77,9 @@ class SupportsProcess(Protocol): class SupportsClose(Protocol): def close(self, error: Exception) -> None: ... +class SupportsOption(Protocol): + def option(self, key: str, value: OptionalPrimitiveType) -> Any: ... + class UserDefinedFunctionLike(Protocol): func: Callable[..., Any] evalType: int diff --git a/python/pyspark/sql/connect/_typing.py b/python/pyspark/sql/connect/_typing.py index 50bb02be09484..95719e490c749 100644 --- a/python/pyspark/sql/connect/_typing.py +++ b/python/pyspark/sql/connect/_typing.py @@ -70,6 +70,10 @@ ] +class SupportsOption(Protocol): + def option(self, key: str, value: OptionalPrimitiveType) -> Any: ... + + class UserDefinedFunctionLike(Protocol): func: Callable[..., Any] evalType: int diff --git a/python/pyspark/sql/connect/readwriter.py b/python/pyspark/sql/connect/readwriter.py index 5c2c0c80ccdfb..eb3bd6a469b1b 100644 --- a/python/pyspark/sql/connect/readwriter.py +++ b/python/pyspark/sql/connect/readwriter.py @@ -45,7 +45,7 @@ if TYPE_CHECKING: from pyspark.sql.connect.dataframe import DataFrame - from pyspark.sql.connect._typing import ColumnOrName, OptionalPrimitiveType + from pyspark.sql.connect._typing import ColumnOrName, OptionalPrimitiveType, SupportsOption from pyspark.sql.connect.session import SparkSession from pyspark.sql.metrics import ExecutionInfo @@ -57,7 +57,7 @@ class OptionUtils: def _set_opts( - self, + self: "SupportsOption", schema: Optional[Union[StructType, str]] = None, **options: "OptionalPrimitiveType", ) -> None: @@ -68,7 +68,7 @@ def _set_opts( self.schema(schema) # type: ignore[attr-defined] for k, v in options.items(): if v is not None: - self.option(k, v) # type: ignore[attr-defined] + self.option(k, v) class DataFrameReader(OptionUtils): @@ -130,15 +130,15 @@ def load( self.schema(schema) self.options(**options) - paths = path - if isinstance(path, str): - paths = [path] + paths: Optional[List[str]] = None + if path is not None: + paths = [path] if isinstance(path, str) else path plan = DataSource( format=self._format, schema=self._schema, options=self._options, - paths=paths, # type: ignore[arg-type] + paths=paths, ) return self._df(plan) diff --git a/python/pyspark/sql/connect/streaming/readwriter.py b/python/pyspark/sql/connect/streaming/readwriter.py index 130844309ae4c..1c829476cbf0c 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 ( @@ -510,13 +510,15 @@ def partitionBy(self, *cols: str) -> "DataStreamWriter": ... @overload def partitionBy(self, __cols: List[str]) -> "DataStreamWriter": ... - def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + def partitionBy(self, *cols: Union[str, List[str]]) -> "DataStreamWriter": if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] + columns: Sequence[str] = cols[0] + else: + columns = cast("Sequence[str]", cols) # Clear any existing columns (if any). while len(self._write_proto.partitioning_column_names) > 0: self._write_proto.partitioning_column_names.pop() - self._write_proto.partitioning_column_names.extend(cast(List[str], cols)) + self._write_proto.partitioning_column_names.extend(columns) return self partitionBy.__doc__ = PySparkDataStreamWriter.partitionBy.__doc__ @@ -527,13 +529,15 @@ def clusterBy(self, *cols: str) -> "DataStreamWriter": ... @overload def clusterBy(self, __cols: List[str]) -> "DataStreamWriter": ... - def clusterBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + def clusterBy(self, *cols: Union[str, List[str]]) -> "DataStreamWriter": if len(cols) == 1 and isinstance(cols[0], (list, tuple)): - cols = cols[0] + columns: Sequence[str] = cols[0] + else: + columns = cast("Sequence[str]", cols) # Clear any existing columns (if any). while len(self._write_proto.clustering_column_names) > 0: self._write_proto.clustering_column_names.pop() - self._write_proto.clustering_column_names.extend(cast(List[str], cols)) + self._write_proto.clustering_column_names.extend(columns) return self clusterBy.__doc__ = PySparkDataStreamWriter.clusterBy.__doc__ diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py index afe8000b5c456..b64e28ab06aa8 100644 --- a/python/pyspark/sql/readwriter.py +++ b/python/pyspark/sql/readwriter.py @@ -26,7 +26,7 @@ if TYPE_CHECKING: from py4j.java_gateway import JavaObject from pyspark.core.rdd import RDD - from pyspark.sql._typing import OptionalPrimitiveType, ColumnOrName + from pyspark.sql._typing import OptionalPrimitiveType, ColumnOrName, SupportsOption from pyspark.sql.session import SparkSession from pyspark.sql.dataframe import DataFrame from pyspark.sql.streaming import StreamingQuery @@ -39,7 +39,7 @@ class OptionUtils: def _set_opts( - self, + self: "SupportsOption", schema: Optional[Union[StructType, str]] = None, **options: "OptionalPrimitiveType", ) -> None: @@ -50,7 +50,7 @@ def _set_opts( self.schema(schema) # type: ignore[attr-defined] for k, v in options.items(): if v is not None: - self.option(k, v) # type: ignore[attr-defined] + self.option(k, v) class DataFrameReader(OptionUtils): diff --git a/python/pyspark/sql/streaming/readwriter.py b/python/pyspark/sql/streaming/readwriter.py index 6b7faa6222076..bcc1ea76a4064 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 @@ -1195,7 +1195,7 @@ def partitionBy(self, *cols: str) -> "DataStreamWriter": ... @overload def partitionBy(self, __cols: List[str]) -> "DataStreamWriter": ... - def partitionBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + def partitionBy(self, *cols: Union[str, List[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 @@ -1241,8 +1241,10 @@ 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] - self._jwrite = self._jwrite.partitionBy(_to_seq(self._spark._sc, cols)) + columns: Sequence[str] = cols[0] + else: + columns = cast("Sequence[str]", cols) + self._jwrite = self._jwrite.partitionBy(_to_seq(self._spark._sc, columns)) return self @overload @@ -1251,7 +1253,7 @@ def clusterBy(self, *cols: str) -> "DataStreamWriter": ... @overload def clusterBy(self, __cols: List[str]) -> "DataStreamWriter": ... - def clusterBy(self, *cols: str) -> "DataStreamWriter": # type: ignore[misc] + def clusterBy(self, *cols: Union[str, List[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 @@ -1298,8 +1300,10 @@ 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] - self._jwrite = self._jwrite.clusterBy(_to_seq(self._spark._sc, cols)) + columns: Sequence[str] = cols[0] + else: + columns = cast("Sequence[str]", cols) + self._jwrite = self._jwrite.clusterBy(_to_seq(self._spark._sc, columns)) return self def queryName(self, queryName: str) -> "DataStreamWriter":