[SPARK-58358][SQL] Add tag validation to removeTag

[jubins]

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 which call SparkContext.throwIfInvalidTag.

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?

• 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

[jubins]

hi @uros-b @cloud-fan can I please get review on this?

[cloud-fan]

1 blocking, 1 non-blocking, 0 nits.
The Classic fix is well placed, but the new cross-mode tests expose one unhandled Connect contract violation and use an assertion broad enough to mask unrelated failures.

Correctness (2)

• python/pyspark/sql/tests/test_job_cancellation.py:40: The shared Python coverage omits valid absent-tag removal, which currently violates the documented no-op contract in Connect. -- see inline
• python/pyspark/sql/tests/test_job_cancellation.py:49: The invalid-tag test accepts every exception instead of verifying the intended validation failures. -- see inline

Verification

I traced Classic removeTag through SparkContext.throwIfInvalidTag and compared it with both the JVM and Python Connect implementations. Invalid tags fail before mutation in both modes. Valid absent removal remains a no-op in Classic, while Python Connect calls set.remove and therefore raises KeyError, contrary to the shared public documentation.

PR description suggestions

• Update the brief change log: it says check_job_cancellation was changed to re-raise non-cancellation exceptions, but that helper is unchanged in the current diff.

[cloud-fan]

Please cover removal of a valid absent tag here, as the Scala test does. This mixin also runs against Connect, where remove_tag calls set.remove and raises KeyError despite SparkSession.removeTag documenting a no-op; changing it to discard would restore parity.

[jubins]

good catch, switched it to set.discard and added the valid-absent-tag case here (mirroring the Scala test), so the mixin now locks the no-op contract in both modes. thank you!

[cloud-fan]

Please assert the expected mode-specific types instead of Exception. A tuple of captured IllegalArgumentException and PySparkValueError keeps this shared while preventing unrelated session or transport failures from satisfying the test.

[jubins]

fixed, dropped the stale line about check_job_cancellation (unchanged in this diff) and updated the description to cover both fixes: Classic tag validation and the Connect absent-tag no-op.

[jubins]

@cloud-fan thanks for the thorough review, both points addressed.

[cloud-fan]

2 addressed, 0 remaining, 0 new.
0 blocking, 0 non-blocking, 0 nits.
The two issues from the prior review are resolved, and the current implementation and tests match the documented cross-mode contract.

Verification

I traced the public SparkSession.removeTag documentation to both implementations. Classic calls SparkContext.throwIfInvalidTag before map removal; Connect calls its validation helper before set.discard. The updated shared Python test asserts the mode-specific validation exception types and directly exercises valid absent-tag removal, while the Scala test covers the Classic path.

[cloud-fan]

the ci error is unrelated, thanks, merging to master/4.x/4.3

[cloud-fan]

Merge Summary:

• merged into master 8c22eff
• merged into branch-4.x ac1849c
• merged into branch-4.3 c9eee57

Posted by merge_spark_pr.py