This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 38d3af27535d [SPARK-54954][PYTHON] Fix remote related type hints in
util.py
38d3af27535d is described below
commit 38d3af27535de2604e137a4d71246210030eaeab
Author: Tian Gao <[email protected]>
AuthorDate: Fri Jan 9 07:01:42 2026 +0900
[SPARK-54954][PYTHON] Fix remote related type hints in util.py
### What changes were proposed in this pull request?
* Add overload stubs for `remote_only`.
* Distinguish remote and normal SparkSession
* Do some type assertions so we don't need all the ignore comments
### Why are the changes needed?
When we do something like
```python
property
remote_only
def method(self):
...
```
The latest version of mypy is not happy because it believes `remote_only`
can return a `property`. We can't do `property` on a `property`.
By adding the explicit overload, mypy knows that `remote_only` always
return a callable when given a callable.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Wait for CI.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Cursor (claude-4.5-opus-high)
Closes #53722 from gaogaotiantian/fix-remote-only-overload.
Authored-by: Tian Gao <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
python/pyspark/sql/utils.py | 13 ++++++++++++-
python/pyspark/util.py | 24 ++++++++++++++----------
2 files changed, 26 insertions(+), 11 deletions(-)
diff --git a/python/pyspark/sql/utils.py b/python/pyspark/sql/utils.py
index b0782d04cba3..d61a6ff642a4 100644
--- a/python/pyspark/sql/utils.py
+++ b/python/pyspark/sql/utils.py
@@ -24,6 +24,7 @@ from typing import (
Dict,
Optional,
List,
+ overload,
Sequence,
TYPE_CHECKING,
cast,
@@ -488,7 +489,17 @@ class NumpyHelper:
return [start + step * i for i in range(num)]
-def remote_only(func: Union[Callable, property]) -> Union[Callable, property]:
+@overload
+def remote_only(func: property) -> property:
+ ...
+
+
+@overload
+def remote_only(func: FuncT) -> FuncT:
+ ...
+
+
+def remote_only(func: Union[FuncT, property]) -> Union[FuncT, property]:
"""
Decorator to mark a function or method as only available in Spark Connect.
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index a9c36fb2ae6b..9f5fad57cb14 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -393,11 +393,15 @@ def inheritable_thread_target(f: Optional[Union[Callable,
"SparkSession"]] = Non
# Spark Connect
if is_remote():
+ from pyspark.sql.connect.session import SparkSession as
RemoteSparkSession
+
session = f
- assert session is not None, "Spark Connect session must be provided."
def outer(ff: Callable) -> Callable:
- thread_local = session.client.thread_local # type:
ignore[union-attr, operator]
+ assert isinstance(
+ session, RemoteSparkSession
+ ), "f is expected to be SparkSession for spark connect"
+ thread_local = session.client.thread_local
session_client_thread_local_attrs = [
(attr, copy.deepcopy(value))
for (
@@ -409,15 +413,15 @@ def inheritable_thread_target(f: Optional[Union[Callable,
"SparkSession"]] = Non
@functools.wraps(ff)
def inner(*args: Any, **kwargs: Any) -> Any:
# Propagates the active remote spark session to the current
thread.
- from pyspark.sql.connect.session import SparkSession as
RemoteSparkSession
+ assert isinstance(
+ session, RemoteSparkSession
+ ), "f is expected to be SparkSession for spark connect"
- RemoteSparkSession._set_default_and_active_session(
- session # type: ignore[arg-type]
- )
+ RemoteSparkSession._set_default_and_active_session(session)
# Set thread locals in child thread.
for attr, value in session_client_thread_local_attrs:
setattr(
- session.client.thread_local, # type:
ignore[union-attr, operator]
+ session.client.thread_local,
attr,
value,
)
@@ -437,7 +441,6 @@ def inheritable_thread_target(f: Optional[Union[Callable,
"SparkSession"]] = Non
if isinstance(f, SparkSession):
session = f
- assert session is not None
tags = set(session.getTags())
# Local properties are copied when wrapping the function.
assert SparkContext._active_spark_context is not None
@@ -561,7 +564,8 @@ class InheritableThread(threading.Thread):
def copy_local_properties(*a: Any, **k: Any) -> Any:
# Set tags in child thread.
assert hasattr(self, "_tags")
- thread_local = session.client.thread_local # type:
ignore[union-attr, operator]
+ assert session is not None
+ thread_local = session.client.thread_local
thread_local.tags = self._tags # type: ignore[has-type]
return target(*a, **k)
@@ -596,7 +600,7 @@ class InheritableThread(threading.Thread):
if is_remote():
# Spark Connect
assert hasattr(self, "_session")
- thread_local = self._session.client.thread_local # type:
ignore[union-attr, operator]
+ thread_local = self._session.client.thread_local
if not hasattr(thread_local, "tags"):
thread_local.tags = set()
self._tags = set(thread_local.tags)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
