This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 8c22eff93916 [SPARK-58358][SQL] Add tag validation to removeTag
8c22eff93916 is described below
commit 8c22eff939162e7e3c5079da0abe4558c747f2bb
Author: Jubin Soni <[email protected]>
AuthorDate: Fri Aug 7 22:18:34 2026 +0800
[SPARK-58358][SQL] Add tag validation to removeTag
### What is the purpose of the change
Fixes SPARK-58358 - aligns `SparkSession.removeTag` with its documented
contract across both Spark Classic and Spark Connect. Two related gaps are
addressed:
1. **Classic did not reject invalid tags.** `removeTag` went straight to
the tag map (`managedJobTags.get().remove(tag)`) and skipped validation, so
`removeTag("")` and `removeTag("a,b")` silently succeeded. This contradicts the
method's own docstring ("Cannot contain ',' (comma) character or be an empty
string"), diverges from Spark Connect (which validates both `addTag` and
`removeTag`), and is inconsistent with the sibling `addTag` and with
`SparkContext.removeJobTags`, all of whic [...]
2. **Connect did not honor the no-op contract for absent tags.** The
Connect client's `remove_tag` used `set.remove`, which raises `KeyError` when
the tag is absent, whereas `SparkSession.removeTag` documents removal of an
unknown tag as a no-op (and Classic behaves that way).
Both are surfaced and locked down by new tests.
### Brief change log
-
`sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala`:
`removeTag` now calls `SparkContext.throwIfInvalidTag(tag)` before removing,
mirroring `addTag`.
- `python/pyspark/sql/connect/client/core.py`: Connect `remove_tag` uses
`set.discard` instead of `set.remove`, so removing an absent tag is a no-op,
matching Classic and the documented contract.
-
`sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala`:
added a test asserting `addTag` and `removeTag` reject empty /
comma-containing tags, that rejected tags do not leak into the tag set, and
that removing a valid absent tag stays a no-op.
- `python/pyspark/sql/tests/test_job_cancellation.py`: added
`test_invalid_tags` covering the same contract in both Spark Classic and Spark
Connect modes (via the shared mixin) - invalid tags raise the mode-specific
type (`IllegalArgumentException` in Classic, `PySparkValueError` in Connect),
and removing a valid absent tag is a no-op.
### Verifying this change
This change adds two correctness fixes plus test coverage.
- New Scala test `SPARK-58358: addTag and removeTag reject invalid tags`
passes with the fix and fails without it (reverting the one-line Classic fix
produces "Expected exception java.lang.IllegalArgumentException to be thrown,
but noexception was thrown" at the `removeTag` assertion), confirming the test
exercises
the fixed behavior.
- New PySpark test `test_invalid_tags` runs in both Classic and Connect
modes; its
absent-tag no-op assertion fails against the previous Connect `set.remove`
(`KeyError`) and passes with `set.discard`.
- Existing tag/cancellation tests in the affected suites continue to pass.
Run:
- build/sbt 'sql/testOnly
org.apache.spark.sql.SparkSessionJobTaggingAndCancellationSuite'
- python/run-tests --testnames pyspark.sql.tests.test_job_cancellation
- python/run-tests --testnames
pyspark.sql.tests.connect.test_parity_job_cancellation
### Does this pull request potentially affect one of the following parts
- Dependencies (does it add or upgrade a dependency): no
- The public API, i.e., is any changed class annotated with
Public/Evolving: yes - `SparkSession.removeTag` now throws
`IllegalArgumentException` on invalid tags in Classic where it previously
ignored them, and the Spark Connect client no longer raises `KeyError` when
removing an absent tag; both align the API with its documented contract.
- The serializers: no
- The runtime per-record code paths (performance sensitive): no
- Anything that affects deployment or recovery: no
- The S3 file system connector: no
### Documentation
Does this pull request introduce a new feature? No - this is a bug fix that
enforces existing documented behaviour.
### Was generative AI tooling used to co-author this PR?
- [x] Yes - Claude Code was used as a pair-programming assistant. All
changes were reviewed and verified by the author. Generated-by: Claude Opus 4.8
Closes #57550 from jubins/j-SPARK-58358-validate-invalid-tags.
Authored-by: Jubin Soni <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
python/pyspark/sql/connect/client/core.py | 4 +++-
python/pyspark/sql/tests/test_job_cancellation.py | 22 ++++++++++++++++++++++
.../apache/spark/sql/classic/SparkSession.scala | 5 ++++-
...parkSessionJobTaggingAndCancellationSuite.scala | 21 +++++++++++++++++++++
4 files changed, 50 insertions(+), 2 deletions(-)
diff --git a/python/pyspark/sql/connect/client/core.py
b/python/pyspark/sql/connect/client/core.py
index 5c1dc029f826..0ff6f0ae90f9 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -2237,7 +2237,9 @@ class SparkConnectClient(object):
self._throw_if_invalid_tag(tag)
if not hasattr(self.thread_local, "tags"):
self.thread_local.tags = set()
- self.thread_local.tags.remove(tag)
+ # Use discard, not remove: removing an absent tag is a documented no-op
+ # (see SparkSession.removeTag), matching the Classic behavior.
+ self.thread_local.tags.discard(tag)
def get_tags(self) -> Set[str]:
if not hasattr(self.thread_local, "tags"):
diff --git a/python/pyspark/sql/tests/test_job_cancellation.py
b/python/pyspark/sql/tests/test_job_cancellation.py
index c5be232afd3b..0bd9b267f87a 100644
--- a/python/pyspark/sql/tests/test_job_cancellation.py
+++ b/python/pyspark/sql/tests/test_job_cancellation.py
@@ -20,6 +20,7 @@ import threading
import time
from pyspark import InheritableThread, inheritable_thread_target
+from pyspark.errors import IllegalArgumentException, PySparkValueError
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import eventually
@@ -37,6 +38,27 @@ class JobCancellationTestsMixin:
self.assertEqual(self.spark.getTags(), set())
self.spark.clearTags()
+ def test_invalid_tags(self):
+ # A tag cannot be an empty string or contain the ',' separator
(documented on
+ # SparkSession.addTag / removeTag). Both the classic and Spark Connect
paths reject
+ # such tags. The two paths raise different types -- a JVM-backed
+ # IllegalArgumentException in Classic and a PySparkValueError in
Connect -- so accept
+ # either while still rejecting unrelated session/transport failures.
+ invalid_tag_errors = (IllegalArgumentException, PySparkValueError)
+ self.spark.clearTags()
+ for invalid_tag in ["", "a,b", ","]:
+ with self.assertRaises(invalid_tag_errors):
+ self.spark.addTag(invalid_tag)
+ with self.assertRaises(invalid_tag_errors):
+ self.spark.removeTag(invalid_tag)
+ # A rejected tag must not have been recorded.
+ self.assertEqual(self.spark.getTags(), set())
+
+ # Removing a valid but absent tag is a no-op in both modes (does not
raise).
+ self.spark.removeTag("absent")
+ self.assertEqual(self.spark.getTags(), set())
+ self.spark.clearTags()
+
def test_tags_multithread(self):
output1 = None
output2 = None
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
index c5f01de37000..ffbfed21bb29 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
@@ -763,7 +763,10 @@ class SparkSession private(
}
/** @inheritdoc */
- override def removeTag(tag: String): Unit = managedJobTags.get().remove(tag)
+ override def removeTag(tag: String): Unit = {
+ SparkContext.throwIfInvalidTag(tag)
+ managedJobTags.get().remove(tag)
+ }
/** @inheritdoc */
override def getTags(): Set[String] = managedJobTags.get().keySet.toSet
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala
index d7b2511eac2a..85de26b7a567 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala
@@ -81,6 +81,27 @@ class SparkSessionJobTaggingAndCancellationSuite
assert(session.getTags() == Set("one"))
}
+ test("SPARK-58358: addTag and removeTag reject invalid tags") {
+ val session = classic.SparkSession.builder().master("local").getOrCreate()
+
+ // A tag cannot be empty or contain the ',' separator. Both addTag and
removeTag must
+ // reject such tags (removeTag previously skipped this validation).
+ Seq("", "a,b", ",").foreach { invalidTag =>
+ intercept[IllegalArgumentException] {
+ session.addTag(invalidTag)
+ }
+ intercept[IllegalArgumentException] {
+ session.removeTag(invalidTag)
+ }
+ }
+ // A rejected tag must not leak into the tag set.
+ assert(session.getTags() == Set())
+
+ // Removing a valid, absent tag stays a no-op.
+ session.removeTag("absent")
+ assert(session.getTags() == Set())
+ }
+
test("Tags set from session are prefixed with session UUID") {
sc = new SparkContext("local[2]", "test")
val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
