[SPARK-58244][PYTHON] Remove type: ignore[name-defined] for DataFrame.plot and streaming _test() in pyspark.sql

[Spenserrrr]

What changes were proposed in this pull request?

This PR removes part of the# type: ignore[name-defined] suppressions and their paired# noqa: F821 in pyspark.sql. It has two parts:

1. DataFrame.plot forward-reference annotation. The plot property in classic/dataframe.py and connect/dataframe.py returns "PySparkPlotAccessor", but that type was only imported inside the method body. Adding it to the module's if TYPE_CHECKING: block makes the forward reference resolve cleanly.
2. Fix an undefined sc in two doctest runners. In the _test() functions of python/pyspark/sql/streaming/query.py and python/pyspark/sql/streaming/listener.py, the except Py4JError: fallback calls SparkSession(sc), but sc was never defined in the function. The equivalent _test() helpers in readwriter.py and observation.py define sc = SparkContext("local[4]", "PythonTest") before the try; this PR applies the same, and drops the suppression.

Note: the [name-defined] suppressions on StreamingQueryListener._set_spark_session and its spark getter/setter are intentionally left in place. That class is shared by both classic and Connect, but the session-injection mechanism is Connect-only; annotating it properly needs a generics-based design.

Why are the changes needed?

Part 1 documents the real type and lets the type checker verify it, instead of suppressing an unresolvable annotation. Part 2 fixes a small bug: if SparkSession._getActiveSessionOrCreate() raised Py4JError, the fallback would fail with NameError: name 'sc' is not defined instead of creating a session.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

• mypy --config-file python/mypy.ini over python/pyspark tree.
  Success: no issues found with no remaining [name-defined] suppressions.
• Runtime imports of the touched modules succeed.
• The affected doctests run via python/run-tests
  (pyspark.sql.streaming.query, pyspark.sql.streaming.listener).

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 4.8)

[gaogaotiantian]

I think the dataframe.py changes are correct. However, the original _test() function is really weird - is sc expected to be injected somehow or it's just missed? @HyukjinKwon as the original author.

[Spenserrrr]

In both _test() functions, sc is not defined in the file and assigned into globs, and no doctest in these modules references sc. They only use spark, which is injected via globs["spark"]. So, on the except Py4JError: fallback path, I believe thatSparkSession(sc) would raise an error that 'sc' is not defined.
The other _test() helpers that use the same try/except fallback (readwriter.py, observation.py, merge.py) define sc = SparkContext("local[4]", "PythonTest") before the try. So I added those two lines to the file.

[gaogaotiantian]

Ah good, then it's probably just missed. Let's wait for Hyukjin's confirmation before merging. Overall I think it's good.

[Spenserrrr]

CI found an error in my fix. For StreamingQueryListener._set_spark_session's annotation, I originally pinned it to the classic SparkSession. CI reports a type error here, and I found out that itsonly callers are in Spark Connect (connect/streaming/query.py, connect/streaming/worker/listener_worker.py), which pass a connect SparkSession.
Thus, I typed all three as ConnectSparkSession instead. Union["SparkSession", "ConnectSparkSession"] would also work, and I see that this is a base class that is used by both classic and connect. I can switch to this if you prefer to keep the base class backed-agnostic. @HyukjinKwon

[gaogaotiantian]

Okay this is where things get tricky. Normally we don't import connect related stuff in non connect directory. I think we need to figure out the actual purpose here. Whether this class is meant to be used by connect-only, or both classic and connect. Type annotation is to better-define our code, not just to pass the static checker.

[Spenserrrr]

Yeah that is a good point. For the StreamingQueryListener class, it is used by both classic and connect, but with different methods:

• Classics use the _jlistener property, in sql/streaming.query.py, line 683: self._jsqm.addListener(listener._jlistener)
• Connect use the _set_spark_session(...), in sql/connect/streaming/query.py and connect/streaming/worker/listener_worker.py
  So the class is shared, but the _set_spark_session + spark getter/setter mechanism is only used by Connect.
  Another way I think of is to type it as something like Union[SparkSession, ConnectSparkSession], but it still imports connect into a non-connect dir.
  Given this, do you think the original way of neglecting this with # type: ignore[name-defined] is better?

[gaogaotiantian]

Union[SparkSession, ConnectSparkSession] is actually not the best way. The proper annotation needs some generics, otherwise the caller does not know what the returned value is. I guess that's the reason why it was ignored in the first place - like I said, there are cases where it's not easy to annotate. We can leave the question here and revert the changes to that file - if we have answers to help us move forward we can do it later.

[Spenserrrr]

Makes sense. I've reverted listener.py back to the original # type: ignore on those three members. Happy to revisit the generics approach separately if we land on a clean design later.

[dongjoon-hyun]

Thank you, @Spenserrrr .

However, the PR title sounds a little misleading because this PR doesn't remove all instances in pyspark.sql.

BEFORE THIS PR

$ git grep 'type: ignore\[name-defined]'
python/pyspark/sql/classic/dataframe.py:    def plot(self) -> "PySparkPlotAccessor":  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/connect/dataframe.py:    def plot(self) -> "PySparkPlotAccessor":  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/streaming/listener.py:        session: "SparkSession",  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/streaming/listener.py:    def spark(self) -> Optional["SparkSession"]:  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/streaming/listener.py:    def spark(self, session: "SparkSession") -> None:  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/streaming/listener.py:        spark = SparkSession(sc)  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/streaming/query.py:        spark = SparkSession(sc)  # type: ignore[name-defined] # noqa: F821

AFTER THIS PR

$ git grep 'type: ignore\[name-defined]'
python/pyspark/sql/streaming/listener.py:        session: "SparkSession",  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/streaming/listener.py:    def spark(self) -> Optional["SparkSession"]:  # type: ignore[name-defined] # noqa: F821
python/pyspark/sql/streaming/listener.py:    def spark(self, session: "SparkSession") -> None:  # type: ignore[name-defined] # noqa: F821

If we want to keep the PR title, we need to fix them all. Otherwise, please revise the PR title and description properly.

[Spenserrrr]

Thank @dongjoon-hyun for catching this! I have revised the title and description to reflect my actual scope.

[dongjoon-hyun]

+1, LGTM. I guess we need @gaogaotiantian 's final sign-off. I'll leave this to him.

[gaogaotiantian]

Thank you @dongjoon-hyun for catching the title issue!

[gaogaotiantian]

Merge Summary:

• merged into master 7d8a190
• merged into branch-4.x de4453d

Posted by merge_spark_pr.py