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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
