[SPARK-61183][CONNECT][PYTHON] Give back the SparkSession object of a DataFrame in DataFrame.sparkSession, in the Python client as PySpark does

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))

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<Py<crate::session::PySparkSession>> {
+        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<Py<crate::session::PySparkSession>> {
+        self.spark_session(py)
     }
 
     /// Not supported on Spark Connect, as PySpark Connect's: `NOT_IMPLEMENTED` of `rdd`, which

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<PySparkSession>,
 }
 
 impl PyDataSourceRegistration {
-    pub fn new(session: PySparkSession) -> Self {
+    pub fn new(session: Py<PySparkSession>) -> 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<PySparkSession> {
+        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(())

crates/pyspark-rs/src/session.rs

@@ -470,11 +470,83 @@ fn new_session(py: Python<'_>, url: &str) -> PyResult<Py<PySparkSession>> {
         .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<String, Vec<Py<PyAny>>>;
+
+/// The registry of the Python session objects.
+fn session_objects() -> &'static std::sync::Mutex<SessionObjects> {
+    static OBJECTS: std::sync::OnceLock<std::sync::Mutex<SessionObjects>> =
+        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<PyAny>) -> Option<Py<PySparkSession>> {
+    let object = weak.bind(py).call0().ok()?;
+    let session = object.cast::<PySparkSession>().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<Py<PySparkSession>> {
+    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<Py<PySparkSession>> {
+    {
+        let registry = session_objects().lock().unwrap_or_else(|e| e.into_inner());
+        if let Some(object) = registered_session_object(py, &registry, 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, &registry, 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<PySparkSession>) {
@@ -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<PyDataSourceRegistration> {
+        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<Py<PySparkSession>> {
+        session_object(py, &self.session.new_session())
     }
 
     /// Clone this session on the server, as PySpark's `cloneSession(new_session_id=None)`: a
@@ -2447,15 +2522,15 @@ impl PySparkSession {
         &self,
         py: Python<'_>,
         new_session_id: Option<String>,
-    ) -> PyResult<PySparkSession> {
+    ) -> PyResult<Py<PySparkSession>> {
         let session = py
             .detach(|| {
                 self.session
                     .clone_session_with_id(new_session_id.as_deref())
                     .map_err(Box::new)
             })
             .to_pyerr()?;
-        Ok(PySparkSession::new(session))
+        session_object(py, &session)
     }
 }
 

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<PyAny>) -> 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,))?;

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()

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()