From b291657b9701ddd82f28da96a559d46e36e940f0 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Sun, 8 Nov 2026 14:13:12 +0900 Subject: [PATCH] [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 PySpark Connect's inheritable_thread_target(session) copies the thread_local of the client, of the tags, when it decorates a function, and every thread that runs the function gets that one copy, so the threads share it. The client copied the tags into each thread instead, so a thread that the function started from another thread of the function lacked the tags that that thread added. The tags of a thread in the core are now a set that threads can share: the session copies those of the current thread (copy_thread_tags) and gives the copy to a thread (share_thread_tags); clear_tags gives the thread a new set of its own, as PySpark's clearTags does. The decorator copies the tags when it decorates the function and gives the copy to each of its threads. Co-authored-by: Isaac --- crates/pyspark-rs/src/session.rs | 19 +++++ crates/spark-connect/src/session.rs | 116 +++++++++++++++++++++++++--- python/pyspark/util.py | 10 +-- python/tests/test_dropin_offline.py | 39 ++++++++++ 4 files changed, 168 insertions(+), 16 deletions(-) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index b19a8445..696ec42d 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/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 { 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 diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index cbbbf63f..952c45c3 100644 --- a/crates/spark-connect/src/session.rs +++ b/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>>, + by_thread: Mutex>, } +/// 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>>); + impl ThreadTags { /// The tags of the current thread. fn current(&self) -> Vec { - 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)) { 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(©); + 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(©); + 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; diff --git a/python/pyspark/util.py b/python/pyspark/util.py index 5708a518..bf8877e8 100644 --- a/python/pyspark/util.py +++ b/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 diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index c83fdb00..2d754766 100644 --- a/python/tests/test_dropin_offline.py +++ b/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