[SPARK-58447][CONNECT] Preserve operation tags in Python foreachBatch callbacks

[j1wonpark]

What changes were proposed in this pull request?

This PR propagates the operation tags from the writeStream.start() request to the Spark Connect session used by a Python foreachBatch callback.

The tags are transferred through the internal server-worker protocol when the Python worker starts and restored before processing micro-batches. A regression test verifies that the callback session preserves the tags and can execute a nested action with them.

Why are the changes needed?

Spark Connect executes Python foreachBatch callbacks in a separate Python worker that creates a local Connect client. Operation tags are client thread-local, so the new client does not contain the tags from the originating request.

As a result, operations executed inside the callback, such as actions, writes, and SQL commands, cannot be identified or interrupted using those tags.

Does this PR introduce any user-facing change?

Yes. Operations executed inside a Spark Connect Python foreachBatch callback now inherit the operation tags captured when the streaming query starts. This does not introduce a new API.

How was this patch tested?

• Added test_streaming_foreach_batch_operation_tags and verified that it fails before the fix and passes after the fix.
• python/run-tests --testnames pyspark.sql.tests.connect.streaming.test_parity_foreach_batch
• build/sbt -Phive package
• build/sbt connect/scalastyle
• SparkConnectSessionHolderSuite with the python foreachBatch process test filter

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-5.6)

[uros-b]

The new test is added to StreamingTestsForeachBatchMixin, which is inherited by both the Connect parity suite (StreamingForeachBatchParityTests) and the classic suite (StreamingTestsForeachBatch). In classic mode the foreachBatch callback runs on the Py4J callback server thread, a persistent pre-created thread. Classic SparkSession.addTag() sets tags in the JVM via InheritableThreadLocal; that Py4J thread was created before addTag() is called in the test body, so the tag will not be visible in the callback. The assertion assert tag in batch_df.sparkSession.getTags() will fail for the classic suite. The test exercises Connect-specific behaviour (the Python subprocess tag-injection mechanism) and belongs in python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py, not in the shared mixin.

[j1wonpark]

Although the fix is specific to the Spark Connect foreachBatch worker, the same test also passes in Classic through its existing session-cloning path. I kept it in the shared mixin to verify that both implementations preserve tags inside foreachBatch callbacks. Would you prefer moving it to the Connect-only suite to keep the regression test scoped to the changed implementation?

[HyukjinKwon]

0 blocking, 0 non-blocking, 0 nits.
Clean, protocol-symmetric tag propagation; well tested and backward compatible.

Verification

Confirmed the wire protocol is symmetric: the server writes writeInt(size) then a sorted writeUTF per tag; the worker reads read_int then loads each and addTags them onto the callback session before the micro-batch while-loop (set once per worker init). The new sessionTags parameter defaults to Set.empty so existing paths are unaffected, and the sorted write gives deterministic order. The regression test asserts the tag is visible via getTags() inside the callback.