[SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect by LuciferYang · Pull Request #55410 · apache/spark

LuciferYang · 2026-04-19T05:49:24Z

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

Add a `dataFrameQueryIndex` map to `SessionHolder` that tracks the active cached DataFrame ID per streaming query. Before caching a new batch DataFrame in `dataFrameCachingWrapper`, we check for any stale entry from a previous batch and remove it, logging a warning. This addresses TODO 1 in StreamingForeachBatchHelper. Changes: - `SessionHolder`: Add `dataFrameQueryIndex` ConcurrentMap - `StreamingForeachBatchHelper`: Accept `queryIdRef` in `dataFrameCachingWrapper`, perform sanity check, clean up mapping in finally block. Update `scalaForeachBatchWrapper` and `pythonForeachBatchWrapper` return types to include `AtomicReference[String]`. - `SparkConnectPlanner`: Destructure new return types, set query id on the AtomicReference after query starts. - `SparkConnectSessionHolderSuite`: Update test to destructure 3-element tuple from `pythonForeachBatchWrapper`.

…nHolder for foreachBatch When a streaming query starts, `StreamExecution` clones the SparkSession. The DataFrame passed to foreachBatch has this cloned session, but the `SessionHolder` still references the original session. This causes: 1. Session mismatch: batch DataFrames operate on a different session 2. Session lifetime leak: Python worker keeps original session alive 3. CachedRemoteRelation resolves against the wrong session This commit fixes all three issues by: - Adding `registerExistingSession` to `SparkConnectSessionManager` to wrap an existing SparkSession (the stream clone) under a new session ID without creating another clone. - Adding `ForeachBatchSessionManager` that lazily creates a stream-level `SessionHolder` on the first batch invocation and reuses it for subsequent batches. - Adding `ForeachBatchCleaner` that closes both the Python runner and the stream `SessionHolder` on query termination. - Sending the stream session ID per-batch to the Python worker so it resolves `CachedRemoteRelation` against the correct session. - Updating `foreach_batch_worker.py` to read the stream session ID and create/switch the SparkSession accordingly. The stream `SessionHolder` uses `customInactiveTimeoutMs = -1` (never expire by inactivity); its lifecycle is managed by query termination via `CleanerCache`. Addresses TODO 2 and TODO 3 in StreamingForeachBatchHelper, removing all SPARK-44462 TODO comments.

HyukjinKwon · 2026-04-19T22:57:40Z

@heyihong @HeartSaVioR FYI

…chHelper Run scalafmt to fix formatting violations detected by CI linter.

xiaoxuandev · 2026-04-20T21:17:05Z

+              .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+            val streamSessionId = UUID.randomUUID().toString
+            _streamSessionHolder = SparkConnectService.sessionManager
+              .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)


The registerExistingSession call uses parentSessionHolder.userId and a new random UUID, but nothing ties this stream session back to the parent session's artifacts or classloader. If the user's foreachBatch function relies on session-level artifacts (added jars, UDFs, etc.) from the parent session, would those be missing in the stream SessionHolder? The underlying SparkSession is the cloned one from StreamExecution, so Spark-level config should carry over, but Connect-level SessionHolder state (artifact manager, plan cache, etc.) would be fresh/empty.

Checked this. StreamExecution clones via sparkSession.cloneSession() (StreamExecution.scala:164), and SparkSession.cloneSession() force-copies ArtifactManager and its resources (SparkSession.scala:295). Our cloned SessionHolder wraps that cloned classic SparkSession, and SessionHolder.artifactManager just delegates to session.artifactManager, so jars and the UDF classloader added to the root session before the stream starts are inherited. Artifacts added after the stream is running would not auto-sync, but that mirrors classic mode.

xiaoxuandev · 2026-04-20T21:35:04Z

+        # The per-batch stream session id will be received with each batch and used to create
+        # or switch to the correct stream-level session.
+        connect_url_init = connect_url + ";session_id=" + session_id
+        spark_connect_session = SparkSession.builder.remote(connect_url_init).getOrCreate()


The initial spark_connect_session is created but is never assigned to spark and never used afterward, t only serves as a validation? Is this intended to
remain open for the lifetime of the worker?

I guess we need a session to read the FEB func, but cloned session is not yet available so we have to use the original session first

+1 to @bogao007. We need a session to unpickle the foreachBatch function before the first batch arrives, at which point the cloned session id is not yet known. Once batches start flowing, spark is reassigned to the cloned session via .create(), so the initial one is bootstrap-only.

bogao007 · 2026-04-20T23:14:42Z

    log_name = "Streaming ForeachBatch worker"

-    def process(df_id, batch_id):  # type: ignore[no-untyped-def]
+    def process(df_id, batch_id, stream_session_id):  # type: ignore[no-untyped-def]


Can we use clone_session_id to indicate that this is a cloned session? To differentiate, maybe we could name the previous session to be root_session_id.

Renamed to cloned_session_id / root_session_id (matches the Scala side: clonedSession* / rootSessionHolder).

bogao007 · 2026-04-20T23:16:16Z

+        # The per-batch stream session id will be received with each batch and used to create
+        # or switch to the correct stream-level session.
+        connect_url_init = connect_url + ";session_id=" + session_id
+        spark_connect_session = SparkSession.builder.remote(connect_url_init).getOrCreate()


I guess we need a session to read the FEB func, but cloned session is not yet available so we have to use the original session first

bogao007 · 2026-04-20T23:21:47Z

        def curried_function(df):
            def inner(batch_df, batch_id):
-                df.createOrReplaceTempView("updates")
-                batch_df.createOrReplaceTempView("batch_updates")


Can we still keep the createOrReplaceTempView() to verify that root session cannot access the temp view created by the cloned session? I think this is an important test case and we should make sure the behavior aligns with the classic ForeachBatch

Added test_temp_view_is_isolated_from_root_session. It creates a temp view inside FEB, writes a sentinel table so the test can prove collect_batch actually ran, then asserts the root session hits TABLE_OR_VIEW_NOT_FOUND when reading that view. Kept test_nested_dataframes on saveAsTable — that one is about cross-session persistent tables.

bogao007 · 2026-04-20T23:27:34Z

+          log"[session: ${MDC(SESSION_ID, effectiveHolder.sessionId)}] " +
+            log"Caching DataFrame with id ${MDC(DATAFRAME_ID, dfId)}")
+
+        // Sanity check: remove any stale DataFrame left over from a previous batch for this query.


Why do we need this check? We are already removing the DF from the cache after the batch completes

It is defensive. Happy path is handled by the finally, but if a previous batch skipped its cleanup (async interruption / unchecked error path), a stale dfId could still be sitting in dataFrameQueryIndex. This also lines up with the original TODO(SPARK-44462) that the sanity-check commit replaced. I tightened the comment to spell this out. Fine to drop it if you prefer; it is strictly belt-and-suspenders.

I see, yeah I'm okay with it since it's fixing an existing TODO.

bogao007 · 2026-04-20T23:31:43Z

+        // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
+        val queryId = queryIdRef.get()
+
+        effectiveHolder.cacheDataFrameById(dfId, df)


For a case where a table created outside FEB func() but used inside FEB() to perform a Spark operation, it would probably fail because it's created by the root session and would use root session id to find the cached dataframe. How do we handle that case?

Depends on what kind of table:

persistent / managed tables (Hive metastore etc.): fine, visible across sessions, no CachedRemoteRelation lookup involved.

temp views from the root session: not visible in the cloned session — intentional isolation, same as classic.

DataFrames cached via df.persist() in the root session and captured in the closure: this one would fail because the CachedRemoteRelation id is only registered in the root SessionHolder, not the cloned one. That is a real gap but outside the scope of this PR; I will file a follow-up JIRA. WDYT?

Sure, I'm okay with a follow up Jira, thanks!

AI-review / LLM disclosure: This comment was drafted by an LLM on behalf of the reviewer. Please verify the analysis independently.

[Blocker] Following up with a confirmed reproducer: using root.sql("SELECT value FROM {batch}", batch=batch_df) in foreachBatch fails on this PR head with CONNECT_INVALID_PLAN.DATAFRAME_NOT_FOUND. The worker reports the batch relation ID under the cloned session; the original session then cannot resolve that same ID. A plain uncached DataFrame captured before the query also fails when joined to batch_df, with SESSION_NOT_SAME.

Both tests first verify that the batch and captured objects work independently. See test-only PR #36 and its CI: HeartSaVioR#36 and https://github.com/HeartSaVioR/spark/actions/runs/36569505476/job/109417033218.

This is broader than a persisted-DataFrame corner case. Could we resolve or explicitly define the compatibility behavior in this PR before enabling it by default?

bogao007 · 2026-04-20T23:32:35Z

+      extends Logging {
+    @volatile private var _streamSessionHolder: SessionHolder = null
+
+    def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {


Same as above: can we use clonedSession as naming to indicate that this is a cloned session?

Done — renamed streamSession* → clonedSession* and the ctor param parentSessionHolder → rootSessionHolder so Scala and Python stay symmetric.

- Rename streamSession/parentSessionHolder -> clonedSession/rootSessionHolder across Scala and Python for consistency with StreamExecution's semantics. - Tighten close() to clear the reference before calling closeSession so a failing close cannot leave a half-closed holder; add success/failure logs with MDC. - Eagerly close the cleaner in registerCleanerForQuery when the query is already inactive, so the cloned SessionHolder (which never expires by inactivity) cannot leak if the query terminates before cleaner registration. - Add a parity test test_temp_view_is_isolated_from_root_session that writes a sentinel table from inside foreachBatch to prove the batch ran, then asserts the root session hits TABLE_OR_VIEW_NOT_FOUND for a temp view created only in the cloned session.

- Stub `query.isActive` in StreamingForeachBatchHelperSuite's mockQuery so the eager-cleanup branch added in the previous commit doesn't immediately purge the registered cleaners. - Rewrap the ForeachBatchSessionManager docstring so no line exceeds scalafmt's maxColumn=98.

bogao007 · 2026-04-22T00:02:25Z

+   * The returned SessionHolder has `customInactiveTimeoutMs = Some(-1)` (never expire by
+   * inactivity); its lifecycle is managed by the streaming query termination.
+   */
+  private[connect] def registerExistingSession(


@hvanhovell could you help review the changes here? Thanks!

bogao007

LGTM, thanks for making the fix! I've asked @hvanhovell to help review changes as well.

LuciferYang · 2026-04-22T00:04:50Z

Thank you @bogao007 ~

juanpaulH · 2026-07-14T18:55:48Z

Any tangible resolution or path forward for this critical issue?

Resolved: CleanerCache.registerCleanerForQuery keeps both post-insert cleanup guards (master's isClosing re-check, this branch's !query.isActive); SparkConnectSessionHolderSuite keeps master's listener capture with the ForeachBatchCleaner rename; the foreachBatch worker and parity test follow master's import layout (uuid added, unused unittest dropped).

- Make ForeachBatchSessionManager.close() terminal: a batch arriving after cleanup can no longer register another SessionHolder that never expires by inactivity, and registration is refused while the root session is closing. - Read the cloned holder under the lock so a concurrent close() cannot make getOrCreateClonedSessionHolder return null. - Close the internal cloned session with allowReconnect = true so it does not take a closedSessionsCache slot from a real client session. - Point the worker's default/active session at the cloned session, so SparkSession.active() inside the user function is the session the batch runs on, and drop the cloned-session-id assertion that can never fire. - Tests: terminal close and a closing root session, a cleaner registered for an already terminated query, and an active-session assertion in the parity test.

LuciferYang · 2026-09-01T11:24:31Z

The failing pyspark-connect-old-client check comes from the behavior change, not from flakiness. That job runs branch-4.0's pyspark-connect tests against a server from this build, and branch-4.0's copy of test_nested_dataframes still reads the batch DataFrame's temp view back through the client session. It fails on the spark.sql call after processAllAvailable() returns, so the old client drives the new server to completion; only that assertion breaks.

#58446 guards that test on branch-4.0 with SPARK_SKIP_CONNECT_COMPAT_TESTS, as SPARK-52077 did for ArrowUDTFParityTests. It needs to land before this job can go green here.

… in the Spark Connect compatibility test ### What changes were proposed in this pull request? Skip `StreamingForeachBatchParityTests.test_nested_dataframes` when `SPARK_SKIP_CONNECT_COMPAT_TESTS` is set. ### Why are the changes needed? The Spark Connect compatibility jobs run this branch's `pyspark-connect` tests against a server built from `master`. The `pyspark-connect-old-client` job in `build_and_test.yml` and `build_python_connect40.yml` both clone `branch-4.0` for the client side. #55410 (SPARK-44462) changes the server so the DataFrame handed to `foreachBatch` is bound to the session `StreamExecution` clones for the query, which is what classic `foreachBatch` already does. `test_nested_dataframes` creates a temp view from the batch DataFrame and reads it back through the client session, so against that server it fails: ``` File ".../connect/streaming/test_parity_foreach_batch.py", line 49, in test_nested_dataframes collected = self.spark.sql("select * from batch_updates").collect() pyspark.errors.exceptions.connect.AnalysisException: [TABLE_OR_VIEW_NOT_FOUND] The table or view `batch_updates` cannot be found. ``` Failing run: https://github.com/LuciferYang/spark/actions/runs/33346837991/job/99354964415 The assertion holds for a 4.0 server, so the guard is scoped to the compatibility jobs and this branch's own CI keeps running the test. SPARK-52077 skipped `ArrowUDTFParityTests` the same way. ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? `test_nested_dataframes` still runs with `SPARK_SKIP_CONNECT_COMPAT_TESTS` unset and is skipped with it set to `1`. This branch's CI cannot exercise the case the guard is for, since it runs against a 4.0 server; the `pyspark-connect-old-client` job on #55410 is what verifies it, once this lands. Pass GitHub Actions. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code Closes #58446 from LuciferYang/skip-nested-dataframes-compat-40. Authored-by: YangJie <yangjie01@baidu.com> Signed-off-by: yangjie01 <yangjie01@baidu.com>

LuciferYang · 2026-09-02T11:42:46Z

CI passed

LuciferYang · 2026-09-02T12:12:10Z

@HyukjinKwon

What changes is Python foreachBatch under Connect. The batch DataFrame now belongs to the session StreamExecution clones for the query, so session-scoped state created from it stops reaching the client session:

def process(batch_df, batch_id):
    batch_df.createOrReplaceTempView("v")
...
spark.sql("SELECT * FROM v")   # before: resolves.  after: TABLE_OR_VIEW_NOT_FOUND

Scala and classic are unaffected; both already hand the user function the clone. So Connect is catching up here, but working code does break. Not hypothetically: branch-4.0's test_nested_dataframes broke on it and needed SPARK-59148.

master and branch-4.x can absorb that with a migration note. The reason to go deeper is that on 4.0 through 4.2, whatever the foreachBatch body does with the batch DataFrame is planned under the client session, without the configs the engine sets on the clone, CBO and AQE among them. But going deeper means changing session semantics inside a patch release, which I would rather not do bare.

So: a conf selecting the session, new behavior on master and branch-4.x and old on the maintenance branches? Or conf-free, stopping the backport at branch-4.x?

HyukjinKwon

I think it's fine LGTM

HeartSaVioR · 2026-09-02T21:49:35Z

This is a critical fix, thanks for driving the fix! Would you mind giving me some time to review this? I know Bo reviewed it, but would love to give myself a review.

LuciferYang · 2026-09-03T02:31:19Z

This is a critical fix, thanks for driving the fix! Would you mind giving me some time to review this? I know Bo reviewed it, but would love to give myself a review.

OK, looking forward to your review. Thank you @HeartSaVioR

szehon-ho · 2026-09-29T01:05:16Z

Can we merge it to 4.2.1 as well to keep in sync with 4.3 : https://lists.apache.org/thread/bl3zbvbmrhz2m6jz25ogfwfd6v3wjo4q

HeartSaVioR

AI-review / LLM disclosure: This comment was drafted by an LLM on behalf of the reviewer. Please verify the analysis independently.

I reviewed the Connect foreachBatch session change at cf3458c107dc. It fixes the stateful batch case by running batch-derived work in the stream's cloned session. The following issues remain before this can be enabled by default.

Blocker: ordinary mixed-session operations fail

A callback can capture the original SparkSession and use the batch DataFrame in SQL:

root = spark

def process(batch_df, _):
    root.sql("SELECT value FROM {batch}", batch=batch_df).collect()

root remains bound to the original session after callback deserialization. SQL interpolation includes batch_df's cached relation, but root.sql submits it through the original session. This patch registers that relation only in the cloned holder, so the original holder cannot resolve its ID. Changing SparkSession.active() does not rebind the captured root. The same mechanism applies to other captured session-bound objects that accept the batch as an operation's input.

An independent case uses a plain, uncached DataFrame captured before the query:

lookup = spark.range(2).selectExpr(
    "CASE id WHEN 0 THEN 'hello' ELSE 'this' END AS value")

def process(batch_df, _):
    batch_df.join(lookup, "value").collect()

This fails at DataFrame._check_same_session with SESSION_NOT_SAME, before execution. The revised test_nested_dataframes exercises the captured and batch DataFrames separately, so it does not cover either form of composition.

The same three regression tests are available in test-only PRs against master and this PR's head. Their master-based CI and PR-head CI show:

Test	Master-based run	This PR's head
Stateful AQE / stream session	Fails: root and batch share a session, with AQE enabled in both	Passes
Batch DataFrame in captured-session SQL	Passes	Fails: `CONNECT_INVALID_PLAN.DATAFRAME_NOT_FOUND`
Batch join with captured DataFrame	Passes	Fails: `SESSION_NOT_SAME`

The SQL failure names the same batch relation ID that the worker announced under the cloned session, but reports it missing from the original session. In both failing tests, the batch and captured objects worked independently before the mixed operation failed. Please include these regression tests and resolve the mixed-session behavior in this PR.
If the operations are intentionally unsupported, the compatibility decision and migration path should be explicit before changing the default. Copying the relation into the original holder alone would avoid the lookup error but would still allow operations to bypass stream-specific configuration.

Major: clarify the scope of the configuration guarantee

Reassigning SparkSession.active() does not alter session-bound objects captured by the callback. Such objects can still run independent work under the original session's configuration. Please state that the stream-session guarantee currently applies to operations through batch_df or the active cloned session, and test that boundary.

Major: document the temp-view compatibility change

A temp view created through batch_df is now isolated from the original session. That is consistent with the cloned-session model, but changes callbacks that read the view through a previously captured session. Please document the affected behavior and a migration path for callers relying on visibility from the original session.

Minor: cleanup is not recoverable after a close failure

ForeachBatchSessionManager.close() clears its holder reference before calling closeSession; closeSession removes the holder from its store before SessionHolder.close() completes. If cleanup throws, there is no handle left for a retry. Please add a failure-path test or retain a recoverable cleanup state.

Minor: avoid an unused remote holder for Scala callbacks

The Scala callback receives the classic DataFrame directly and does not use the cached remote relation ID, yet the wrapper registers a separate SessionHolder for it. Please consider creating that holder only for paths that resolve a remote batch relation.

Nit: use the session-ID logging key

The cloned session ID is logged with STREAM_ID. Please use the session-ID key so the structured field matches the value being logged.

HeartSaVioR · 2026-09-29T23:31:32Z

+        // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
+        val queryId = queryIdRef.get()
+
+        effectiveHolder.cacheDataFrameById(dfId, df)


AI-review / LLM disclosure: This comment was drafted by an LLM on behalf of the reviewer. Please verify the analysis independently.

[Blocker] Following up with a confirmed reproducer: using root.sql("SELECT value FROM {batch}", batch=batch_df) in foreachBatch fails on this PR head with CONNECT_INVALID_PLAN.DATAFRAME_NOT_FOUND. The worker reports the batch relation ID under the cloned session; the original session then cannot resolve that same ID. A plain uncached DataFrame captured before the query also fails when joined to batch_df, with SESSION_NOT_SAME.

Both tests first verify that the batch and captured objects work independently. See test-only PR #36 and its CI: HeartSaVioR#36 and https://github.com/HeartSaVioR/spark/actions/runs/36569505476/job/109417033218.

This is broader than a persisted-DataFrame corner case. Could we resolve or explicitly define the compatibility behavior in this PR before enabling it by default?

HeartSaVioR · 2026-09-29T23:32:08Z

+            # unpickled, which requires the active session to still be the root one.
+            with SparkSession._lock:
+                SparkSession._default_session = spark
+            SparkSession._active_session.session = spark


AI-review / LLM disclosure: This comment was drafted by an LLM on behalf of the reviewer. Please verify the analysis independently.

[Major] Updating the active/default session changes what new calls through SparkSession.active() see, but it does not rebind session-bound objects already captured and unpickled under the root session. Those objects can still run independently with root-session configuration. Please scope the stream-session guarantee accordingly and add a test for this boundary.

HeartSaVioR · 2026-09-29T23:32:35Z

+            q.processAllAvailable()
+            self.assertEqual(1, self.spark.table(sentinel_table).count())
+            with self.assertRaisesRegex(AnalysisException, "TABLE_OR_VIEW_NOT_FOUND"):
+                self.spark.sql(f"SELECT * FROM {view_name}").collect()


AI-review / LLM disclosure: This comment was drafted by an LLM on behalf of the reviewer. Please verify the analysis independently.

[Major] This test establishes that a temp view created through batch_df is no longer visible through the original session. That fits the cloned-session model, but it changes callbacks that previously read the view through a captured session. Please document this compatibility change and the migration path.

HeartSaVioR · 2026-09-29T23:33:05Z

-    dataFrameCachingWrapper(
+    val sessionManager = new ForeachBatchSessionManager(sessionHolder)
+    val queryIdRef = new AtomicReference[String]()
+    val wrappedFn = dataFrameCachingWrapper(


AI-review / LLM disclosure: This comment was drafted by an LLM on behalf of the reviewer. Please verify the analysis independently.

[Minor] The Scala callback receives the classic DataFrame directly, and dfId is unused on this path, yet dataFrameCachingWrapper creates a remote SessionHolder and caches the batch relation. Could this holder creation be limited to paths that resolve a remote batch relation?

HeartSaVioR · 2026-09-29T23:33:26Z

+            allowReconnect = true)
+          logInfo(
+            log"[rootSession: ${MDC(SESSION_ID, rootSessionHolder.sessionId)}] " +
+              log"Closed cloned SessionHolder ${MDC(STREAM_ID, holder.sessionId)}.")


AI-review / LLM disclosure: This comment was drafted by an LLM on behalf of the reviewer. Please verify the analysis independently.

[Nit] holder.sessionId is a session ID but is logged with STREAM_ID. Please use the session-ID logging key for this field.

HeartSaVioR · 2026-09-29T23:40:14Z

I think the blocker finding really matters; the query which worked would break with this change. We can't simply defer this to the out of the scope for this PR because we indicate we are going to break some queries. We wouldn't be able to make this change by default if we don't address it.

…chBatch-session-fix

…de, holder/log cleanups - Turn the mixed-session probes into contract tests: root.sql referencing the batch DataFrame and joining it with a captured root-session DataFrame fail with DATAFRAME_NOT_FOUND / SESSION_NOT_SAME. Add migration tests that run the same work through batch_df.sparkSession. - Document the behavior change in the Structured Streaming migration guide (4.3). - Scala foreachBatch no longer registers an unused cloned SessionHolder. - Log the cloned session id under CLONED_SESSION_ID instead of STREAM_ID.

uros-b

Hi folks, thank you for working on this! I understand this is a blocker for Spark 4.3, however - could you please explain whether you intend to backport this to 4.0-4.2 as well? i.e. should @szehon-ho and myself treat this as a release blocker for maintenance releases, or not? It is my understanding that this is a risky change, and probably no backports. Could you please help confirm this? cc @HeartSaVioR @HyukjinKwon

HeartSaVioR · 2026-10-01T05:55:13Z

Can we introduce a config and disable the fix by default? I don't think we want to fix something and break the other thing. If we want to make this by default, I would rather say "Let's just document what the fix does not solve" isn't an acceptable state. I don't want users to read the migration guide and change their query, till we are very sure we tried our best to not break it.

Either land this partial fix with disabling by default (moving the content of migration guide to the config doc) and aim for better fix in future, or hold this fix till we address them.

HeartSaVioR · 2026-10-01T06:14:05Z

I'll need to lower down the tone, sorry about it.

Making a decision to break the existing query isn't a cheap action. Writing the breaking change into migration guide isn't really a no-brainer solution; that only works when they can blindly change the query in the suggested way to fix the issue. If users encounter the issue, and have to think hard on how to migrate the query to unblock, we are not in a good direction to go. That decision must have taken seriously (mostly involving discussion), rather than taken by LLM.

HeartSaVioR · 2026-10-01T06:28:34Z

Sorry for back and forth, I hit the enter button too soon.

The case we describe in migration guide is actually a behavior change from fixing the bug on Spark Connect, and it doesn't seem to be a valid pattern (even for Spark Classic). Maybe the baseline can be set to Spark Classic; if it doesn't work in Spark Classic, it's definitely no-go. If it works in Spark Classic but it's "expected" to not work in Spark Connect (due to technical issue e.g. serialization), it's also no-go. Otherwise I think we need to try fixing it.

Maybe I could spend some time to validate my comment by myself. My apologize for the incomplete feedback.

At least let's have a SQL config to go back to the previous behavior. It's a wrong behavior, but it's always happening that someone cares more about not seeing the breakage.

LuciferYang · 2026-10-01T06:37:07Z

No problem, I didn't leave any comment after pushing b0fad56 because I was also thinking about whether the current approach is sound, if there is a more elegant solution, or whether we need to add a new conf.

HeartSaVioR · 2026-10-01T06:41:06Z

Let's just have a SQL config and mention it in the migration guide. (Fix is effective by default, and they can disable it.) I confirmed that both patterns are valid to be blocked. Our diligence for this is probably just to give users a chance to save their time to migrate the query via disabling the fix while they can still upgrade the cluster.

Thanks for understanding. This is really a critical issue and thanks for working on this and patiently follow this up over multiple months.

… on) Add spark.sql.connect.streaming.foreachBatch.useClonedSession (default true, 4.3.0). When false, the batch DataFrame is bound to the root session (the pre-fix behavior), so it can be combined with the root session, at the cost of not running under the stream session's configuration. Mention it in the Structured Streaming migration guide and add a parity test covering the disabled path.

HeartSaVioR

+1 pending CI

If registerNewStreamingQuery throws after a batch has already created the never-expiring cloned SessionHolder, registering the cleaner first ensures the holder is still reaped (via onQueryTerminated / cleanUpAll / the inactive-query eager close) instead of leaking until server shutdown.

LuciferYang · 2026-10-01T08:15:20Z

Thanks @HeartSaVioR, that works for me. Added spark.sql.connect.streaming.foreachBatch.useClonedSession (default true, so the fix is effective by default). Setting it to false restores the previous behavior: the batch DataFrame stays on the root session, so it can still be combined with the root session, at the cost of not running under the stream session's configuration. The migration guide points at the config, so users who hit the two blocked patterns can disable it and keep their queries while they migrate the cluster.

On the close()-recoverability nit: the cloned holder is torn down via the shared closeSession, which removes it from the session store before SessionHolder.close() runs, and SessionHolder.close() is single-shot (the closedTimeMs guard). So retaining the reference would not enable a real retry without changing that shared teardown ordering, which is out of scope here. close() runs on query termination and is intentionally best-effort: a failure is caught and logged as a possible leak, the same as any session whose close() fails. I would keep it as is rather than add a brittle failure-path test that has to mock the global session manager.

HeartSaVioR · 2026-10-02T05:07:17Z

I was waiting for CI status but it was already succeeded; the result just wasn't reflected in this PR.

HeartSaVioR · 2026-10-02T05:07:29Z

Thanks! Merging to master/4.x/~~4.3~~.

…Spark Connect ### 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 Closes #55410 from LuciferYang/SPARK-44462-foreachBatch-session-fix. Lead-authored-by: YangJie <yangjie01@baidu.com> Co-authored-by: reviewer <reviewer@local> Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com> (cherry picked from commit 362ee55) Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

HeartSaVioR · 2026-10-02T05:10:45Z

@LuciferYang Could you please submit a PR for 4.3? I see there is a merge conflict and this is valuable enough to not slip through. Thanks!

LuciferYang · 2026-10-02T06:20:23Z

@LuciferYang Could you please submit a PR for 4.3? I see there is a merge conflict and this is valuable enough to not slip through. Thanks!

ok, I'll do this

LuciferYang · 2026-10-02T06:22:13Z

Thanks for your reviews. @HeartSaVioR @HyukjinKwon @uros-b @bogao007 @xiaoxuandev

…gisters the cloned session ### What changes were proposed in this pull request? Followup to #55410. The two `ForeachBatchSessionManager` cleanup tests in `StreamingForeachBatchHelperSuite` checked that `close()` is terminal and that the registered cleaner runs, but neither asserted the side effect `close()` exists to guarantee: that the cloned `SessionHolder` is removed from `SparkConnectService.sessionManager`. Because `close()` sets its `closed` flag independently of the `closeSession` call, a future change that dropped or broke the unregister step would still pass both tests while re-introducing the never-expiring-holder leak. This adds assertions that the cloned `SessionKey` is present in the session manager before `close()` and absent after. ### Why are the changes needed? To make the cloned-session leak-prevention contract regression-sensitive. Spotted in review of #59200, the branch-4.3 backport. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GitHub Actions. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code Closes #59202 from LuciferYang/SPARK-44462-assert-cloned-session-unregister. Lead-authored-by: YangJie <yangjie01@baidu.com> Co-authored-by: reviewer <reviewer@local> Signed-off-by: yangjie01 <yangjie01@baidu.com>

…gisters the cloned session ### What changes were proposed in this pull request? Followup to #55410. The two `ForeachBatchSessionManager` cleanup tests in `StreamingForeachBatchHelperSuite` checked that `close()` is terminal and that the registered cleaner runs, but neither asserted the side effect `close()` exists to guarantee: that the cloned `SessionHolder` is removed from `SparkConnectService.sessionManager`. Because `close()` sets its `closed` flag independently of the `closeSession` call, a future change that dropped or broke the unregister step would still pass both tests while re-introducing the never-expiring-holder leak. This adds assertions that the cloned `SessionKey` is present in the session manager before `close()` and absent after. ### Why are the changes needed? To make the cloned-session leak-prevention contract regression-sensitive. Spotted in review of #59200, the branch-4.3 backport. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GitHub Actions. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code Closes #59202 from LuciferYang/SPARK-44462-assert-cloned-session-unregister. Lead-authored-by: YangJie <yangjie01@baidu.com> Co-authored-by: reviewer <reviewer@local> Signed-off-by: yangjie01 <yangjie01@baidu.com> (cherry picked from commit 38422ae) Signed-off-by: yangjie01 <yangjie01@baidu.com>

LuciferYang added 2 commits April 19, 2026 02:56

HyukjinKwon changed the title ~~[SPARK-44462][CONNECT] Fix the session passed to foreachBatch in Spark Connect~~ [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect Apr 19, 2026

LuciferYang added 3 commits April 20, 2026 08:01

[SPARK-44462][CONNECT] Fix scalafmt formatting in StreamingForeachBat…

617898d

…chHelper Run scalafmt to fix formatting violations detected by CI linter.

fix format

2aa36de

fix: Reformat Python files with black

eac2656

xiaoxuandev reviewed Apr 20, 2026

View reviewed changes

bogao007 reviewed Apr 20, 2026

View reviewed changes

LuciferYang added 2 commits April 21, 2026 14:42

bogao007 reviewed Apr 22, 2026

View reviewed changes

bogao007 approved these changes Apr 22, 2026

View reviewed changes

LuciferYang added 2 commits August 31, 2026 03:13

HyukjinKwon approved these changes Aug 31, 2026

View reviewed changes

LuciferYang mentioned this pull request Sep 1, 2026

[SPARK-59148][PYTHON][CONNECT][TEST][4.0] Skip test_nested_dataframes in the Spark Connect compatibility test #58446

Closed

Merge branch 'apache:master' into SPARK-44462-foreachBatch-session-fix

6360843

HyukjinKwon approved these changes Sep 2, 2026

View reviewed changes

Merge branch 'apache:master' into SPARK-44462-foreachBatch-session-fix

cf3458c

HeartSaVioR reviewed Sep 29, 2026

View reviewed changes

reviewer added 2 commits September 30, 2026 16:47

Merge remote-tracking branch 'upstream/master' into SPARK-44462-forea…

2e46999

…chBatch-session-fix

uros-b reviewed Sep 30, 2026

View reviewed changes

HeartSaVioR approved these changes Oct 1, 2026

View reviewed changes

reviewer added 2 commits October 1, 2026 15:21

Tighten the foreachBatch migration note

f8bcd3b

reviewer added 3 commits October 1, 2026 18:23

style: apply scalafmt to foreachBatch session manager call sites

cd39c4f

style: collapse sql_conf call to satisfy ruff format

cec09ab

fix: declare NOT_APPLICABLE binding policy for useClonedSession conf

eb7914f

HeartSaVioR closed this in 362ee55 Oct 2, 2026

LuciferYang mentioned this pull request Oct 2, 2026

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

Closed

LuciferYang mentioned this pull request Oct 2, 2026

[SPARK-44462][FOLLOWUP][SS][CONNECT] Assert foreachBatch close() unregisters the cloned session #59202

Closed

Conversation

LuciferYang commented Apr 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

HyukjinKwon commented Apr 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bogao007 left a comment

Choose a reason for hiding this comment

Uh oh!

LuciferYang commented Apr 22, 2026

Uh oh!

juanpaulH commented Jul 14, 2026

Uh oh!

LuciferYang commented Sep 1, 2026

Uh oh!

LuciferYang commented Sep 2, 2026

Uh oh!

LuciferYang commented Sep 2, 2026

Uh oh!

HyukjinKwon left a comment

Choose a reason for hiding this comment

Uh oh!

HeartSaVioR commented Sep 2, 2026

Uh oh!

LuciferYang commented Sep 3, 2026

Uh oh!

szehon-ho commented Sep 29, 2026

Uh oh!

HeartSaVioR left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Blocker: ordinary mixed-session operations fail

Major: clarify the scope of the configuration guarantee

Major: document the temp-view compatibility change

Minor: cleanup is not recoverable after a close failure

LuciferYang commented Apr 19, 2026 •

edited

Loading

HeartSaVioR left a comment •

edited

Loading

HeartSaVioR commented Oct 1, 2026 •

edited

Loading

HeartSaVioR commented Oct 1, 2026 •

edited

Loading

HeartSaVioR commented Oct 1, 2026 •

edited

Loading

HeartSaVioR commented Oct 2, 2026 •

edited

Loading