From 83f8648fbd9db3f3b193cd694f9277068e9cc6d1 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sun, 8 Nov 2026 19:18:20 +0900 Subject: [PATCH] [SPARK-61183][CONNECT][PYTHON] Give back the SparkSession object of a DataFrame in DataFrame.sparkSession, in the Python client as PySpark does In PySpark, the DataFrames of a session, its streams and its client give back the session object itself, so DataFrame.sparkSession is spark. The client built a new session object at each access, so the identity was lost, and so were the streams and the client cached on it. Session objects are now registered by session ID, as weak references, by one function that every session-making call and session-giving getter uses: getOrCreate, create, newSession, cloneSession, DataFrame.sparkSession, dataSource, addListener and registerFunction. A getter gives back the live object of its session, and makes and registers one only when none is alive, so the registry keeps no session alive. DataSourceRegistration also gets the sparkSession attribute that PySpark has. Co-authored-by: Isaac --- crates/pyspark-rs/src/catalog.rs | 4 +- crates/pyspark-rs/src/dataframe.rs | 10 ++-- crates/pyspark-rs/src/datasource.rs | 13 ++++- crates/pyspark-rs/src/session.rs | 89 ++++++++++++++++++++++++++--- crates/pyspark-rs/src/streaming.rs | 2 +- crates/spark-connect/src/session.rs | 6 ++ python/tests/test_dropin_offline.py | 55 ++++++++++++++++++ 7 files changed, 161 insertions(+), 18 deletions(-) diff --git a/crates/pyspark-rs/src/catalog.rs b/crates/pyspark-rs/src/catalog.rs index 80680ba5..0fd6dbf9 100644 --- a/crates/pyspark-rs/src/catalog.rs +++ b/crates/pyspark-rs/src/catalog.rs @@ -495,8 +495,8 @@ impl PyCatalog { future_warning, ), )?; - let session = crate::session::PySparkSession::new(self.catalog.spark_session().clone()); - Py::new(py, session)? + let session = crate::session::session_object(py, self.catalog.spark_session())?; + session .bind(py) .getattr("udf")? .call_method1("register", (name, f, returnType)) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 68bf6924..7953a9ba 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -593,7 +593,7 @@ impl PyDataFrame { let this = slf.borrow(); py.import("pyspark.sql.connect.dataframe")? .getattr("_reduce_dataframe")? - .call1((this.spark_session(), this._relation_bytes(py))) + .call1((this.spark_session(py)?, this._relation_bytes(py))) } /// The DataFrame of `relation`, the bytes of a relation of `session`, as a pickled DataFrame @@ -1840,15 +1840,15 @@ impl PyDataFrame { /// The SparkSession that created this DataFrame. #[pyo3(name = "sparkSession")] #[getter] - fn spark_session(&self) -> crate::session::PySparkSession { - crate::session::PySparkSession::new(self.dataframe.spark_session()) + fn spark_session(&self, py: Python<'_>) -> PyResult> { + crate::session::session_object(py, &self.dataframe.spark_session()) } /// The SparkSession of this DataFrame, as PySpark Connect's `DataFrame._session`, which the /// modules of PySpark read, such as `pyspark.ml`, `pyspark.sql.plot` and `pyspark.pandas`. #[getter] - fn _session(&self) -> crate::session::PySparkSession { - self.spark_session() + fn _session(&self, py: Python<'_>) -> PyResult> { + self.spark_session(py) } /// Not supported on Spark Connect, as PySpark Connect's: `NOT_IMPLEMENTED` of `rdd`, which diff --git a/crates/pyspark-rs/src/datasource.rs b/crates/pyspark-rs/src/datasource.rs index 521d921c..20aa3849 100644 --- a/crates/pyspark-rs/src/datasource.rs +++ b/crates/pyspark-rs/src/datasource.rs @@ -19,17 +19,23 @@ use crate::session::PySparkSession; dict )] pub struct PyDataSourceRegistration { - pub(crate) session: PySparkSession, + pub(crate) session: Py, } impl PyDataSourceRegistration { - pub fn new(session: PySparkSession) -> Self { + pub fn new(session: Py) -> Self { PyDataSourceRegistration { session } } } #[pymethods] impl PyDataSourceRegistration { + /// The session of the registration, as PySpark's `DataSourceRegistration.sparkSession`. + #[getter(sparkSession)] + fn spark_session(&self, py: Python<'_>) -> Py { + self.session.clone_ref(py) + } + /// Register a custom data source class. /// /// Mirrors `pyspark.sql.SparkSession.dataSource.register`. @@ -48,7 +54,8 @@ impl PyDataSourceRegistration { let expr = CommonInlineUserDefinedDataSourceExpression::new(name, payload); // Register with the session (release GIL during the RPC) - py.detach(|| self.session.session.register_data_source(expr)) + let session = self.session.borrow(py).session.clone(); + py.detach(|| session.register_data_source(expr)) .to_pyerr()?; Ok(()) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 94aa92dd..73672cc4 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -470,11 +470,83 @@ fn new_session(py: Python<'_>, url: &str) -> PyResult> { .remote(url) .get_or_create() .to_pyerr()?; - let session = Py::new(py, PySparkSession::new(session))?; + let session = session_object(py, &session)?; set_default_and_active(py, &session); Ok(session) } +/// The Python session objects by session ID, held weakly, so that the registry keeps none +/// alive. Sessions made apart may share a session ID, such as two of an explicit `session_id`, +/// so an ID has a list of them. +type SessionObjects = std::collections::HashMap>>; + +/// The registry of the Python session objects. +fn session_objects() -> &'static std::sync::Mutex { + static OBJECTS: std::sync::OnceLock> = + std::sync::OnceLock::new(); + OBJECTS.get_or_init(|| std::sync::Mutex::new(SessionObjects::new())) +} + +/// The session object that the weak reference `weak` refers to, unless it was dropped. +fn live_session_object(py: Python<'_>, weak: &Py) -> Option> { + let object = weak.bind(py).call0().ok()?; + let session = object.cast::().ok()?; + Some(session.clone().unbind()) +} + +/// The live object of `session` itself in `registry`, if any, and not that of another session +/// of its session ID. +fn registered_session_object( + py: Python<'_>, + registry: &SessionObjects, + session: &SparkSession, +) -> Option> { + registry + .get(session.session_id())? + .iter() + .filter_map(|weak| live_session_object(py, weak)) + .find(|object| { + object + .bind(py) + .try_borrow() + .is_ok_and(|object| object.session.is_same_session(session)) + }) +} + +/// The Python object of `session`, one per session as PySpark has them: its live object, or +/// else a new one, which the registry holds weakly. So the DataFrames, the streams and the +/// clients of a session give back the session's object, as PySpark's do. +pub(crate) fn session_object( + py: Python<'_>, + session: &SparkSession, +) -> PyResult> { + { + let registry = session_objects().lock().unwrap_or_else(|e| e.into_inner()); + if let Some(object) = registered_session_object(py, ®istry, session) { + return Ok(object); + } + } + let object = Py::new(py, PySparkSession::new(session.clone()))?; + let weak = py + .import("weakref")? + .call_method1("ref", (&object,))? + .unbind(); + let mut registry = session_objects().lock().unwrap_or_else(|e| e.into_inner()); + // Another thread may have made the object of this session meanwhile; then it is that one. + if let Some(existing) = registered_session_object(py, ®istry, session) { + return Ok(existing); + } + registry.retain(|_, objects| { + objects.retain(|weak| live_session_object(py, weak).is_some()); + !objects.is_empty() + }); + registry + .entry(session.session_id().to_string()) + .or_default() + .push(weak); + Ok(object) +} + /// Make `session` the default session, and the active session of this thread, when they are /// not set, as PySpark's `SparkSession._set_default_and_active_session`. fn set_default_and_active(py: Python<'_>, session: &Py) { @@ -2244,8 +2316,11 @@ impl PySparkSession { /// /// Mirrors `SparkSession.dataSource.register`. Used to register custom data sources. #[getter] - fn dataSource(&self) -> PyDataSourceRegistration { - PyDataSourceRegistration::new(PySparkSession::new(self.session.clone())) + fn dataSource(&self, py: Python<'_>) -> PyResult { + Ok(PyDataSourceRegistration::new(session_object( + py, + &self.session, + )?)) } /// Build and register a resource profile with the server. @@ -2434,8 +2509,8 @@ impl PySparkSession { } #[pyo3(name = "newSession")] - fn new_session(&self) -> PySparkSession { - PySparkSession::new(self.session.new_session()) + fn new_session(&self, py: Python<'_>) -> PyResult> { + session_object(py, &self.session.new_session()) } /// Clone this session on the server, as PySpark's `cloneSession(new_session_id=None)`: a @@ -2447,7 +2522,7 @@ impl PySparkSession { &self, py: Python<'_>, new_session_id: Option, - ) -> PyResult { + ) -> PyResult> { let session = py .detach(|| { self.session @@ -2455,7 +2530,7 @@ impl PySparkSession { .map_err(Box::new) }) .to_pyerr()?; - Ok(PySparkSession::new(session)) + session_object(py, &session) } } diff --git a/crates/pyspark-rs/src/streaming.rs b/crates/pyspark-rs/src/streaming.rs index 58a3c9c4..8376be3b 100644 --- a/crates/pyspark-rs/src/streaming.rs +++ b/crates/pyspark-rs/src/streaming.rs @@ -1245,7 +1245,7 @@ impl PyStreamingQueryManager { /// listener added twice is in the bus twice. Should the server fail to add its listener, this /// warns, as PySpark Connect's does. fn addListener(&self, py: Python<'_>, listener: Py) -> PyResult<()> { - let session = crate::session::PySparkSession::new(self.session.clone()); + let session = crate::session::session_object(py, &self.session)?; listener .bind(py) .call_method1("_set_spark_session", (session,))?; diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index caf94421..fab61e1c 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -391,6 +391,12 @@ impl SparkSession { &self.client } + /// Whether `other` is this session or a clone of it, which shares its client. Two sessions + /// made apart are not the same session, even of the same session ID. + pub fn is_same_session(&self, other: &SparkSession) -> bool { + Arc::ptr_eq(&self.client, &other.client) + } + /// Create a builder for a new SparkSession. pub fn builder() -> SparkSessionBuilder { SparkSessionBuilder::new() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index abdebcb8..fe38d1b4 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -1029,6 +1029,61 @@ def test_session_builders_set_their_options_and_validate_their_urls_as_pyspark(m +def test_dataframes_give_back_their_session_object_as_pyspark(offline_spark): + # As PySpark: the DataFrames of a session, its streams and its client give back the session + # object itself, and two sessions do not mix. The registry holds the session objects + # weakly, so a session that is dropped is freed. + import gc + import weakref + + def check(actual, expected): + # Not an assert: pytest reprs a DataFrame in a failed assert, which asks the server. + if actual is not expected: + raise AssertionError("not the session object that PySpark gives back") + + spark = offline_spark + df = spark.range(2) + check(df.sparkSession, spark) + check(df._session, spark) + check(df.select("id").sparkSession, spark) + check(df.groupBy("id").count().sparkSession, spark) + check(df.na.drop().sparkSession, spark) + check(spark.read.format("rate").load().sparkSession, spark) + check(spark.udf.sparkSession, spark) + check(spark.udtf.sparkSession, spark) + check(spark.dataSource.sparkSession, spark) + assert df.sparkSession == spark + check(df.sparkSession.streams, spark.streams) + check(df.sparkSession.client, spark.client) + other = spark.newSession() + check(other.range(1).sparkSession, other) + if other.range(1).sparkSession is spark: + raise AssertionError("two sessions mixed") + check(other.range(1).sparkSession.streams, other.streams) + if other.streams is spark.streams: + raise AssertionError("two sessions share their streams") + ref = weakref.ref(other) + del other + gc.collect() + assert ref() is None + check(df.sparkSession, spark) + # Two sessions made apart, of the same session ID, are two objects, each that of its own + # DataFrames. create() sets the default session, so this runs in a process of its own. + import subprocess + import sys + + subprocess.run([sys.executable, "-c", """ +import uuid +from pyspark.sql import SparkSession +url = f"sc://localhost:1/;session_id={uuid.uuid4()}" +first = SparkSession.builder.remote(url).create() +second = SparkSession.builder.remote(url).create() +assert first is not second +assert first.range(1).sparkSession is first +assert second.range(1).sparkSession is second +"""], check=True) + + def test_session_builder_class_constructor(): from pyspark.sql import SparkSession b = SparkSession.Builder()