[SPARK-61166][CONNECT][PYTHON] Make the threads of inheritable_thread_target share the copy of the tags of the session, in the Rust and Python clients, as PySpark does

crates/pyspark-rs/src/session.rs

@@ -548,6 +548,11 @@ impl PyHook {
 }
 
 /// Python wrapper for a Spark session.
+/// A copy of the tags of a thread, which the threads that `_share_thread_tags` gives it share,
+/// for `pyspark.util.inheritable_thread_target`.
+#[pyclass(name = "_SharedTags", module = "pyspark.sql.session", frozen)]
+pub struct PySharedTags(spark_connect::session::SharedTags);
+
 /// A tag of `addTag` or `removeTag`: `CANNOT_BE_NONE` for None, as PySpark Connect's.
 fn tag_of(py: Python<'_>, tag: &Bound<'_, PyAny>) -> PyResult<String> {
     if tag.is_none() {
@@ -2340,6 +2345,20 @@ impl PySparkSession {
         self.session.clear_tags()
     }
 
+    /// A copy of the tags of the current thread, which `pyspark.util.inheritable_thread_target`
+    /// makes when it decorates a function, as PySpark Connect's copies the `thread_local` of the
+    /// client.
+    fn _copy_thread_tags(&self) -> PySharedTags {
+        PySharedTags(self.session.copy_thread_tags())
+    }
+
+    /// Make `tags` those of the current thread, which shares them with the other threads given
+    /// them, as `pyspark.util.inheritable_thread_target` gives every thread that runs its
+    /// function the copy of the `thread_local` of the client of PySpark Connect.
+    fn _share_thread_tags(&self, tags: &PySharedTags) {
+        self.session.share_thread_tags(&tags.0)
+    }
+
     /// Add local files, or `file:` URIs, as artifacts of the session, as PySpark's
     /// `addArtifacts(*path, pyfile=False, archive=False, file=False)`: each path is named by
     /// the flags and its extension, and the errors for a path that is not supported are

crates/spark-connect/src/session.rs

@@ -41,32 +41,72 @@ pub struct ExecutionInfo {
 
 /// The operation tags of a session, per thread, as PySpark Connect keeps them in the
 /// `thread_local` of its client: those that a thread added, which the operations that it
-/// starts carry. A new thread has none.
+/// starts carry. A new thread has none. The tags of a thread are a set that it may share with
+/// other threads ([`SharedTags`]), as the threads of PySpark's `inheritable_thread_target` share
+/// the set of tags that it copied.
 #[derive(Default)]
 pub(crate) struct ThreadTags {
-    by_thread: Mutex<std::collections::HashMap<std::thread::ThreadId, Vec<String>>>,
+    by_thread: Mutex<std::collections::HashMap<std::thread::ThreadId, SharedTags>>,
 }
 
+/// A set of operation tags of threads: a copy of the tags of a thread, of
+/// [`SparkSession::copy_thread_tags`], which the threads given it by
+/// [`SparkSession::share_thread_tags`] share, as the threads of PySpark's
+/// `inheritable_thread_target` share the copy of the `thread_local` of the client of the thread
+/// that decorated their function. The tags that one of them adds or removes are those of all.
+#[derive(Clone, Default)]
+pub struct SharedTags(Arc<Mutex<Vec<String>>>);
+
 impl ThreadTags {
     /// The tags of the current thread.
     fn current(&self) -> Vec<String> {
-        let by_thread = self.by_thread.lock().unwrap_or_else(|e| e.into_inner());
-        by_thread
+        let shared = self
+            .by_thread
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
             .get(&std::thread::current().id())
-            .cloned()
+            .cloned();
+        shared
+            .map(|tags| tags.0.lock().unwrap_or_else(|e| e.into_inner()).clone())
             .unwrap_or_default()
     }
 
-    /// Change the tags of the current thread, and forget a thread without tags.
+    /// Change the tags of the current thread, and forget a thread without tags that it shares
+    /// with no other.
     fn update(&self, change: impl FnOnce(&mut Vec<String>)) {
         let mut by_thread = self.by_thread.lock().unwrap_or_else(|e| e.into_inner());
         let thread = std::thread::current().id();
-        let tags = by_thread.entry(thread).or_default();
-        change(tags);
-        if tags.is_empty() {
+        let shared = by_thread.entry(thread).or_default().clone();
+        let mut tags = shared.0.lock().unwrap_or_else(|e| e.into_inner());
+        change(&mut tags);
+        // The set is held by the map and here only: no other thread or copy shares it.
+        if tags.is_empty() && Arc::strong_count(&shared.0) == 2 {
+            drop(tags);
             by_thread.remove(&thread);
         }
     }
+
+    /// A new empty set of tags of the current thread, which shares no tags with another, as
+    /// PySpark Connect's `clearTags` sets the tags of the `thread_local` to a new set.
+    fn clear(&self) {
+        self.by_thread
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .remove(&std::thread::current().id());
+    }
+
+    /// A copy of the tags of the current thread.
+    fn copy(&self) -> SharedTags {
+        SharedTags(Arc::new(Mutex::new(self.current())))
+    }
+
+    /// Make `tags` those of the current thread.
+    fn share(&self, tags: &SharedTags) {
+        self.by_thread
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .insert(std::thread::current().id(), tags.clone());
+    }
 }
 
 /// A Spark Connect session for interacting with a remote Spark cluster.
@@ -158,9 +198,25 @@ impl SparkSession {
         self.tags.current()
     }
 
-    /// Clear the tags that this thread set on this session. Mirrors `SparkSession.clearTags`.
+    /// Clear the tags that this thread set on this session. Mirrors `SparkSession.clearTags`:
+    /// the thread gets a new empty set of tags, which it shares with no other thread.
     pub fn clear_tags(&self) {
-        self.tags.update(Vec::clear);
+        self.tags.clear();
+    }
+
+    /// A copy of the tags of the current thread, as of now, which the threads that are given it
+    /// with [`SparkSession::share_thread_tags`] share. Mirrors how PySpark's
+    /// `inheritable_thread_target(session)` copies the tags of the thread that decorates a
+    /// function, which every thread that runs the function then has.
+    pub fn copy_thread_tags(&self) -> SharedTags {
+        self.tags.copy()
+    }
+
+    /// Make `tags` the tags of the current thread, which it shares with the other threads given
+    /// them: a tag that one of them adds or removes is added or removed for all, until a thread
+    /// clears its tags.
+    pub fn share_thread_tags(&self, tags: &SharedTags) {
+        self.tags.share(tags);
     }
 
     /// Register a handler of the progress of the executions of this session: it is called with
@@ -2599,6 +2655,44 @@ mod tests {
         assert!(spark.tags.by_thread.lock().unwrap().is_empty());
     }
 
+    #[test]
+    fn threads_given_a_copy_of_tags_share_it() {
+        // As the threads of PySpark's inheritable_thread_target: they share the copy of the tags
+        // of the thread that decorated their function, of then, until one clears its tags.
+        let spark = session();
+        spark.add_tag("a").unwrap();
+        let copy = spark.copy_thread_tags();
+        spark.add_tag("later").unwrap();
+        let run = |copy: SharedTags, tag: &'static str| {
+            let spark = spark.clone();
+            std::thread::spawn(move || {
+                spark.share_thread_tags(&copy);
+                spark.add_tag(tag).unwrap();
+                spark.get_tags()
+            })
+            .join()
+            .unwrap()
+        };
+        assert_eq!(run(copy.clone(), "b"), vec!["a", "b"]);
+        assert_eq!(run(copy.clone(), "c"), vec!["a", "b", "c"]);
+        assert_eq!(spark.get_tags(), vec!["a", "later"]);
+        let cleared = {
+            let spark = spark.clone();
+            let copy = copy.clone();
+            std::thread::spawn(move || {
+                spark.share_thread_tags(&copy);
+                spark.remove_tag("c");
+                spark.clear_tags();
+                spark.add_tag("d").unwrap();
+                spark.get_tags()
+            })
+            .join()
+            .unwrap()
+        };
+        assert_eq!(cleared, vec!["d"]);
+        assert_eq!(run(copy, "e"), vec!["a", "b", "e"]);
+    }
+
     #[test]
     fn arrow_local_relation_preserves_batches_and_rejects_schema_mismatch() {
         use arrow::array::Array;

python/pyspark/util.py

@@ -141,7 +141,9 @@ def inheritable_thread_target(f: Union[Callable, "SparkSession"]) -> Callable:
     >>> Thread(target=inheritable_thread_target(session)(target_func)).start()  # doctest: +SKIP
     """
     # As PySpark's of Spark Connect, the only mode of this client: `f` is the session, whose
-    # tags, those of this thread, the thread inherits, with the session as its active one.
+    # tags, those of this thread when the function is decorated, the threads of the function
+    # inherit, with the session as their active one. As PySpark's, which gives each thread the
+    # one copy of the `thread_local` of the client, the threads share that copy of the tags.
     from pyspark.sql import SparkSession
 
     session = f
@@ -150,7 +152,7 @@ def outer(ff: Callable) -> Callable:
         assert isinstance(session, SparkSession), (
             "f is expected to be SparkSession for spark connect"
         )
-        tags = set(session.getTags())
+        tags = session._copy_thread_tags()
 
         @functools.wraps(ff)
         def inner(*args: Any, **kwargs: Any) -> Any:
@@ -161,9 +163,7 @@ def inner(*args: Any, **kwargs: Any) -> Any:
 
             SparkSession._set_default_and_active_session(session)
             # Set the tags of the session in the child thread.
-            session.clearTags()
-            for tag in tags:
-                session.addTag(tag)
+            session._share_thread_tags(tags)
             return ff(*args, **kwargs)
 
         return inner

python/tests/test_dropin_offline.py

@@ -5655,6 +5655,45 @@ def child(key):
         inheritable_thread_target(child)(child)
 
 
+def test_threads_of_inheritable_thread_targets_share_their_tags_as_pyspark(mock_connect):
+    # As PySpark Connect's inheritable_thread_target, as Apache Spark's
+    # test_inheritable_tags_with_deco: the threads of the function share the copy of the tags of
+    # the thread that decorated it, of then, so a thread that the function starts sees the tags
+    # that the thread that started it added, while the decorating thread keeps its own.
+    import threading
+
+    from pyspark import inheritable_thread_target
+
+    spark, _ = mock_connect
+    spark.addTag("a")
+    try:
+        @inheritable_thread_target(spark)
+        def func(target):
+            return target()
+
+        spark.addTag("later")
+        first, second = set(), set()
+
+        def inner():
+            spark.addTag("c")
+            second.update(spark.getTags())
+
+        def outer():
+            spark.addTag("b")
+            first.update(spark.getTags())
+            nested = threading.Thread(target=func, args=(inner,))
+            nested.start()
+            nested.join()
+
+        thread = threading.Thread(target=func, args=(outer,))
+        thread.start()
+        thread.join()
+        assert spark.getTags() == {"a", "later"}
+        assert first == {"a", "b"} and second == {"a", "b", "c"}
+    finally:
+        spark.clearTags()
+
+
 def test_row_asdict():
     """Test Row.asDict() method."""
     from pyspark.sql import Row