This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new b7503aa80f3d [SPARK-44462][SS][CONNECT] Fix the session passed to
foreachBatch in Spark Connect
b7503aa80f3d is described below
commit b7503aa80f3dd1187ac73dd47b249c46197c10ef
Author: YangJie <[email protected]>
AuthorDate: Fri Oct 2 14:08:34 2026 +0900
[SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in 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.conne [...]
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 hold [...]
- `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
[...]
### 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.us [...]
### 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 <[email protected]>
Co-authored-by: reviewer <reviewer@local>
Signed-off-by: Jungtaek Lim <[email protected]>
(cherry picked from commit 362ee5500dd7c05328886ce46f90934ab32682e5)
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../java/org/apache/spark/internal/LogKeys.java | 1 +
docs/streaming/ss-migration-guide.md | 4 +
.../streaming/worker/foreach_batch_worker.py | 54 ++--
.../connect/streaming/test_parity_foreach_batch.py | 218 ++++++++++++++++-
.../org/apache/spark/sql/internal/SQLConf.scala | 14 ++
.../sql/connect/planner/SparkConnectPlanner.scala | 30 ++-
.../planner/StreamingForeachBatchHelper.scala | 272 +++++++++++++++------
.../spark/sql/connect/service/SessionHolder.scala | 5 +
.../service/SparkConnectSessionManager.scala | 43 ++++
.../planner/StreamingForeachBatchHelperSuite.scala | 56 +++++
.../service/SparkConnectSessionHolderSuite.scala | 11 +-
11 files changed, 605 insertions(+), 103 deletions(-)
diff --git
a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
index 9c533e4e106e..04cdd718b613 100644
--- a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
+++ b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
@@ -102,6 +102,7 @@ public enum LogKeys implements LogKey {
CLASS_PATHS,
CLAUSES,
CLEANUP_LOCAL_DIRS,
+ CLONED_SESSION_ID,
CLUSTER_CENTROIDS,
CLUSTER_ID,
CLUSTER_LABEL,
diff --git a/docs/streaming/ss-migration-guide.md
b/docs/streaming/ss-migration-guide.md
index 026dde80026b..f31eae9fc2ca 100644
--- a/docs/streaming/ss-migration-guide.md
+++ b/docs/streaming/ss-migration-guide.md
@@ -30,6 +30,10 @@ Please refer [Migration Guide: SQL, Datasets and
DataFrame](../sql-migration-gui
unsupported-operation checks. Existing queries can still restart from their
persistent
checkpoints using the compatibility value recorded in the offset log.
+## Upgrading from Structured Streaming 4.2 to 4.3
+
+- Since Spark 4.3, in Spark Connect the DataFrame passed to a Python
`foreachBatch` function is bound to the streaming query's own session, a clone
of the session that started the query, matching classic `foreachBatch`. This
runs the batch under the configuration the streaming engine pins on the clone
rather than the root session's. As a result, the batch DataFrame can no longer
be combined with the root session inside the function: a temporary view created
from it is not visible to the [...]
+
## Upgrading from Structured Streaming 4.1 to 4.2
- Since Spark 4.2, restarting a streaming query from a checkpoint whose
metadata file is missing while the offset or commit logs contain data fails
with `STREAMING_CHECKPOINT_MISSING_METADATA_FILE`, instead of silently
generating a new query ID (which can duplicate data in exactly-once sinks).
Restore the metadata file or use a new checkpoint location. To restore the
previous behavior, set
`spark.sql.streaming.checkpoint.verifyMetadataExists.enabled` to `false`. (See
[SPARK-55058](https: [...]
diff --git
a/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py
b/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py
index 998805aa8cfc..e7f7c83bfbdc 100644
--- a/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py
+++ b/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py
@@ -24,6 +24,7 @@ import os
from typing import IO
from pyspark import worker
+from pyspark.errors import PySparkAssertionError
from pyspark.serializers import (
CPickleSerializer,
UTF8Deserializer,
@@ -46,13 +47,32 @@ def main(infile: IO, outfile: IO) -> None:
log_name = "Streaming ForeachBatch worker"
- def process(df_id, batch_id): # type: ignore[no-untyped-def]
+ def process(df_id, batch_id, cloned_session_id): # type:
ignore[no-untyped-def]
global spark
- print(f"{log_name} Started batch {batch_id} with DF id {df_id} and
session id {session_id}")
- batch_df = spark_connect_session._create_remote_dataframe(df_id)
+ # Lazily create or switch the SparkSession when the cloned session
changes.
+ # Use create() instead of getOrCreate() because the latter returns the
cached
+ # active/default session, ignoring the cloned session ID in the URL.
+ if spark is None or spark.session_id != cloned_session_id:
+ cloned_url = connect_url + ";session_id=" + cloned_session_id
+ spark = SparkSession.builder.remote(cloned_url).create()
+ # create() fills the default/active slots only when they are
empty, and the bootstrap
+ # session below already claimed them. Point them at the cloned
session so that
+ # SparkSession.active() and getActiveSession() inside the user
function return the
+ # session the batch actually runs on. This has to happen after the
user function was
+ # unpickled, which requires the active session to still be the
root one.
+ with SparkSession._lock:
+ SparkSession._default_session = spark
+ SparkSession._active_session.session = spark
+ print(f"{log_name} Created new session for cloned_session_id
{cloned_session_id}")
+ print(
+ f"{log_name} Started batch {batch_id} with DF id {df_id} "
+ f"and session id {cloned_session_id}"
+ )
+ batch_df = spark._create_remote_dataframe(df_id)
func(batch_df, batch_id)
print(
- f"{log_name} Completed batch {batch_id} with DF id {df_id} and
session id {session_id}"
+ f"{log_name} Completed batch {batch_id} with DF id {df_id} "
+ f"and session id {cloned_session_id}"
)
try:
@@ -62,15 +82,20 @@ def main(infile: IO, outfile: IO) -> None:
os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1"
connect_url = os.environ["SPARK_CONNECT_LOCAL_URL"]
- session_id = utf8_deserializer.loads(infile)
+ root_session_id = utf8_deserializer.loads(infile)
- print(f"{log_name} is starting with url {connect_url} and sessionId
{session_id}.")
+ print(f"{log_name} is starting with url {connect_url} and sessionId
{root_session_id}.")
- # To attach to the existing SparkSession, we're setting the session_id
in the URL.
- connect_url = connect_url + ";session_id=" + session_id
- spark_connect_session =
SparkSession.builder.remote(connect_url).getOrCreate()
- assert spark_connect_session.session_id == session_id
- spark = spark_connect_session
+ # Bootstrap SparkSession on the root session id: needed to unpickle
the user's
+ # foreachBatch function before any batch (and the cloned session id)
arrives.
+ # Per-batch work runs against the cloned session created in
`process()`.
+ connect_url_init = connect_url + ";session_id=" + root_session_id
+ spark_connect_session =
SparkSession.builder.remote(connect_url_init).getOrCreate()
+ if spark_connect_session.session_id != root_session_id:
+ raise PySparkAssertionError(
+ f"Root session ID mismatch: expected {root_session_id}, "
+ f"got {spark_connect_session.session_id}"
+ )
func = worker.read_command(pickle_ser, infile)
write_int(0, outfile)
@@ -79,9 +104,10 @@ def main(infile: IO, outfile: IO) -> None:
while True:
df_ref_id = utf8_deserializer.loads(infile)
batch_id = read_long(infile)
- # Handle errors inside Python worker. Write 0 to outfile if no
errors and write -2 with
- # traceback string if error occurs.
- process(df_ref_id, int(batch_id))
+ cloned_session_id = utf8_deserializer.loads(infile)
+ # Handle errors inside Python worker. Write 0 to outfile if no
errors and write -2
+ # with traceback string if error occurs.
+ process(df_ref_id, int(batch_id), cloned_session_id)
write_int(0, outfile)
outfile.flush()
except Exception as e:
diff --git
a/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py
b/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py
index b6dd23353eb9..42fcd151cfa8 100644
--- a/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py
+++ b/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py
@@ -16,14 +16,18 @@
#
import time
+import uuid
-from pyspark.errors import PySparkPicklingError
+from pyspark.errors import PySparkPicklingError, StreamingQueryException
from pyspark.sql.tests.streaming.test_streaming_foreach_batch import
StreamingTestsForeachBatchMixin
from pyspark.testing.connectutils import ReusedConnectTestCase,
should_test_connect
from pyspark.testing.utils import eventually, timeout
if should_test_connect:
- from pyspark.errors.exceptions.connect import
StreamingPythonRunnerInitializationException
+ from pyspark.errors.exceptions.connect import (
+ AnalysisException,
+ StreamingPythonRunnerInitializationException,
+ )
class StreamingForeachBatchParityTests(StreamingTestsForeachBatchMixin,
ReusedConnectTestCase):
@@ -43,24 +47,224 @@ class
StreamingForeachBatchParityTests(StreamingTestsForeachBatchMixin, ReusedCo
self.assertIsNone(q.exception(), "No exception has to be propagated.")
def test_nested_dataframes(self):
+ # Tests that closured DataFrames and batch DataFrames can both be used
+ # inside foreachBatch. In Connect, batch_df runs in a separate cloned
session,
+ # so we use saveAsTable (visible cross-session) instead of temp views.
def curried_function(df):
def inner(batch_df, batch_id):
- df.createOrReplaceTempView("updates")
- batch_df.createOrReplaceTempView("batch_updates")
+
df.write.format("parquet").mode("overwrite").saveAsTable("nested_df_updates")
+ batch_df.write.format("parquet").mode("overwrite").saveAsTable(
+ "nested_df_batch_updates"
+ )
return inner
+ q = None
try:
df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
other_df = self.spark.range(100)
q = df.writeStream.foreachBatch(curried_function(other_df)).start()
q.processAllAvailable()
- collected = self.spark.sql("select * from batch_updates").collect()
- self.assertTrue(len(collected), 2)
- self.assertEqual(100, self.spark.sql("select * from
updates").count())
+ collected = self.spark.sql("select * from
nested_df_batch_updates").collect()
+ self.assertTrue(len(collected) > 0)
+ self.assertEqual(100, self.spark.sql("select * from
nested_df_updates").count())
finally:
if q:
q.stop()
+ self.spark.sql("DROP TABLE IF EXISTS nested_df_updates")
+ self.spark.sql("DROP TABLE IF EXISTS nested_df_batch_updates")
+
+ def test_temp_view_is_isolated_from_root_session(self):
+ # A temp view created inside foreachBatch lives in the cloned session
and must
+ # not be visible from the root session.
+ suffix = uuid.uuid4().hex
+ view_name = f"cloned_only_view_{suffix}"
+ sentinel_table = f"cloned_session_ran_{suffix}"
+
+ def collect_batch(batch_df, _):
+ from pyspark.sql import SparkSession
+
+ # The batch runs on the cloned session, so that is what the
standard entry point
+ # must hand back inside the function.
+ assert SparkSession.active() is batch_df.sparkSession, (
+ f"active session {SparkSession.active().session_id} is not the
batch session "
+ f"{batch_df.sparkSession.session_id}"
+ )
+ batch_df.createOrReplaceTempView(view_name)
+ # Read through the cloned session; a resolution failure fails the
query.
+ batch_df.sparkSession.sql(f"SELECT * FROM {view_name}").collect()
+ # Persist a sentinel so the test can prove collect_batch actually
ran,
+ # otherwise the negative assertion below would trivially pass.
+ batch_df.sparkSession.createDataFrame([(1,)], ["v"]).write.mode(
+ "overwrite"
+ ).saveAsTable(sentinel_table)
+
+ q = None
+ try:
+ df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+ q = df.writeStream.foreachBatch(collect_batch).start()
+ 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()
+ finally:
+ if q:
+ q.stop()
+ self.spark.sql(f"DROP TABLE IF EXISTS {sentinel_table}")
+ # Best-effort cleanup in case a regression ever leaks the view.
+ try:
+ self.spark.catalog.dropTempView(view_name)
+ except Exception:
+ pass
+
+ def test_batch_uses_stream_session_for_stateful_aqe(self):
+ # SPARK-44462: the batch runs on the stream's cloned session, which
pins the stateful
+ # configs (AQE off) independently of the root session. Proves the
isolation this PR adds.
+ root = self.spark
+
+ def process(batch_df, _):
+ batch_count = batch_df.count()
+ root_aqe = root.conf.get("spark.sql.adaptive.enabled")
+ batch_aqe =
batch_df.sparkSession.conf.get("spark.sql.adaptive.enabled")
+ same_session = root.session_id == batch_df.sparkSession.session_id
+ if same_session or root_aqe != "true" or batch_aqe != "false":
+ raise AssertionError(
+ f"batch_count={batch_count}, same_session={same_session}, "
+ f"AQE root={root_aqe}, batch={batch_aqe}"
+ )
+
+ q = None
+ with self.sql_conf({"spark.sql.adaptive.enabled": "true"}):
+ try:
+ df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+ q = (
+ df.groupBy("value")
+ .count()
+ .writeStream.outputMode("complete")
+ .foreachBatch(process)
+ .start()
+ )
+ q.processAllAvailable()
+ self.assertIsNone(q.exception())
+ finally:
+ if q:
+ q.stop()
+
+ def test_batch_dataframe_in_root_session_sql_is_unsupported(self):
+ # SPARK-44462 contract: the batch DataFrame is bound to the stream's
cloned session, so its
+ # cached relation cannot be resolved through a captured root session.
Documented behavior.
+ root = self.spark
+
+ def process(batch_df, _):
+ root.sql("SELECT value FROM {batch}", batch=batch_df).collect()
+
+ q = None
+ try:
+ df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+ q = df.writeStream.foreachBatch(process).start()
+ self.assertRaises(StreamingQueryException, q.processAllAvailable)
+ self.assertIn("DATAFRAME_NOT_FOUND", str(q.exception()))
+ finally:
+ if q:
+ q.stop()
+
+ def
test_batch_dataframe_join_with_root_session_dataframe_is_unsupported(self):
+ # SPARK-44462 contract: the batch DataFrame (cloned session) cannot be
combined with a
+ # DataFrame captured from the root session; the client rejects the
cross-session join.
+ lookup = self.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()
+
+ q = None
+ try:
+ df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+ q = df.writeStream.foreachBatch(process).start()
+ self.assertRaises(StreamingQueryException, q.processAllAvailable)
+ self.assertIn("SESSION_NOT_SAME", str(q.exception()))
+ finally:
+ if q:
+ q.stop()
+
+ def test_batch_dataframe_in_batch_session_sql_is_supported(self):
+ # SPARK-44462 migration path: resolve the batch DataFrame through its
own (cloned) session
+ # instead of a captured root session.
+ def process(batch_df, _):
+ expected = sorted(row.value for row in
batch_df.select("value").collect())
+ actual = sorted(
+ row.value
+ for row in batch_df.sparkSession.sql(
+ "SELECT value FROM {batch}", batch=batch_df
+ ).collect()
+ )
+ if actual != expected:
+ raise AssertionError(f"SQL rows {actual} != batch rows
{expected}")
+
+ q = None
+ try:
+ df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+ q = df.writeStream.foreachBatch(process).start()
+ q.processAllAvailable()
+ self.assertTrue(any(p.numInputRows > 0 for p in q.recentProgress))
+ self.assertIsNone(q.exception())
+ finally:
+ if q:
+ q.stop()
+
+ def test_batch_dataframe_join_in_batch_session_is_supported(self):
+ # SPARK-44462 migration path: rebuild the lookup on the batch's own
(cloned) session so the
+ # join stays within a single session.
+ def process(batch_df, _):
+ spark = batch_df.sparkSession
+ lookup = spark.range(2).selectExpr(
+ "CASE id WHEN 0 THEN 'hello' ELSE 'this' END AS value"
+ )
+ expected = sorted(row.value for row in
batch_df.select("value").collect())
+ actual = sorted(
+ row.value for row in batch_df.join(lookup,
"value").select("value").collect()
+ )
+ if actual != expected:
+ raise AssertionError(f"Joined rows {actual} != batch rows
{expected}")
+
+ q = None
+ try:
+ df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+ q = df.writeStream.foreachBatch(process).start()
+ q.processAllAvailable()
+ self.assertTrue(any(p.numInputRows > 0 for p in q.recentProgress))
+ self.assertIsNone(q.exception())
+ finally:
+ if q:
+ q.stop()
+
+ def test_disabling_fix_restores_root_session_behavior(self):
+ # With the fix disabled, the batch DataFrame is bound to the root
session again, so it can
+ # be combined with a DataFrame captured from the root session.
+ with
self.sql_conf({"spark.sql.connect.streaming.foreachBatch.useClonedSession":
"false"}):
+ lookup = self.spark.range(2).selectExpr(
+ "CASE id WHEN 0 THEN 'hello' ELSE 'this' END AS value"
+ )
+
+ def process(batch_df, _):
+ expected = sorted(row.value for row in
batch_df.select("value").collect())
+ actual = sorted(
+ row.value for row in batch_df.join(lookup,
"value").select("value").collect()
+ )
+ if actual != expected:
+ raise AssertionError(f"Joined rows {actual} != batch rows
{expected}")
+
+ q = None
+ try:
+ df =
self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+ q = df.writeStream.foreachBatch(process).start()
+ q.processAllAvailable()
+ self.assertTrue(any(p.numInputRows > 0 for p in
q.recentProgress))
+ self.assertIsNone(q.exception())
+ finally:
+ if q:
+ q.stop()
def test_pickling_error(self):
class NoPickle:
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 1c4c93dfebe6..1c50810684ff 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -4988,6 +4988,20 @@ object SQLConf {
.version("4.1.0")
.fallbackConf(SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED)
+ val CONNECT_STREAMING_FOREACH_BATCH_USE_CLONED_SESSION =
+ buildConf("spark.sql.connect.streaming.foreachBatch.useClonedSession")
+ .doc("When true, the DataFrame passed to a Python foreachBatch function
under Spark " +
+ "Connect is bound to the streaming query's own session, which is a
clone of the " +
+ "session that started the query, matching classic foreachBatch. This
runs the batch " +
+ "under the configuration the streaming engine pins on the clone rather
than the root " +
+ "session's. Set to false to restore the previous behavior, where the
batch DataFrame " +
+ "is bound to the root session: this lets it be combined with the root
session inside " +
+ "the function, but the batch no longer runs under the stream session's
configuration.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(true)
+
val THRIFTSERVER_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED =
buildConf("spark.sql.thriftserver.shuffleDependency.fileCleanup.enabled")
.doc("When enabled, shuffle files will be cleaned up at the end of
Thrift server " +
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
index faa70c099c2c..3412dca56817 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.connect.planner
import java.util.{HashMap, Properties, UUID}
+import java.util.concurrent.atomic.AtomicReference
import scala.collection.mutable
import scala.jdk.CollectionConverters._
@@ -3605,22 +3606,28 @@ class SparkConnectPlanner(
}
}
- // This is filled when a foreach batch runner started for Python.
+ // This is filled when a foreachBatch runner is started (Python or Scala).
var foreachBatchRunnerCleaner: Option[AutoCloseable] = None
+ // Filled when foreachBatch is used, to set query id after query starts.
+ var foreachBatchQueryIdRef: Option[AtomicReference[String]] = None
if (writeOp.hasForeachBatch) {
val foreachBatchFn = writeOp.getForeachBatch.getFunctionCase match {
case StreamingForeachFunction.FunctionCase.PYTHON_FUNCTION =>
val pythonFn =
transformPythonFunction(writeOp.getForeachBatch.getPythonFunction)
- val (fn, cleaner) =
+ val (fn, cleaner, queryIdRef) =
StreamingForeachBatchHelper.pythonForeachBatchWrapper(pythonFn,
sessionHolder)
foreachBatchRunnerCleaner = Some(cleaner)
+ foreachBatchQueryIdRef = Some(queryIdRef)
fn
case StreamingForeachFunction.FunctionCase.SCALA_FUNCTION =>
- StreamingForeachBatchHelper.scalaForeachBatchWrapper(
+ val (fn, cleaner, queryIdRef) =
StreamingForeachBatchHelper.scalaForeachBatchWrapper(
writeOp.getForeachBatch.getScalaFunction.getPayload.toByteArray,
sessionHolder)
+ foreachBatchRunnerCleaner = Some(cleaner)
+ foreachBatchQueryIdRef = Some(queryIdRef)
+ fn
case other =>
throw InvalidInputErrors.invalidOneOfField(
@@ -3647,18 +3654,23 @@ class SparkConnectPlanner(
throw ex
}
+ // Set the query id so the sanity check in dataFrameCachingWrapper can use
it.
+ foreachBatchQueryIdRef.foreach(_.set(query.id.toString))
+ // Register the cleaner with the query if foreachBatch is used. Do this
before
+ // registerNewStreamingQuery so that if that call throws, the cleaner is
already wired to reap
+ // the cloned SessionHolder a batch may have created (it never expires by
inactivity, so an
+ // unregistered one would leak until server shutdown).
+ foreachBatchRunnerCleaner.foreach { cleaner =>
+
sessionHolder.streamingForeachBatchRunnerCleanerCache.registerCleanerForQuery(
+ query,
+ cleaner)
+ }
// Register the new query so that its reference is cached and is stopped
on session timeout.
SparkConnectService.streamingSessionManager.registerNewStreamingQuery(
sessionHolder,
query,
executeHolder.sparkSessionTags,
executeHolder.operationId)
- // Register the runner with the query if Python foreachBatch is enabled.
- foreachBatchRunnerCleaner.foreach { cleaner =>
-
sessionHolder.streamingForeachBatchRunnerCleanerCache.registerCleanerForQuery(
- query,
- cleaner)
- }
executeHolder.eventsManager.postFinished()
val resultBuilder = WriteStreamOperationStartResult
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala
index 29086b182c4f..8bc584a7aefd 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala
@@ -20,6 +20,7 @@ import java.io.EOFException
import java.util.UUID
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.ConcurrentMap
+import java.util.concurrent.atomic.AtomicReference
import scala.jdk.CollectionConverters._
import scala.util.control.NonFatal
@@ -27,15 +28,16 @@ import scala.util.control.NonFatal
import org.apache.spark.SparkException
import org.apache.spark.api.python.{PythonException, PythonWorkerUtils,
SimplePythonFunction, SpecialLengths, StreamingPythonRunner}
import org.apache.spark.internal.Logging
-import org.apache.spark.internal.LogKeys.{DATAFRAME_ID, PYTHON_EXEC, QUERY_ID,
RUN_ID_STRING, SESSION_ID, USER_ID}
+import org.apache.spark.internal.LogKeys.{CLONED_SESSION_ID, DATAFRAME_ID,
PYTHON_EXEC, QUERY_ID, RUN_ID_STRING, SESSION_ID, USER_ID}
import org.apache.spark.sql.{DataFrame, Dataset}
import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder,
AgnosticEncoders}
import org.apache.spark.sql.connect.IllegalStateErrors
import org.apache.spark.sql.connect.common.{ForeachWriterPacket,
UdfSerialization}
import org.apache.spark.sql.connect.config.Connect
-import org.apache.spark.sql.connect.service.SessionHolder
+import org.apache.spark.sql.connect.service.{SessionHolder, SessionKey}
import org.apache.spark.sql.connect.service.SparkConnectService
import org.apache.spark.sql.execution.python.PythonWorkerEnvironment
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.StreamingQuery
import org.apache.spark.sql.streaming.StreamingQueryListener
import org.apache.spark.util.Utils
@@ -47,47 +49,174 @@ object StreamingForeachBatchHelper extends Logging {
type ForeachBatchFnType = (DataFrame, Long) => Unit
- // Visible for testing.
- /** An AutoClosable to clean up resources on query termination. Stops Python
worker. */
- private[connect] case class RunnerCleaner(runner: StreamingPythonRunner)
extends AutoCloseable {
+ /**
+ * Wraps the per-stream cloned `SparkSession` produced by
+ * `StreamExecution.sparkSessionForStream` in its own `SessionHolder` and
registers it with
+ * `SparkConnectSessionManager`. The cloned session id is pinned on the
first batch and shared
+ * by all subsequent batches of the same streaming query, so the Python
worker can resolve
+ * `CachedRemoteRelation` for the batch DataFrame against the cloned session
rather than the
+ * root session.
+ *
+ * Lifecycle is driven by the query: created on first batch, closed on query
termination via
+ * `ForeachBatchCleaner`. Not designed for concurrent batches (streaming
runs them serially).
+ *
+ * When `useClonedSession` is false (the fix disabled via
+ * `spark.sql.connect.streaming.foreachBatch.useClonedSession`), no cloned
holder is created and
+ * the root holder is used instead, restoring the pre-fix behavior.
+ */
+ private[connect] class ForeachBatchSessionManager(
+ rootSessionHolder: SessionHolder,
+ useClonedSession: Boolean)
+ extends Logging {
+ @volatile private var _clonedSessionHolder: SessionHolder = null
+ // Set by close(). Registration must not happen again afterwards: the
holder never expires by
+ // inactivity, and the cleaner that would have closed it is already gone
from the cache. Read
+ // and written under this instance's lock, so it needs no @volatile.
+ private var closed = false
+
+ def getOrCreateClonedSessionHolder(batchDf: DataFrame): SessionHolder = {
+ val existing = _clonedSessionHolder
+ if (!useClonedSession) {
+ // Gated off: keep the pre-fix behavior by caching the batch DataFrame
in the root holder
+ // and binding the worker to 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).
+ rootSessionHolder
+ } else if (existing != null) {
+ existing
+ } else {
+ synchronized {
+ if (_clonedSessionHolder == null) {
+ if (closed) {
+ throw IllegalStateErrors.streamLifecycleAlreadyCompleted(
+ "getOrCreateClonedSessionHolder")
+ }
+ if (rootSessionHolder.isClosing) {
+ throw
IllegalStateErrors.sessionAlreadyClosed(rootSessionHolder.key.toString)
+ }
+ val clonedSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val clonedSessionId = UUID.randomUUID().toString
+ _clonedSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(rootSessionHolder.userId,
clonedSessionId, clonedSession)
+ logInfo(
+ log"[rootSession: ${MDC(SESSION_ID,
rootSessionHolder.sessionId)}] " +
+ log"Registered cloned SessionHolder " +
+ log"${MDC(CLONED_SESSION_ID, clonedSessionId)} for
foreachBatch.")
+ }
+ // Read under the lock so close() cannot null the field out between
the registration
+ // above and the return; a null holder would NPE in the caller
instead of failing the
+ // batch.
+ _clonedSessionHolder
+ }
+ }
+ }
+
+ def close(): Unit = synchronized {
+ closed = true
+ if (_clonedSessionHolder != null) {
+ val holder = _clonedSessionHolder
+ // Clear the reference before closing so a failure cannot leave a
half-closed holder.
+ _clonedSessionHolder = null
+ try {
+ // allowReconnect = true: no client ever reconnects with this
server-generated id, so a
+ // tombstone in closedSessionsCache would only take a slot from a
real client session.
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(holder.userId, holder.sessionId),
+ allowReconnect = true)
+ logInfo(
+ log"[rootSession: ${MDC(SESSION_ID, rootSessionHolder.sessionId)}]
" +
+ log"Closed cloned SessionHolder ${MDC(CLONED_SESSION_ID,
holder.sessionId)}.")
+ } catch {
+ case NonFatal(ex) =>
+ logWarning(
+ log"[rootSession: ${MDC(SESSION_ID,
rootSessionHolder.sessionId)}] " +
+ log"Error closing cloned SessionHolder " +
+ log"${MDC(CLONED_SESSION_ID, holder.sessionId)} for
foreachBatch; " +
+ log"it may be leaked.",
+ ex)
+ }
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the
cloned SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: Option[ForeachBatchSessionManager])
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
- catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ sessionManager.foreach { sm =>
+ try sm.close()
+ catch {
+ case NonFatal(_) => // already logged inside close()
+ }
}
}
}
- private case class FnArgsWithId(dfId: String, df: DataFrame, batchId: Long)
+ private case class FnArgsWithId(
+ dfId: String,
+ df: DataFrame,
+ batchId: Long,
+ clonedSessionId: String)
/**
* Return a new ForeachBatch function that wraps `fn`. It sets up DataFrame
cache so that the
* user function can access it. The cache is cleared once ForeachBatch
returns.
+ *
+ * On the first batch, lazily creates a cloned-session-level SessionHolder
via sessionManager.
+ * Batch DataFrames are cached in the cloned SessionHolder rather than the
root session.
*/
private def dataFrameCachingWrapper(
fn: FnArgsWithId => Unit,
- sessionHolder: SessionHolder): ForeachBatchFnType = { (df: DataFrame,
batchId: Long) =>
- {
- val dfId = UUID.randomUUID().toString
- // TODO: Add query id to the log.
- logInfo(
- log"[session: ${MDC(SESSION_ID, sessionHolder.sessionId)}] " +
- log"Caching DataFrame with id ${MDC(DATAFRAME_ID, dfId)}")
-
- // TODO(SPARK-44462): Sanity check there is no other active DataFrame
for this query.
- // The query id needs to be saved in the cache for this check.
-
- sessionHolder.cacheDataFrameById(dfId, df)
- try {
- fn(FnArgsWithId(dfId, df, batchId))
- } finally {
+ sessionManager: ForeachBatchSessionManager,
+ queryIdRef: AtomicReference[String]): ForeachBatchFnType = {
+ (df: DataFrame, batchId: Long) =>
+ {
+ val effectiveHolder = sessionManager.getOrCreateClonedSessionHolder(df)
+ val dfId = UUID.randomUUID().toString
logInfo(
- log"[session: ${MDC(SESSION_ID, sessionHolder.sessionId)}] " +
- log"Removing DataFrame with id ${MDC(DATAFRAME_ID, dfId)} from the
cache")
- sessionHolder.removeCachedDataFrame(dfId)
+ log"[session: ${MDC(SESSION_ID, effectiveHolder.sessionId)}] " +
+ log"Caching DataFrame with id ${MDC(DATAFRAME_ID, dfId)}")
+
+ // Defensive: evict any stale dfId from a prior batch whose cleanup
was skipped
+ // (e.g., async interruption). No-op on the happy path.
+ val queryId = queryIdRef.get()
+
+ effectiveHolder.cacheDataFrameById(dfId, df)
+ if (queryId != null) {
+ Option(effectiveHolder.dataFrameQueryIndex.put(queryId,
dfId)).foreach { staleDfId =>
+ logWarning(
+ log"[session: ${MDC(SESSION_ID, effectiveHolder.sessionId)}] " +
+ log"[queryId: ${MDC(QUERY_ID, queryId)}] " +
+ log"Stale DataFrame ${MDC(DATAFRAME_ID, staleDfId)} found in
cache. Removing it.")
+ effectiveHolder.removeCachedDataFrame(staleDfId)
+ }
+ }
+
+ try {
+ fn(FnArgsWithId(dfId, df, batchId, effectiveHolder.sessionId))
+ } finally {
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, effectiveHolder.sessionId)}] " +
+ log"Removing DataFrame with id ${MDC(DATAFRAME_ID, dfId)} from
the cache")
+ effectiveHolder.removeCachedDataFrame(dfId)
+ // Clean up query-to-dfId mapping.
+ if (queryId != null) {
+ effectiveHolder.dataFrameQueryIndex.remove(queryId, dfId)
+ }
+ }
}
- }
}
/**
@@ -95,48 +224,45 @@ object StreamingForeachBatchHelper extends Logging {
* provided foreachBatch function `fn`.
*
* HACK ALERT: This version does not actually set up Spark Connect session.
Directly passes the
- * DataFrame, so the user code actually runs with legacy DataFrame and
session..
+ * DataFrame, so the user code actually runs with legacy DataFrame and
session. The batch
+ * DataFrame is not cached in a SessionHolder: the Scala callback never
resolves it as a remote
+ * relation, so it needs neither the cloned SessionHolder nor the DataFrame
cache.
*/
- def scalaForeachBatchWrapper(
- payloadBytes: Array[Byte],
- sessionHolder: SessionHolder): ForeachBatchFnType = {
+ def scalaForeachBatchWrapper(payloadBytes: Array[Byte], sessionHolder:
SessionHolder)
+ : (ForeachBatchFnType, AutoCloseable, AtomicReference[String]) = {
val foreachBatchPkt =
UdfSerialization
.deserialize[ForeachWriterPacket](payloadBytes,
Utils.getContextOrSparkClassLoader)
val fn = foreachBatchPkt.foreachWriter.asInstanceOf[(Dataset[Any], Long)
=> Unit]
val encoder =
foreachBatchPkt.datasetEncoder.asInstanceOf[AgnosticEncoder[Any]]
- // TODO(SPARK-44462): Set up Spark Connect session.
- // Do we actually need this for the first version?
- dataFrameCachingWrapper(
- (args: FnArgsWithId) => {
- // dfId is not used, see hack comment above.
- try {
- val ds = if (AgnosticEncoders.UnboundRowEncoder == encoder) {
- // When the dataset is a DataFrame (Dataset[Row).
- args.df.asInstanceOf[Dataset[Any]]
- } else {
- // Recover the Dataset from the DataFrame using the encoder.
- args.df.as(encoder)
- }
- fn(ds, args.batchId)
- } catch {
- case t: Throwable =>
- logError(s"Calling foreachBatch fn failed", t)
- throw t
+ val wrappedFn: ForeachBatchFnType = (df: DataFrame, batchId: Long) => {
+ try {
+ val ds = if (AgnosticEncoders.UnboundRowEncoder == encoder) {
+ // When the dataset is a DataFrame (Dataset[Row]).
+ df.asInstanceOf[Dataset[Any]]
+ } else {
+ // Recover the Dataset from the DataFrame using the encoder.
+ df.as(encoder)
}
- },
- sessionHolder)
+ fn(ds, batchId)
+ } catch {
+ case t: Throwable =>
+ logError(s"Calling foreachBatch fn failed", t)
+ throw t
+ }
+ }
+ (wrappedFn, ForeachBatchCleaner(None, None), new AtomicReference[String]())
}
/**
* Starts up Python worker and initializes it with Python function. Returns
a foreachBatch
- * function that sets up the session and Dataframe cache and and interacts
with the Python
- * worker to execute user's function. In addition, it returns an
AutoClosable. The caller must
- * ensure it is closed so that worker process and related resources are
released.
+ * function that sets up the session and DataFrame cache and interacts with
the Python worker to
+ * execute user's function. In addition, it returns an AutoClosable and an
AtomicReference for
+ * setting the query id. The caller must ensure it is closed so that worker
process and related
+ * resources are released.
*/
- def pythonForeachBatchWrapper(
- pythonFn: SimplePythonFunction,
- sessionHolder: SessionHolder): (ForeachBatchFnType, AutoCloseable) = {
+ def pythonForeachBatchWrapper(pythonFn: SimplePythonFunction, sessionHolder:
SessionHolder)
+ : (ForeachBatchFnType, AutoCloseable, AtomicReference[String]) = {
val port = SparkConnectService.localPort
var connectUrl = s"sc://localhost:$port/;user_id=${sessionHolder.userId}"
@@ -158,18 +284,18 @@ object StreamingForeachBatchHelper extends Logging {
val (dataOut, dataIn) = runner.init()
- val foreachBatchRunnerFn: FnArgsWithId => Unit = (args: FnArgsWithId) => {
+ val useClonedSession = sessionHolder.session.sessionState.conf
+ .getConf(SQLConf.CONNECT_STREAMING_FOREACH_BATCH_USE_CLONED_SESSION)
+ val sessionManager = new ForeachBatchSessionManager(sessionHolder,
useClonedSession)
+ val queryIdRef = new AtomicReference[String]()
- // TODO(SPARK-44462): A new session id pointing to args.df.sparkSession
needs to be created.
- // This is because MicroBatch execution clones the session during
start.
- // The session attached to the foreachBatch dataframe is different
from the one the one
- // the query was started with. `sessionHolder` here contains the
latter.
- // Another issue with not creating new session id: foreachBatch
worker keeps
- // the session alive. The session mapping at Connect server does not
expire and query
- // keeps running even if the original client disappears. This keeps
the query running.
+ val foreachBatchRunnerFn: FnArgsWithId => Unit = (args: FnArgsWithId) => {
+ // The cloned session id lets the worker resolve CachedRemoteRelation
against the
+ // cloned session rather than the root one.
PythonWorkerUtils.writeUTF(args.dfId, dataOut)
dataOut.writeLong(args.batchId)
+ PythonWorkerUtils.writeUTF(args.clonedSessionId, dataOut)
dataOut.flush()
try {
@@ -205,7 +331,10 @@ object StreamingForeachBatchHelper extends Logging {
}
}
- (dataFrameCachingWrapper(foreachBatchRunnerFn, sessionHolder),
RunnerCleaner(runner))
+ (
+ dataFrameCachingWrapper(foreachBatchRunnerFn, sessionManager,
queryIdRef),
+ ForeachBatchCleaner(Some(runner), Some(sessionManager)),
+ queryIdRef)
}
/**
@@ -282,6 +411,13 @@ object StreamingForeachBatchHelper extends Logging {
cleanupStreamingRunner(key)
removeListenerIfRegistered()
}
+
+ // Handle the other miss the comment above describes: the query
terminated before this
+ // registration, so onQueryTerminated will never fire for it. Close
eagerly so the cloned
+ // SessionHolder (which never expires by inactivity) is not leaked.
+ if (!query.isActive) {
+ cleanupStreamingRunner(key)
+ }
}
/** Cleans up all the registered runners. */
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala
index e97bcb0c786e..b2175407c3e0 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala
@@ -131,6 +131,11 @@ case class SessionHolder(userId: String, sessionId:
String, session: SparkSessio
private[spark] lazy val dataFrameCache: ConcurrentMap[String, DataFrame] =
new ConcurrentHashMap()
+ // Mapping from streaming query ID to the active cached DataFrame ID for
that query.
+ // Used for foreachBatch to detect stale DataFrames from previous batches.
+ private[connect] lazy val dataFrameQueryIndex: ConcurrentMap[String, String]
=
+ new ConcurrentHashMap()
+
// ML model cache
@volatile private var mlCacheInitialized = false
private[connect] lazy val mlCache = {
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala
index c7558465a0b9..79e0acd2b5ea 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala
@@ -394,6 +394,49 @@ class SparkConnectSessionManager extends Logging {
}
}
+ /**
+ * Register an existing SparkSession under a new session ID. Unlike
[[cloneSession]], this wraps
+ * the *exact* SparkSession instance (no cloning). Used for foreachBatch to
register the
+ * stream-level SparkSession created by StreamExecution.
+ *
+ * The returned SessionHolder has `customInactiveTimeoutMs = Some(-1)`
(never expire by
+ * inactivity); its lifecycle is managed by the streaming query termination.
+ */
+ private[connect] def registerExistingSession(
+ userId: String,
+ newSessionId: String,
+ existingSession: SparkSession): SessionHolder = {
+
+ // Validate UUID format.
+ try {
+ UUID.fromString(newSessionId).toString
+ } catch {
+ case _: IllegalArgumentException =>
+ throw new SparkSQLException(
+ errorClass = "INVALID_HANDLE.FORMAT",
+ messageParameters = Map("handle" -> newSessionId))
+ }
+
+ val key = SessionKey(userId, newSessionId)
+
+ // Validate that session with this key has not been already closed.
+ if (closedSessionsCache.getIfPresent(key) != null) {
+ throw new SparkSQLException(
+ errorClass = "INVALID_HANDLE.SESSION_CLOSED",
+ messageParameters = Map("handle" -> newSessionId))
+ }
+
+ schedulePeriodicChecks()
+ sessionStore.computeIfAbsent(
+ key,
+ _ => {
+ val h = SessionHolder(userId, newSessionId, existingSession)
+ h.initializeSession()
+ h.setCustomInactiveTimeoutMs(Some(-1L)) // Never expire; query manages
lifetime.
+ h
+ })
+ }
+
/**
* Used for testing
*/
diff --git
a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala
index 6a9656e5d6f8..9cad09887f3d 100644
---
a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala
+++
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelperSuite.scala
@@ -25,6 +25,7 @@ import org.mockito.Mockito.verify
import org.mockito.Mockito.when
import org.scalatestplus.mockito.MockitoSugar
+import org.apache.spark.SparkIllegalStateException
import org.apache.spark.sql.connect.SparkConnectTestUtils
import org.apache.spark.sql.connect.service.SessionHolder
import org.apache.spark.sql.streaming.StreamingQuery
@@ -47,6 +48,7 @@ class StreamingForeachBatchHelperSuite extends
SharedSparkSession with MockitoSu
val (queryId, runId) = (UUID.randomUUID(), UUID.randomUUID())
when(query.id).thenReturn(queryId)
when(query.runId).thenReturn(runId)
+ when(query.isActive).thenReturn(true)
query
}
@@ -114,6 +116,60 @@ class StreamingForeachBatchHelperSuite extends
SharedSparkSession with MockitoSu
assert(spark.streams.listListeners().length == listenersBefore)
}
+ test("ForeachBatchSessionManager: close is terminal, and a closing root
session is refused") {
+ // The cloned holder is registered with customInactiveTimeoutMs = -1, so
it never expires by
+ // inactivity: once the cleaner ran, a later batch must not be able to
register another one.
+ // Use a child session, not the shared one -- closing the holder tears its
session down.
+ val batchDf = spark.newSession().range(1).toDF()
+
+ val manager = new StreamingForeachBatchHelper.ForeachBatchSessionManager(
+ SparkConnectTestUtils.createDummySessionHolder(spark),
+ useClonedSession = true)
+ val clonedHolder = manager.getOrCreateClonedSessionHolder(batchDf)
+ // The id is pinned for the whole query: every batch resolves to the same
holder.
+ assert(manager.getOrCreateClonedSessionHolder(batchDf) eq clonedHolder)
+
+ manager.close()
+ checkError(
+ exception =
+
intercept[SparkIllegalStateException](manager.getOrCreateClonedSessionHolder(batchDf)),
+ condition =
"SPARK_CONNECT_ILLEGAL_STATE.STREAM_LIFECYCLE_ALREADY_COMPLETED",
+ parameters = Map("operation" -> "getOrCreateClonedSessionHolder"))
+
+ // Same for a root session that is already closing: the query has no owner
to clean up after
+ // it, so registering a holder for it would strand that holder.
+ val closingRoot = unregisteredSessionHolder()
+ closingRoot.close()
+ assert(closingRoot.isClosing)
+ val managerOnClosingRoot =
+ new StreamingForeachBatchHelper.ForeachBatchSessionManager(
+ closingRoot,
+ useClonedSession = true)
+ checkError(
+ exception = intercept[SparkIllegalStateException](
+ managerOnClosingRoot.getOrCreateClonedSessionHolder(batchDf)),
+ condition =
"SPARK_CONNECT_ILLEGAL_STATE.SESSION_MANAGEMENT_SESSION_ALREADY_CLOSED",
+ parameters = Map("key" -> closingRoot.key.toString))
+ }
+
+ test("CleanerCache: a cleaner registered for an already terminated query is
cleaned up") {
+ // onQueryTerminated only fires for a query that terminates after the
registration, so a
+ // cleaner registered for a query that is already gone has no reaper left.
It holds the cloned
+ // SessionHolder, which never expires by inactivity, so it must be closed
here.
+ val cleaner = mock[AutoCloseable]
+ val query = mockQuery()
+ when(query.isActive).thenReturn(false)
+ val cache = new StreamingForeachBatchHelper.CleanerCache(
+ SparkConnectTestUtils.createDummySessionHolder(spark))
+
+ cache.registerCleanerForQuery(query, cleaner)
+
+ verify(cleaner, times(1)).close()
+ assert(cache.listEntriesForTesting().isEmpty)
+
+ cache.cleanUpAll() // Drops the listener this registration added.
+ }
+
test("CleanerCache.cleanUpAll unregisters the streaming listener") {
// close() does not remove the StreamingRunnerCleanerListener (it is not
tracked in the
// session's listenerCache), so cleanUpAll() must drop it; otherwise the
listener keeps the
diff --git
a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala
index cff5f345d257..f34f46101136 100644
---
a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala
+++
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala
@@ -38,7 +38,7 @@ import org.apache.spark.sql.connect.{PythonTestDepsChecker,
SparkConnectTestUtil
import org.apache.spark.sql.connect.common.InvalidPlanInput
import org.apache.spark.sql.connect.config.Connect
import org.apache.spark.sql.connect.planner.{PythonStreamingQueryListener,
SparkConnectPlanner, StreamingForeachBatchHelper}
-import
org.apache.spark.sql.connect.planner.StreamingForeachBatchHelper.RunnerCleaner
+import
org.apache.spark.sql.connect.planner.StreamingForeachBatchHelper.ForeachBatchCleaner
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.pipelines.graph.{DataflowGraph,
PipelineUpdateContextImpl}
import org.apache.spark.sql.pipelines.logging.PipelineEvent
@@ -345,9 +345,9 @@ class SparkConnectSessionHolderSuite extends
SharedSparkSession {
capturedServer = SparkConnectService.server
val pythonFn =
dummyPythonFunction(sessionHolder)(streamingForeachBatchFunction)
- val (fn1, cleaner1) =
+ val (fn1, cleaner1, _) =
StreamingForeachBatchHelper.pythonForeachBatchWrapper(pythonFn,
sessionHolder)
- val (fn2, cleaner2) =
+ val (fn2, cleaner2, _) =
StreamingForeachBatchHelper.pythonForeachBatchWrapper(pythonFn,
sessionHolder)
val query1 = spark.readStream
@@ -377,8 +377,9 @@ class SparkConnectSessionHolderSuite extends
SharedSparkSession {
// listeners we added so finally only removes ours, not a concurrent
attempt's.
ourNewListeners = spark.streams.listListeners().toSet --
baselineListeners
- val (runner1, runner2) =
- (cleaner1.asInstanceOf[RunnerCleaner].runner,
cleaner2.asInstanceOf[RunnerCleaner].runner)
+ val (runner1, runner2) = (
+ cleaner1.asInstanceOf[ForeachBatchCleaner].runner.get,
+ cleaner2.asInstanceOf[ForeachBatchCleaner].runner.get)
// assert both python processes are running
assert(!runner1.isWorkerStopped().get)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
