[SPARK-61097][CONNECT][PYTHON] Cancel listener work during session shutdown

crates/pyspark-rs/src/functions.rs

@@ -333,9 +333,11 @@ fn numpy_ndarray_ndim(obj: &Bound<'_, PyAny>) -> PyResult<Option<usize>> {
 /// The literal of a NumPy scalar, typed as PySpark's `LiteralExpression._infer_type` types it: a
 /// bool, a signed int of 8 to 64 bits (tinyint to bigint), or a float32 or float64. A str or bytes
 /// scalar is `None`, for the string and binary conversions to take. Any other NumPy scalar (an
-/// unsigned int, float16, longdouble, complex, datetime64 or timedelta64) is
-/// UNSUPPORTED_DATA_TYPE, named by its Python type. NumPy is looked up in `sys.modules`, as a
-/// NumPy scalar exists only when NumPy is imported.
+/// unsigned int, float16, complex, datetime64 or timedelta64) is UNSUPPORTED_DATA_TYPE,
+/// named by its Python type. A supported dtype with a scalar class outside PySpark's accepted
+/// classes raises a bare `AssertionError`, such as a distinct `longlong`, or an eight-byte
+/// `longdouble` on macOS arm64 and Windows. Wider `longdouble` dtypes are unsupported.
+/// NumPy is looked up in `sys.modules`, as a NumPy scalar exists only when NumPy is imported.
 fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult<Option<LiteralExpression>> {
     let py = obj.py();
     let modules = py.import("sys")?.getattr("modules")?;
@@ -352,6 +354,36 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult<Option<LiteralExpres
     let dtype = obj.getattr("dtype")?;
     let kind: String = dtype.getattr("kind")?.extract()?;
     let itemsize: usize = dtype.getattr("itemsize")?.extract()?;
+    // Mirror LiteralExpression's accepted scalar classes after dtype inference. Aliases can
+    // share a supported dtype without being instances of its canonical NumPy scalar class.
+    // Check the matching width first; the remaining classes preserve the reference's full set.
+    let scalar_classes: &[&str] = match (kind.as_str(), itemsize) {
+        ("i", 1) => &["int8"],
+        ("i", 2) => &["int16", "int8"],
+        ("i", 4) => &["int32", "int16", "int8"],
+        ("i", 8) => &["int64", "int32", "int16", "int8"],
+        ("f", 4) => &["float32"],
+        ("f", 8) => &["float64", "float32"],
+        _ => &[],
+    };
+    if !scalar_classes.is_empty() {
+        let mut accepted = match kind.as_str() {
+            "i" => obj.is_instance_of::<pyo3::types::PyInt>(),
+            "f" => obj.is_instance_of::<pyo3::types::PyFloat>(),
+            _ => false,
+        };
+        if !accepted {
+            for class in scalar_classes {
+                if obj.is_instance(&numpy.getattr(*class)?)? {
+                    accepted = true;
+                    break;
+                }
+            }
+        }
+        if !accepted {
+            return Err(pyo3::exceptions::PyAssertionError::new_err(()));
+        }
+    }
     let literal = match (kind.as_str(), itemsize) {
         ("U", _) | ("S", _) => return Ok(None),
         ("i", 1) => LiteralExpression::Byte(obj.extract::<i64>()? as i32),

crates/pyspark-rs/src/session.rs

@@ -1799,11 +1799,12 @@ impl PySparkSession {
 
     /// Stop this Spark session. Best-effort and idempotent, mirroring pyspark: a failure
     /// to release the session on the server is logged, not raised. Also clears the
-    /// active session and the cached `SQLContext` if they belong to this session.
+    /// active session and the cached `SQLContext` if they belong to this session. Listener
+    /// callbacks are not joined; the current event's listener snapshot may finish afterward.
     fn stop(&self, py: Python<'_>) -> PyResult<()> {
         crate::progress::forget_session(&self.session);
         crate::profiler::forget_session(&self.session);
-        if let Err(e) = py.detach(|| self.session.stop()) {
+        if let Err(e) = py.detach(|| self.session.stop().map_err(Box::new)) {
             let logger = py
                 .import("logging")?
                 .call_method1("getLogger", ("pyspark",))?;
@@ -2075,8 +2076,9 @@ impl PySparkSession {
     }
 
     /// Close the session without releasing it on the server, for `client.close`.
-    fn _close(&self) {
-        self.session.close();
+    /// Do not join listener callbacks; the current event's listener snapshot may finish afterward.
+    fn _close(&self, py: Python<'_>) {
+        py.detach(|| self.session.close());
     }
 
     /// The client of the session, as PySpark Connect's `SparkSession._client`, which the

crates/pyspark-rs/src/streaming.rs

@@ -1244,6 +1244,8 @@ impl PyStreamingQueryManager {
     /// it. The id of each registration is appended to the listener's `_rust_listener_ids`, as a
     /// listener added twice is in the bus twice. Should the server fail to add its listener, this
     /// warns, as PySpark Connect's does.
+    /// A terminal event-stream end or error clears the registrations; add listeners again
+    /// to open a new subscription.
     fn addListener(&self, py: Python<'_>, listener: Py<PyAny>) -> PyResult<()> {
         let session = crate::session::PySparkSession::new(self.session.clone());
         listener
@@ -1277,6 +1279,9 @@ impl PyStreamingQueryManager {
     /// the bus of the session, the first one that the bus has (as the bus removes the first
     /// occurrence), and the others, which are of other buses, stay. Should the server fail to
     /// remove its listener, this warns, and the listener stays, as PySpark Connect's does.
+    /// A callback coalesced with another caller's in-progress last-listener removal or close
+    /// keeps its ids so a failed operation can be retried. Only a completed removal consumes
+    /// one matching id from this session's bus.
     fn removeListener(&self, py: Python<'_>, listener: Py<PyAny>) -> PyResult<()> {
         let bound = listener.bind(py);
         let Ok(ids) = bound.getattr("_rust_listener_ids") else {
@@ -1319,9 +1324,21 @@ impl PyStreamingQueryManager {
     }
 
     /// Remove all client-side listeners and stop the dispatch thread. Mirrors
-    /// `StreamingQueryManager.close`.
+    /// `StreamingQueryManager.close`. A server removal failure warns and keeps the
+    /// listeners registered, as PySpark Connect's does, so close can be retried.
     fn close(&self, py: Python<'_>) -> PyResult<()> {
-        crate::dataframe::detached(py, || self.inner.close().map_err(Box::new))
+        if let Err(error) = crate::dataframe::detached(py, || self.inner.close().map_err(Box::new))
+        {
+            return warn(
+                py,
+                format!(
+                    "Failed to remove the listener because of exception: {}\n\
+                     The listener is not removed, please remove it again.",
+                    error.value(py)
+                ),
+            );
+        }
+        Ok(())
     }
 
     fn streamListenerEvents(&self, py: Python<'_>) -> PyResult<PyListenerEventStream> {

crates/spark-connect/src/progress.rs

@@ -2,11 +2,14 @@
 //! `pyspark.sql.connect.shell.progress`: the progress handlers of the session are called with
 //! each progress of the server that has tasks, and once more, as done, when the execution ends.
 
+use std::sync::Arc;
+
 use spark_connect_core::client::ReattachableResponseStream;
 use spark_connect_core::error::Result;
+use spark_connect_core::runtime::CallbackGuard;
 use spark_connect_proto as proto;
 
-use crate::session::SparkSession;
+use crate::session::{ProgressHandler, ProgressHandlers, SparkSession};
 
 /// The progress of a stage of an execution, as PySpark's `StageInfo`.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -44,7 +47,7 @@ pub struct ProgressUpdate {
 
 /// The progress of one execution, as PySpark's `Progress` of the execution of a request.
 pub(crate) struct ProgressTracker {
-    session: SparkSession,
+    handlers: ProgressHandlers,
     operation_id: Option<String>,
     stages: Vec<StageInfo>,
     inflight_tasks: i64,
@@ -53,7 +56,7 @@ pub(crate) struct ProgressTracker {
 impl ProgressTracker {
     pub(crate) fn new(session: &SparkSession, operation_id: Option<String>) -> Self {
         ProgressTracker {
-            session: session.clone(),
+            handlers: session.progress_handlers(),
             operation_id,
             stages: Vec::new(),
             inflight_tasks: 0,
@@ -80,12 +83,29 @@ impl ProgressTracker {
     }
 
     fn notify(&self, done: bool) {
-        self.session.notify_progress(|| ProgressUpdate {
+        // Clone handlers out of the registry lock so callbacks can add or remove handlers.
+        let handlers: Vec<ProgressHandler> = self
+            .handlers
+            .lock()
+            .unwrap()
+            .iter()
+            .map(|(_, handler)| Arc::clone(handler))
+            .collect();
+        if handlers.is_empty() {
+            return;
+        }
+        let update = ProgressUpdate {
             stages: self.stages.clone(),
             inflight_tasks: self.inflight_tasks,
             operation_id: self.operation_id.clone(),
             done,
-        });
+        };
+        // A handler may run a query whose progress calls the handlers again. Preserve the
+        // execution callback-depth guard without retaining the session or its listener bus.
+        let _callback = CallbackGuard::enter();
+        for handler in handlers {
+            handler(&update);
+        }
     }
 }
 
@@ -124,6 +144,7 @@ impl ExecutionStream {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
     use std::sync::{Arc, Mutex};
 
     use super::*;
@@ -186,4 +207,50 @@ mod tests {
         );
         session.remove_progress_handler(id);
     }
+
+    #[test]
+    fn detached_progress_registry_preserves_the_execution_recursion_bound() {
+        use spark_connect_core::error::SparkErrorKind;
+        use spark_connect_core::runtime::{check_callback_depth, MAX_CALLBACK_DEPTH};
+
+        let session = SparkSession::builder()
+            .remote("sc://localhost:15002")
+            .get_or_create()
+            .expect("session");
+        let handlers = Arc::downgrade(&session.progress_handlers());
+        let calls = Arc::new(AtomicUsize::new(0));
+        let failures = Arc::new(AtomicUsize::new(0));
+        let seen_calls = Arc::clone(&calls);
+        let seen_failures = Arc::clone(&failures);
+        session.register_progress_handler(move |update| {
+            if update.done {
+                return;
+            }
+            let call = seen_calls.fetch_add(1, Ordering::SeqCst);
+            match check_callback_depth() {
+                Err(error) => {
+                    assert_eq!(error.kind, SparkErrorKind::RecursionError);
+                    seen_failures.fetch_add(1, Ordering::SeqCst);
+                }
+                Ok(()) => {
+                    // Bound this test even if the guard regresses; no network execution or
+                    // unbounded native recursion is needed to observe the missing guard.
+                    assert!(call < MAX_CALLBACK_DEPTH);
+                    let nested = ProgressTracker {
+                        handlers: handlers.upgrade().expect("registry is alive"),
+                        operation_id: Some("nested".to_string()),
+                        stages: Vec::new(),
+                        inflight_tasks: 0,
+                    };
+                    nested.notify(false);
+                }
+            }
+        });
+        let tracker = ProgressTracker::new(&session, None);
+        tracker.notify(false);
+        assert_eq!(calls.load(Ordering::SeqCst), MAX_CALLBACK_DEPTH);
+        assert_eq!(failures.load(Ordering::SeqCst), 1);
+        assert!(check_callback_depth().is_ok());
+        session.clear_progress_handlers();
+    }
 }

crates/spark-connect/src/session.rs

@@ -8,7 +8,7 @@ use std::sync::{Arc, Mutex};
 use spark_connect_core::channel::ChannelBuilder;
 use spark_connect_core::client::{ExecutePlanHook, SparkConnectClient};
 use spark_connect_core::error::{Result, SparkError};
-use spark_connect_core::runtime::{block_on, check_callback_depth, CallbackGuard};
+use spark_connect_core::runtime::{block_on, check_callback_depth};
 use spark_connect_core::thread_entries::forget_when_thread_ends;
 use spark_connect_proto as proto;
 
@@ -24,6 +24,9 @@ use crate::types::DataType;
 /// execution ends. Mirrors `SparkSession.registerProgressHandler`.
 pub type ProgressHandler = Arc<dyn Fn(&ProgressUpdate) + Send + Sync>;
 
+/// The shared registry of progress handlers and their ids, without retaining a session.
+pub(crate) type ProgressHandlers = Arc<Mutex<Vec<(u64, ProgressHandler)>>>;
+
 /// Metrics captured from the last executed query. Mirrors the data behind
 /// `pyspark.sql.DataFrame.executionInfo` / `client.core.ExecutionInfo`.
 #[derive(Debug, Clone, Default)]
@@ -140,7 +143,7 @@ pub struct SparkSession {
     last_execution: Arc<Mutex<Option<ExecutionInfo>>>,
     /// Registered progress handlers, keyed by id (for removal). Invoked with each progress of
     /// an execution, and when it ends. Mirrors `registerProgressHandler`.
-    progress_handlers: Arc<Mutex<Vec<(u64, ProgressHandler)>>>,
+    progress_handlers: ProgressHandlers,
     /// Monotonic id source for `register_progress_handler`.
     progress_handler_id: Arc<AtomicU64>,
     /// The handler of the updates of the Python accumulators of its executions.
@@ -268,29 +271,9 @@ impl SparkSession {
         self.progress_handlers.lock().unwrap().clear();
     }
 
-    /// Invoke all registered progress handlers with the progress of an execution, which
-    /// `update` builds only when there is a handler. Called by the progress of executions; not
-    /// part of the public API.
-    pub(crate) fn notify_progress(&self, update: impl FnOnce() -> ProgressUpdate) {
-        // Clone the handler `Arc`s out of the lock so a handler can't deadlock by
-        // touching the session's handler list.
-        let handlers: Vec<ProgressHandler> = {
-            self.progress_handlers
-                .lock()
-                .unwrap()
-                .iter()
-                .map(|(_, h)| Arc::clone(h))
-                .collect()
-        };
-        if handlers.is_empty() {
-            return;
-        }
-        let update = update();
-        // A handler may run a query of its own, whose progress calls the handlers again.
-        let _callback = CallbackGuard::enter();
-        for h in handlers {
-            h(&update);
-        }
+    /// Share the progress-handler registry without retaining the session or its listener bus.
+    pub(crate) fn progress_handlers(&self) -> ProgressHandlers {
+        Arc::clone(&self.progress_handlers)
     }
 
     /// Execute `request` on the server, with the progress of the execution for the progress
@@ -1141,13 +1124,20 @@ impl SparkSession {
     ///
     /// Idempotent, mirroring pyspark: the session is marked stopped even if releasing it
     /// on the server fails, and stopping an already stopped session is a no-op.
+    ///
+    /// Cancel local listener registration/dispatch first, release the server session, then
+    /// close the client even if release fails. As PySpark's stop, this does not make a
+    /// separate listener-removal RPC: ReleaseSession releases the server's listeners too.
+    /// Teardown does not wait for listener callbacks; the current event's listener snapshot
+    /// may finish after stop returns, while no later event is dispatched.
     pub fn stop(&self) -> Result<()> {
         if self.stopped.swap(true, Ordering::SeqCst) {
             return Ok(());
         }
-        // As PySpark's `stop()`, the client is closed after the session is released, so that
-        // its requests raise `NO_ACTIVE_SESSION` from now on.
+        self.streams().cancel_listener_bus();
         let released = block_on(self.client.release_session());
+        // As PySpark's `stop()`, close the client even if server cleanup fails, so its
+        // future requests raise `NO_ACTIVE_SESSION`.
         self.client.close();
         released
     }
@@ -1159,9 +1149,12 @@ impl SparkSession {
     }
 
     /// Close this session without releasing it on the server, as `SparkConnectClient.close`
-    /// closes its channel: it is stopped, and `stop()` does nothing more.
+    /// closes its channel: it is stopped, and `stop()` does nothing more. Cancel local listener
+    /// registration/dispatch before closing the client. Callbacks are not joined; the current
+    /// event's listener snapshot may finish after close returns.
     pub fn close(&self) {
         self.stopped.store(true, Ordering::SeqCst);
+        self.streams().cancel_listener_bus();
         self.client.close();
     }