[HyukjinKwon]

What changes were proposed in this pull request?

SPARK-61183: PySpark gives back the session object itself from DataFrame.sparkSession and DataFrame._session, so df.sparkSession is spark holds for the DataFrames of a session, and so do df.sparkSession.streams is spark.streams and df.sparkSession.client is spark.client.

• The session objects are registered by session ID in a registry that holds them as weak references (weakref.ref). A getter that gives back a session takes the live object of its session from the registry, and makes and registers a new one only when none is alive. Sessions made apart may share a session ID (two of an explicit session_id), so the registry keeps a list per ID and compares the core sessions with the new SparkSession::is_same_session of the Rust client (a session and its clones share their client). The getters that give back a session are DataFrame.sparkSession and DataFrame._session, the streams and client of a session, dataSource, the listener setup of StreamingQueryManager.addListener, and Catalog.registerFunction. The session-making calls are getOrCreate, create, newSession and cloneSession.
• newSession and cloneSession now return the registered object, not a new one.
• DataSourceRegistration.sparkSession: the session of the registration, which PySpark has and the client lacked.

Why are the changes needed?

Each access made a new session object, so the identity of the session was lost, along with the streams and the client that are cached per object:

spark.range(2).sparkSession is spark
# PySpark 4.3.0: True
# before this change: False
spark.range(2).sparkSession.streams is spark.streams
# PySpark 4.3.0: True
# before this change: False
spark.range(2).sparkSession.client is spark.client
# PySpark 4.3.0: True
# before this change: False
spark.dataSource.sparkSession
# PySpark 4.3.0: the session
# before this change: AttributeError

Does this PR introduce any user-facing change?

Yes. DataFrame.sparkSession and DataFrame._session return the session object of the DataFrame, the same object as spark for its session, and so do the streams and client of it. newSession and cloneSession return the object of the new session, which is registered until it is dropped. DataSourceRegistration.sparkSession is added.

How was this patch tested?

• test_dataframes_give_back_their_session_object_as_pyspark in python/tests/test_dropin_offline.py. It checks the identity of the session for range, select, groupBy().count(), na.drop, read, udf, udtf and dataSource, the identity of streams and client, that two sessions do not mix (their objects and streams differ), and that a dropped newSession object is freed (weakref and gc), while the DataFrames of spark keep giving back spark, and, in a process of its own, that two sessions made with one explicit session_id are two objects, each that of its own DataFrames. It fails without this change, at the first df.sparkSession check. The same body gives the same results with official PySpark 4.3.0 (run as a script with PYTHONPATH on the reference build).
• It uses a check helper that raises, not an assert: pytest reprs the operands of a failed assert, and reprs of DataFrames ask the server for their schema, which hung the baseline run.
• The lib tests of both crates (608 passed), the build of all tests, also with wasm-udf, the offline Python suites (2228 passed, 8 skipped), and the live drop-in and Arrow suites (210 passed) and live cases (10 passed) against a local Spark 4.3.0 server. Clippy reports no diagnostics in the changed lines.

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

Generated-by: Claude Opus 5.5

This pull request and its description were written by Isaac.