You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "LuciferYang (via GitHub)" <gi...@apache.org> on 2026/04/19 05:49:26 UTC
[PR] [SPARK-44462][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
LuciferYang opened a new pull request, #55410:
URL: https://github.com/apache/spark/pull/55410
### What changes were proposed in this pull request?
Two fixes for foreachBatch session handling in Spark Connect:
1. **DataFrame cache sanity check**: Track which DataFrame is cached per streaming query in `SessionHolder`. When caching a new batch DataFrame, detect and remove any stale one left over from a previous batch. Pass the query ID into the foreachBatch wrappers so it can be set after the query starts.
2. **Dedicated stream SessionHolder**: On the first batch, lazily create a new `SessionHolder` wrapping `StreamExecution`'s cloned SparkSession (instead of the original). Register it with `SparkConnectSessionManager` with no inactivity timeout since the streaming query manages its lifecycle. Clean up on query termination via `ForeachBatchCleaner`. On the Python side, each batch now receives `(dfId, batchId, streamSessionId)` from the server and the worker creates a session bound to the stream session ID.
### Why are the changes needed?
`SessionHolder` referenced the original SparkSession, not the `StreamExecution` clone (`sparkSessionForStream`) that batch DataFrames actually use:
```scala
// StreamExecution.scala
protected[sql] val sparkSessionForStream: SparkSession = sparkSession.cloneSession()
```
This caused: (1) batch DataFrames ran against the cloned session but `SessionHolder` pointed at the original, so session-level state was invisible to Connect; (2) the Python foreachBatch worker operated under the original session ID, keeping the parent session active and delaying cleanup after client disconnect; (3) nothing cleaned up stale cached DataFrames from a previous batch that did not exit cleanly.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Updated `SparkConnectSessionHolderSuite` to destructure the new 3-tuple return from `pythonForeachBatchWrapper`.
- Added `test_nested_dataframes` in `StreamingForeachBatchParityTests` -- exercises both closured and batch DataFrames inside foreachBatch, using `saveAsTable` for cross-session visibility.
- Pass Github Actions
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5901148009
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "xiaoxuandev (via GitHub)" <gi...@apache.org>.
xiaoxuandev commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3113806688
##########
python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py:
##########
@@ -67,11 +84,16 @@ def process(df_id, batch_id): # type: ignore[no-untyped-def]
print(f"{log_name} is starting with url {connect_url} and sessionId {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
+ # Create an initial SparkSession using the parent session id for initialization only.
+ # 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()
Review Comment:
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?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3114162351
##########
python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py:
##########
@@ -47,13 +48,29 @@ 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, stream_session_id): # type: ignore[no-untyped-def]
Review Comment:
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`.
##########
python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py:
##########
@@ -34,24 +34,32 @@ def test_streaming_foreach_batch_graceful_stop(self):
super().test_streaming_foreach_batch_graceful_stop()
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 stream 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")
Review Comment:
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
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, parentSessionHolder.sessionId)}] " +
+ log"Created stream SessionHolder with streamSessionId " +
+ log"${MDC(STREAM_ID, streamSessionId)} for foreachBatch.")
+ }
+ }
+ }
+ _streamSessionHolder
+ }
+
+ def close(): Unit = synchronized {
+ if (_streamSessionHolder != null) {
+ try {
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(_streamSessionHolder.userId, _streamSessionHolder.sessionId))
+ } catch {
+ case NonFatal(ex) =>
+ logWarning("Error closing stream SessionHolder for foreachBatch", ex)
+ }
+ _streamSessionHolder = null
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the stream SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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,
+ streamSessionId: 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 stream-level SessionHolder via sessionManager. Batch
+ * DataFrames are cached in the stream SessionHolder rather than the parent 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.getOrCreateStreamSessionHolder(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)}")
+
+ // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
Review Comment:
Why do we need this check? We are already removing the DF from the cache after the batch completes
##########
python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py:
##########
@@ -67,11 +84,16 @@ def process(df_id, batch_id): # type: ignore[no-untyped-def]
print(f"{log_name} is starting with url {connect_url} and sessionId {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
+ # Create an initial SparkSession using the parent session id for initialization only.
+ # 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()
Review Comment:
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
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
Review Comment:
Same as above: can we use `clonedSession` as naming to indicate that this is a cloned session?
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, parentSessionHolder.sessionId)}] " +
+ log"Created stream SessionHolder with streamSessionId " +
+ log"${MDC(STREAM_ID, streamSessionId)} for foreachBatch.")
+ }
+ }
+ }
+ _streamSessionHolder
+ }
+
+ def close(): Unit = synchronized {
+ if (_streamSessionHolder != null) {
+ try {
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(_streamSessionHolder.userId, _streamSessionHolder.sessionId))
+ } catch {
+ case NonFatal(ex) =>
+ logWarning("Error closing stream SessionHolder for foreachBatch", ex)
+ }
+ _streamSessionHolder = null
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the stream SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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,
+ streamSessionId: 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 stream-level SessionHolder via sessionManager. Batch
+ * DataFrames are cached in the stream SessionHolder rather than the parent 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.getOrCreateStreamSessionHolder(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)}")
+
+ // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
+ val queryId = queryIdRef.get()
+
+ effectiveHolder.cacheDataFrameById(dfId, df)
Review Comment:
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?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5927511180
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5882084396
This is a huge change and I don't feel like we should take the risk to make it happen in the maintenance release.
My big apologize for not looking into this in time. I'll take a look really soon (ideally by today).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5945949479
I was waiting for CI status but it was already succeeded; the result just wasn't reflected in this PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3120879798
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, parentSessionHolder.sessionId)}] " +
+ log"Created stream SessionHolder with streamSessionId " +
+ log"${MDC(STREAM_ID, streamSessionId)} for foreachBatch.")
+ }
+ }
+ }
+ _streamSessionHolder
+ }
+
+ def close(): Unit = synchronized {
+ if (_streamSessionHolder != null) {
+ try {
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(_streamSessionHolder.userId, _streamSessionHolder.sessionId))
+ } catch {
+ case NonFatal(ex) =>
+ logWarning("Error closing stream SessionHolder for foreachBatch", ex)
+ }
+ _streamSessionHolder = null
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the stream SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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,
+ streamSessionId: 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 stream-level SessionHolder via sessionManager. Batch
+ * DataFrames are cached in the stream SessionHolder rather than the parent 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.getOrCreateStreamSessionHolder(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)}")
+
+ // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
Review Comment:
I see, yeah I'm okay with it since it's fixing an existing TODO.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "juanpaulH (via GitHub)" <gi...@apache.org>.
juanpaulH commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-4972846101
Any tangible resolution or path forward for this critical issue?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5926094364
No problem, I didn't leave any comment after pushing [b0fad56](https://github.com/apache/spark/pull/55410/commits/b0fad5678caab38c1fe7d4f04a011bf1f5355b7a) 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR closed pull request #55410: [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect
URL: https://github.com/apache/spark/pull/55410
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "szehon-ho (via GitHub)" <gi...@apache.org>.
szehon-ho commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5881656496
Can we merge it to 4.2.1 as well to keep in sync with 4.3 : https://lists.apache.org/thread/bl3zbvbmrhz2m6jz25ogfwfd6v3wjo4q
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5926141294
Let's just have a SQL config and mention it in the migration guide. 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5508972299
<img width="1938" height="1236" alt="image" src="https://github.com/user-attachments/assets/a7283ad9-ecf5-4b8e-b323-f1651bfd4793" />
CI passed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5516916107
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3115540687
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
Review Comment:
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.
##########
python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py:
##########
@@ -67,11 +84,16 @@ def process(df_id, batch_id): # type: ignore[no-untyped-def]
print(f"{log_name} is starting with url {connect_url} and sessionId {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
+ # Create an initial SparkSession using the parent session id for initialization only.
+ # 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()
Review Comment:
+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.
##########
python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py:
##########
@@ -47,13 +48,29 @@ 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, stream_session_id): # type: ignore[no-untyped-def]
Review Comment:
Renamed to `cloned_session_id` / `root_session_id` (matches the Scala side: `clonedSession*` / `rootSessionHolder`).
##########
python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py:
##########
@@ -34,24 +34,32 @@ def test_streaming_foreach_batch_graceful_stop(self):
super().test_streaming_foreach_batch_graceful_stop()
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 stream 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")
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3120866981
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, parentSessionHolder.sessionId)}] " +
+ log"Created stream SessionHolder with streamSessionId " +
+ log"${MDC(STREAM_ID, streamSessionId)} for foreachBatch.")
+ }
+ }
+ }
+ _streamSessionHolder
+ }
+
+ def close(): Unit = synchronized {
+ if (_streamSessionHolder != null) {
+ try {
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(_streamSessionHolder.userId, _streamSessionHolder.sessionId))
+ } catch {
+ case NonFatal(ex) =>
+ logWarning("Error closing stream SessionHolder for foreachBatch", ex)
+ }
+ _streamSessionHolder = null
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the stream SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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,
+ streamSessionId: 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 stream-level SessionHolder via sessionManager. Batch
+ * DataFrames are cached in the stream SessionHolder rather than the parent 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.getOrCreateStreamSessionHolder(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)}")
+
+ // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
+ val queryId = queryIdRef.get()
+
+ effectiveHolder.cacheDataFrameById(dfId, df)
Review Comment:
Sure, I'm okay with a follow up Jira, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5925623046
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. Either land this partial fix with disabling by default and aim for better fix in future, or hold this fix till we address them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3115541819
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, parentSessionHolder.sessionId)}] " +
+ log"Created stream SessionHolder with streamSessionId " +
+ log"${MDC(STREAM_ID, streamSessionId)} for foreachBatch.")
+ }
+ }
+ }
+ _streamSessionHolder
+ }
+
+ def close(): Unit = synchronized {
+ if (_streamSessionHolder != null) {
+ try {
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(_streamSessionHolder.userId, _streamSessionHolder.sessionId))
+ } catch {
+ case NonFatal(ex) =>
+ logWarning("Error closing stream SessionHolder for foreachBatch", ex)
+ }
+ _streamSessionHolder = null
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the stream SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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,
+ streamSessionId: 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 stream-level SessionHolder via sessionManager. Batch
+ * DataFrames are cached in the stream SessionHolder rather than the parent 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.getOrCreateStreamSessionHolder(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)}")
+
+ // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
Review Comment:
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.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, parentSessionHolder.sessionId)}] " +
+ log"Created stream SessionHolder with streamSessionId " +
+ log"${MDC(STREAM_ID, streamSessionId)} for foreachBatch.")
+ }
+ }
+ }
+ _streamSessionHolder
+ }
+
+ def close(): Unit = synchronized {
+ if (_streamSessionHolder != null) {
+ try {
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(_streamSessionHolder.userId, _streamSessionHolder.sessionId))
+ } catch {
+ case NonFatal(ex) =>
+ logWarning("Error closing stream SessionHolder for foreachBatch", ex)
+ }
+ _streamSessionHolder = null
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the stream SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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,
+ streamSessionId: 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 stream-level SessionHolder via sessionManager. Batch
+ * DataFrames are cached in the stream SessionHolder rather than the parent 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.getOrCreateStreamSessionHolder(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)}")
+
+ // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
+ val queryId = queryIdRef.get()
+
+ effectiveHolder.cacheDataFrameById(dfId, df)
Review Comment:
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?
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
Review Comment:
Done — renamed `streamSession*` → `clonedSession*` and the ctor param `parentSessionHolder` → `rootSessionHolder` so Scala and Python stay symmetric.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-4277015797
@heyihong @HeartSaVioR FYI
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-4292639328
Thank you @bogao007 ~
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5945981554
@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!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5946666271
Thanks for your reviews. @HeartSaVioR @HyukjinKwon @uros-b @bogao007 @xiaoxuandev
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5946648320
> @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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5519435665
> 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5509321580
@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:
```python
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`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "xiaoxuandev (via GitHub)" <gi...@apache.org>.
xiaoxuandev commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3113727100
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r3120894643
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala:
##########
@@ -383,6 +383,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(
Review Comment:
@hvanhovell could you help review the changes here? Thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5493181742
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #55410:
URL: https://github.com/apache/spark/pull/55410#discussion_r4139373086
##########
python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py:
##########
@@ -43,24 +47,75 @@ def func(batch_df, _):
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()
Review Comment:
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.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -47,66 +48,179 @@ 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).
+ */
+ private[connect] class ForeachBatchSessionManager(rootSessionHolder: SessionHolder)
+ 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 (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(STREAM_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(STREAM_ID, holder.sessionId)}.")
Review Comment:
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.
##########
python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py:
##########
@@ -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
Review Comment:
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.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -46,66 +47,143 @@ 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 {
+ /**
+ * Manages a stream-level SessionHolder for foreachBatch. On the first batch invocation, lazily
+ * creates a new SessionHolder wrapping the batch DataFrame's SparkSession (the stream session
+ * cloned by StreamExecution). The stream SessionHolder is registered with
+ * SparkConnectSessionManager so the Python worker can resolve CachedRemoteRelation against it.
+ */
+ private[connect] class ForeachBatchSessionManager(parentSessionHolder: SessionHolder)
+ extends Logging {
+ @volatile private var _streamSessionHolder: SessionHolder = null
+
+ def getOrCreateStreamSessionHolder(batchDf: DataFrame): SessionHolder = {
+ if (_streamSessionHolder == null) {
+ synchronized {
+ if (_streamSessionHolder == null) {
+ val streamSession = batchDf.sparkSession
+ .asInstanceOf[org.apache.spark.sql.classic.SparkSession]
+ val streamSessionId = UUID.randomUUID().toString
+ _streamSessionHolder = SparkConnectService.sessionManager
+ .registerExistingSession(parentSessionHolder.userId, streamSessionId, streamSession)
+ logInfo(
+ log"[session: ${MDC(SESSION_ID, parentSessionHolder.sessionId)}] " +
+ log"Created stream SessionHolder with streamSessionId " +
+ log"${MDC(STREAM_ID, streamSessionId)} for foreachBatch.")
+ }
+ }
+ }
+ _streamSessionHolder
+ }
+
+ def close(): Unit = synchronized {
+ if (_streamSessionHolder != null) {
+ try {
+ SparkConnectService.sessionManager.closeSession(
+ SessionKey(_streamSessionHolder.userId, _streamSessionHolder.sessionId))
+ } catch {
+ case NonFatal(ex) =>
+ logWarning("Error closing stream SessionHolder for foreachBatch", ex)
+ }
+ _streamSessionHolder = null
+ }
+ }
+ }
+
+ /**
+ * Composite cleaner that closes both the Python runner (if any) and the stream SessionHolder.
+ * Registered with CleanerCache to clean up on query termination.
+ */
+ private[connect] case class ForeachBatchCleaner(
+ runner: Option[StreamingPythonRunner],
+ sessionManager: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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,
+ streamSessionId: 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 stream-level SessionHolder via sessionManager. Batch
+ * DataFrames are cached in the stream SessionHolder rather than the parent 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.getOrCreateStreamSessionHolder(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)}")
+
+ // Sanity check: remove any stale DataFrame left over from a previous batch for this query.
+ val queryId = queryIdRef.get()
+
+ effectiveHolder.cacheDataFrameById(dfId, df)
Review Comment:
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: https://github.com/HeartSaVioR/spark/pull/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?
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/StreamingForeachBatchHelper.scala:
##########
@@ -47,66 +48,179 @@ 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).
+ */
+ private[connect] class ForeachBatchSessionManager(rootSessionHolder: SessionHolder)
+ 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 (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(STREAM_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(STREAM_ID, holder.sessionId)}.")
+ } catch {
+ case NonFatal(ex) =>
+ logWarning(
+ log"[rootSession: ${MDC(SESSION_ID, rootSessionHolder.sessionId)}] " +
+ log"Error closing cloned SessionHolder ${MDC(STREAM_ID, holder.sessionId)} " +
+ log"for foreachBatch; 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: ForeachBatchSessionManager)
+ extends AutoCloseable {
override def close(): Unit = {
- try runner.stop()
+ runner.foreach { r =>
+ try r.stop()
+ catch {
+ case NonFatal(ex) =>
+ logWarning("Error while stopping streaming Python worker", ex)
+ }
+ }
+ try sessionManager.close()
catch {
- case NonFatal(ex) => // Exception is not propagated.
- logWarning("Error while stopping streaming Python worker", ex)
+ 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)
+ }
+ }
}
- }
}
/**
* Handles setting up Scala remote session and other Spark Connect environment and then runs the
* 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. However, batch
+ * DataFrames are still cached in the cloned SessionHolder for correct session tracking.
*/
- def scalaForeachBatchWrapper(
- payloadBytes: Array[Byte],
- sessionHolder: SessionHolder): ForeachBatchFnType = {
+ def scalaForeachBatchWrapper(payloadBytes: Array[Byte], sessionHolder: SessionHolder)
+ : (ForeachBatchFnType, AutoCloseable, AtomicReference[String]) = {
val foreachBatchPkt =
Utils.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(
+ val sessionManager = new ForeachBatchSessionManager(sessionHolder)
+ val queryIdRef = new AtomicReference[String]()
+ val wrappedFn = dataFrameCachingWrapper(
Review Comment:
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?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5945951285
Thanks! Merging to master/4.x/4.3.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5925825897
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-44462][SS][CONNECT] Fix the session passed to foreachBatch in Spark Connect [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #55410:
URL: https://github.com/apache/spark/pull/55410#issuecomment-5925992946
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org