[SPARK-44462][4.3][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect

[LuciferYang]

What changes were proposed in this pull request?

Spark Connect's Python foreachBatch cached the batch DataFrame in the root session's SessionHolder, but the batch DataFrame belongs to StreamExecution.sparkSessionForStream, a clone of the session that started the query. This PR gives that clone its own SessionHolder and routes per-batch work through it, so the batch runs under the configuration the streaming engine pins on the clone rather than the root session's. The behavior is gated behind a new SQL config spark.sql.connect.streaming.foreachBatch.useClonedSession (default true); setting it to false restores the previous root-session behavior.

On the server, StreamingForeachBatchHelper adds:

• ForeachBatchSessionManager. On the first batch it registers batchDf.sparkSession under a server-generated session id through the new SparkConnectSessionManager.registerExistingSession, with customInactiveTimeoutMs = -1 so it never expires by inactivity (the query owns its lifetime). The id is pinned for the whole query, so every batch resolves to the same holder. close() is terminal. When useClonedSession is false it registers no cloned holder and returns the root holder instead.
• ForeachBatchCleaner, replacing RunnerCleaner. It closes the Python runner (when there is one) and the cloned holder. CleanerCache.registerCleanerForQuery also closes eagerly if the query is already inactive at registration time, since onQueryTerminated never fires for such a query. It is registered before registerNewStreamingQuery so a throw there cannot strand the never-expiring cloned holder.
• A stale-DataFrame check keyed by streaming query id, so a batch DataFrame left in the cache by a previous batch is evicted.

The server writes (dfId, batchId, sessionId) per batch, where the session id is the cloned session's when the fix is on and the root session's when it is off. On the Python side, foreach_batch_worker bootstraps on the root session id to unpickle the user function before the first batch, then creates a session on the provided id and repoints SparkSession._default_session / _active_session at it, so SparkSession.active() inside the user function returns the session the batch runs on.

The Scala foreachBatch path receives the classic DataFrame directly and never resolves a remote batch relation, so it does not register a cloned SessionHolder. The cloned session id is logged under a CLONED_SESSION_ID key so the structured field matches the value.

Why are the changes needed?

StreamExecution runs the query on a clone of the session it was started with, and the engine forces onto that clone the configs stateful correctness depends on: spark.sql.cbo.enabled=false and spark.sql.requireAllClusterKeysForDistribution=false so stateful operators are not rearranged, AQE off for continuous, real-time, and stateful workloads, and the stateful shuffle partition count pinned at start. populateStartOffsets also restores the checkpointed conf from the offset log onto the same session. Connect resolved the batch DataFrame against the root SessionHolder, so operations the user function performed on it were planned and executed under the root session's configuration, with none of those settings in effect. The Python worker also kept the root session active, so the Connect session mapping never expired and the query kept running after the client disappeared.

Does this PR introduce any user-facing change?

Yes, for Python foreachBatch under Spark Connect, with spark.sql.connect.streaming.foreachBatch.useClonedSession at its default of true. The batch DataFrame is now bound to the cloned session, matching classic foreachBatch, which already runs the user function on the clone (ForeachBatchSink.addBatch builds the batch Dataset with ofRows(data.sparkSession, node)). Because the batch DataFrame no longer shares a session with the client:

• SparkSession.active() inside the function returns the batch session rather than the root one.
• A temporary view created from the batch DataFrame is not visible to the root session.
• Referencing the batch DataFrame from spark.sql on a captured root session fails with DATAFRAME_NOT_FOUND, and joining it with a DataFrame captured from the root session fails with SESSION_NOT_SAME.

Combining the batch DataFrame with the root session inside the function is not supported by design: the isolation that fixes the configuration bug is the same thing that separates the two sessions. Run such operations through the batch DataFrame's own session (batch_df.sparkSession), and pass results out of the function through a table or a global temporary view instead of a session-local temporary view. To keep the previous behavior, set spark.sql.connect.streaming.foreachBatch.useClonedSession to false: the batch DataFrame is then bound to the root session again and can be combined with it, at the cost of not running under the stream session's configuration. The Structured Streaming migration guide documents this.

How was this patch tested?

New tests in StreamingForeachBatchParityTests:

• test_batch_uses_stream_session_for_stateful_aqe pins that the batch runs on a different session than the root, with AQE off on the batch session while the root has it on.
• test_batch_dataframe_in_root_session_sql_is_unsupported and test_batch_dataframe_join_with_root_session_dataframe_is_unsupported pin the contract under the default: spark.sql referencing the batch DataFrame from the root session fails with DATAFRAME_NOT_FOUND, and joining it with a captured root-session DataFrame fails with SESSION_NOT_SAME.
• test_batch_dataframe_in_batch_session_sql_is_supported and test_batch_dataframe_join_in_batch_session_is_supported pin the documented migration: the same work run through batch_df.sparkSession succeeds.
• test_temp_view_is_isolated_from_root_session pins that a temp view created inside the function resolves through the cloned session but not the root session, with a sentinel table proving the function ran.
• test_disabling_fix_restores_root_session_behavior pins that with the config set to false, joining the batch DataFrame with a captured root-session DataFrame succeeds again.

New tests in StreamingForeachBatchHelperSuite: ForeachBatchSessionManager.close() is terminal and a closing root session is refused; a cleaner registered for an already terminated query is closed on the spot.

Updated tests: test_nested_dataframes now uses saveAsTable instead of createOrReplaceTempView, and SparkConnectSessionHolderSuite destructures the new 3-tuple from pythonForeachBatchWrapper.

Pass GitHub Actions.

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

Generated-by: Claude Code

[HeartSaVioR]

+1

[HeartSaVioR]

Thanks! Merging to 4.3.

[HeartSaVioR]

Merge Summary:

• merged into branch-4.3 fe00dca
  Posted by merge_spark_pr.py

[LuciferYang]

Thank you @HeartSaVioR and @HyukjinKwon