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

common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java

@@ -102,6 +102,7 @@ public enum LogKeys implements LogKey {
   CLASS_PATHS,
   CLAUSES,
   CLEANUP_LOCAL_DIRS,
+  CLONED_SESSION_ID,
   CLUSTER_CENTROIDS,
   CLUSTER_ID,
   CLUSTER_LABEL,

docs/streaming/ss-migration-guide.md

@@ -23,6 +23,10 @@ Note that this migration guide describes the items specific to Structured Stream
 Many items of SQL migration can be applied when migrating Structured Streaming to higher versions.
 Please refer [Migration Guide: SQL, Datasets and DataFrame](../sql-migration-guide.html).
 
+## Upgrading from Structured Streaming 4.2 to 4.3
+
+- Since Spark 4.3, in Spark Connect the DataFrame passed to a Python `foreachBatch` function is bound to the streaming query's own session, a clone of the session that started the query, matching classic `foreachBatch`. This runs the batch under the configuration the streaming engine pins on the clone rather than the root session's. As a result, the batch DataFrame can no longer be combined with the root session inside the function: a temporary view created from it is not visible to the root session, referencing it from `spark.sql` on the root session fails with `DATAFRAME_NOT_FOUND`, and joining it with a DataFrame captured from the root session fails with `SESSION_NOT_SAME`. Run such work through the batch DataFrame's own session (`batch_df.sparkSession`) instead, and pass results out through a table or a global temporary view rather than a session-local temporary view. To restore the previous behavior, set `spark.sql.connect.streaming.foreachBatch.useClonedSession` to `false`.
+
 ## Upgrading from Structured Streaming 4.1 to 4.2
 
 - Since Spark 4.2, restarting a streaming query from a checkpoint whose metadata file is missing while the offset or commit logs contain data fails with `STREAMING_CHECKPOINT_MISSING_METADATA_FILE`, instead of silently generating a new query ID (which can duplicate data in exactly-once sinks). Restore the metadata file or use a new checkpoint location. To restore the previous behavior, set `spark.sql.streaming.checkpoint.verifyMetadataExists.enabled` to `false`. (See [SPARK-55058](https://issues.apache.org/jira/browse/SPARK-55058) for more details.)

python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py

@@ -21,19 +21,19 @@
 """
 
 import os
+from typing import IO
 
-from pyspark.worker_util import get_sock_file_to_executor
+from pyspark import worker
+from pyspark.errors import PySparkAssertionError
 from pyspark.serializers import (
-    write_int,
-    read_long,
-    UTF8Deserializer,
     CPickleSerializer,
+    UTF8Deserializer,
+    read_long,
+    write_int,
 )
-from pyspark import worker
 from pyspark.sql.connect.session import SparkSession
 from pyspark.util import handle_worker_exception
-from typing import IO
-from pyspark.worker_util import check_python_version
+from pyspark.worker_util import check_python_version, get_sock_file_to_executor
 
 pickle_ser = CPickleSerializer()
 utf8_deserializer = UTF8Deserializer()
@@ -47,13 +47,32 @@ def main(infile: IO, outfile: IO) -> None:
 
     log_name = "Streaming ForeachBatch worker"
 
-    def process(df_id, batch_id):  # type: ignore[no-untyped-def]
+    def process(df_id, batch_id, cloned_session_id):  # type: ignore[no-untyped-def]
         global spark
-        print(f"{log_name} Started batch {batch_id} with DF id {df_id} and session id {session_id}")
-        batch_df = spark_connect_session._create_remote_dataframe(df_id)
+        # Lazily create or switch the SparkSession when the cloned session changes.
+        # Use create() instead of getOrCreate() because the latter returns the cached
+        # active/default session, ignoring the cloned session ID in the URL.
+        if spark is None or spark.session_id != cloned_session_id:
+            cloned_url = connect_url + ";session_id=" + cloned_session_id
+            spark = SparkSession.builder.remote(cloned_url).create()
+            # create() fills the default/active slots only when they are empty, and the bootstrap
+            # session below already claimed them. Point them at the cloned session so that
+            # SparkSession.active() and getActiveSession() inside the user function return the
+            # session the batch actually runs on. This has to happen after the user function was
+            # unpickled, which requires the active session to still be the root one.
+            with SparkSession._lock:
+                SparkSession._default_session = spark
+            SparkSession._active_session.session = spark
+            print(f"{log_name} Created new session for cloned_session_id {cloned_session_id}")
+        print(
+            f"{log_name} Started batch {batch_id} with DF id {df_id} "
+            f"and session id {cloned_session_id}"
+        )
+        batch_df = spark._create_remote_dataframe(df_id)
         func(batch_df, batch_id)
         print(
-            f"{log_name} Completed batch {batch_id} with DF id {df_id} and session id {session_id}"
+            f"{log_name} Completed batch {batch_id} with DF id {df_id} "
+            f"and session id {cloned_session_id}"
         )
 
     try:
@@ -63,15 +82,20 @@ def process(df_id, batch_id):  # type: ignore[no-untyped-def]
         os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1"
 
         connect_url = os.environ["SPARK_CONNECT_LOCAL_URL"]
-        session_id = utf8_deserializer.loads(infile)
+        root_session_id = utf8_deserializer.loads(infile)
 
-        print(f"{log_name} is starting with url {connect_url} and sessionId {session_id}.")
+        print(f"{log_name} is starting with url {connect_url} and sessionId {root_session_id}.")
 
-        # To attach to the existing SparkSession, we're setting the session_id in the URL.
-        connect_url = connect_url + ";session_id=" + session_id
-        spark_connect_session = SparkSession.builder.remote(connect_url).getOrCreate()
-        assert spark_connect_session.session_id == session_id
-        spark = spark_connect_session
+        # Bootstrap SparkSession on the root session id: needed to unpickle the user's
+        # foreachBatch function before any batch (and the cloned session id) arrives.
+        # Per-batch work runs against the cloned session created in `process()`.
+        connect_url_init = connect_url + ";session_id=" + root_session_id
+        spark_connect_session = SparkSession.builder.remote(connect_url_init).getOrCreate()
+        if spark_connect_session.session_id != root_session_id:
+            raise PySparkAssertionError(
+                f"Root session ID mismatch: expected {root_session_id}, "
+                f"got {spark_connect_session.session_id}"
+            )
 
         func = worker.read_command(pickle_ser, infile)
         write_int(0, outfile)
@@ -80,9 +104,10 @@ def process(df_id, batch_id):  # type: ignore[no-untyped-def]
         while True:
             df_ref_id = utf8_deserializer.loads(infile)
             batch_id = read_long(infile)
-            # Handle errors inside Python worker. Write 0 to outfile if no errors and write -2 with
-            # traceback string if error occurs.
-            process(df_ref_id, int(batch_id))
+            cloned_session_id = utf8_deserializer.loads(infile)
+            # Handle errors inside Python worker. Write 0 to outfile if no errors and write -2
+            # with traceback string if error occurs.
+            process(df_ref_id, int(batch_id), cloned_session_id)
             write_int(0, outfile)
             outfile.flush()
     except Exception as e:

python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py

@@ -16,13 +16,18 @@
 #
 
 import time
+import uuid
+
+from pyspark.errors import PySparkPicklingError, StreamingQueryException
 from pyspark.sql.tests.streaming.test_streaming_foreach_batch import StreamingTestsForeachBatchMixin
 from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect
 from pyspark.testing.utils import eventually, timeout
-from pyspark.errors import PySparkPicklingError
 
 if should_test_connect:
-    from pyspark.errors.exceptions.connect import StreamingPythonRunnerInitializationException
+    from pyspark.errors.exceptions.connect import (
+        AnalysisException,
+        StreamingPythonRunnerInitializationException,
+    )
 
 
 class StreamingForeachBatchParityTests(StreamingTestsForeachBatchMixin, ReusedConnectTestCase):
@@ -42,25 +47,225 @@ 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()
+        finally:
+            if q:
+                q.stop()
+            self.spark.sql(f"DROP TABLE IF EXISTS {sentinel_table}")
+            # Best-effort cleanup in case a regression ever leaks the view.
+            try:
+                self.spark.catalog.dropTempView(view_name)
+            except Exception:
+                pass
+
+    def test_batch_uses_stream_session_for_stateful_aqe(self):
+        # SPARK-44462: the batch runs on the stream's cloned session, which pins the stateful
+        # configs (AQE off) independently of the root session. Proves the isolation this PR adds.
+        root = self.spark
+
+        def process(batch_df, _):
+            batch_count = batch_df.count()
+            root_aqe = root.conf.get("spark.sql.adaptive.enabled")
+            batch_aqe = batch_df.sparkSession.conf.get("spark.sql.adaptive.enabled")
+            same_session = root.session_id == batch_df.sparkSession.session_id
+            if same_session or root_aqe != "true" or batch_aqe != "false":
+                raise AssertionError(
+                    f"batch_count={batch_count}, same_session={same_session}, "
+                    f"AQE root={root_aqe}, batch={batch_aqe}"
+                )
+
+        q = None
+        with self.sql_conf({"spark.sql.adaptive.enabled": "true"}):
+            try:
+                df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+                q = (
+                    df.groupBy("value")
+                    .count()
+                    .writeStream.outputMode("complete")
+                    .foreachBatch(process)
+                    .start()
+                )
+                q.processAllAvailable()
+                self.assertIsNone(q.exception())
+            finally:
+                if q:
+                    q.stop()
+
+    def test_batch_dataframe_in_root_session_sql_is_unsupported(self):
+        # SPARK-44462 contract: the batch DataFrame is bound to the stream's cloned session, so its
+        # cached relation cannot be resolved through a captured root session. Documented behavior.
+        root = self.spark
+
+        def process(batch_df, _):
+            root.sql("SELECT value FROM {batch}", batch=batch_df).collect()
+
+        q = None
+        try:
+            df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+            q = df.writeStream.foreachBatch(process).start()
+            self.assertRaises(StreamingQueryException, q.processAllAvailable)
+            self.assertIn("DATAFRAME_NOT_FOUND", str(q.exception()))
         finally:
             if q:
                 q.stop()
 
+    def test_batch_dataframe_join_with_root_session_dataframe_is_unsupported(self):
+        # SPARK-44462 contract: the batch DataFrame (cloned session) cannot be combined with a
+        # DataFrame captured from the root session; the client rejects the cross-session join.
+        lookup = self.spark.range(2).selectExpr(
+            "CASE id WHEN 0 THEN 'hello' ELSE 'this' END AS value"
+        )
+
+        def process(batch_df, _):
+            batch_df.join(lookup, "value").collect()
+
+        q = None
+        try:
+            df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+            q = df.writeStream.foreachBatch(process).start()
+            self.assertRaises(StreamingQueryException, q.processAllAvailable)
+            self.assertIn("SESSION_NOT_SAME", str(q.exception()))
+        finally:
+            if q:
+                q.stop()
+
+    def test_batch_dataframe_in_batch_session_sql_is_supported(self):
+        # SPARK-44462 migration path: resolve the batch DataFrame through its own (cloned) session
+        # instead of a captured root session.
+        def process(batch_df, _):
+            expected = sorted(row.value for row in batch_df.select("value").collect())
+            actual = sorted(
+                row.value
+                for row in batch_df.sparkSession.sql(
+                    "SELECT value FROM {batch}", batch=batch_df
+                ).collect()
+            )
+            if actual != expected:
+                raise AssertionError(f"SQL rows {actual} != batch rows {expected}")
+
+        q = None
+        try:
+            df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+            q = df.writeStream.foreachBatch(process).start()
+            q.processAllAvailable()
+            self.assertTrue(any(p.numInputRows > 0 for p in q.recentProgress))
+            self.assertIsNone(q.exception())
+        finally:
+            if q:
+                q.stop()
+
+    def test_batch_dataframe_join_in_batch_session_is_supported(self):
+        # SPARK-44462 migration path: rebuild the lookup on the batch's own (cloned) session so the
+        # join stays within a single session.
+        def process(batch_df, _):
+            spark = batch_df.sparkSession
+            lookup = spark.range(2).selectExpr(
+                "CASE id WHEN 0 THEN 'hello' ELSE 'this' END AS value"
+            )
+            expected = sorted(row.value for row in batch_df.select("value").collect())
+            actual = sorted(
+                row.value for row in batch_df.join(lookup, "value").select("value").collect()
+            )
+            if actual != expected:
+                raise AssertionError(f"Joined rows {actual} != batch rows {expected}")
+
+        q = None
+        try:
+            df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+            q = df.writeStream.foreachBatch(process).start()
+            q.processAllAvailable()
+            self.assertTrue(any(p.numInputRows > 0 for p in q.recentProgress))
+            self.assertIsNone(q.exception())
+        finally:
+            if q:
+                q.stop()
+
+    def test_disabling_fix_restores_root_session_behavior(self):
+        # With the fix disabled, the batch DataFrame is bound to the root session again, so it can
+        # be combined with a DataFrame captured from the root session.
+        with self.sql_conf({"spark.sql.connect.streaming.foreachBatch.useClonedSession": "false"}):
+            lookup = self.spark.range(2).selectExpr(
+                "CASE id WHEN 0 THEN 'hello' ELSE 'this' END AS value"
+            )
+
+            def process(batch_df, _):
+                expected = sorted(row.value for row in batch_df.select("value").collect())
+                actual = sorted(
+                    row.value for row in batch_df.join(lookup, "value").select("value").collect()
+                )
+                if actual != expected:
+                    raise AssertionError(f"Joined rows {actual} != batch rows {expected}")
+
+            q = None
+            try:
+                df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
+                q = df.writeStream.foreachBatch(process).start()
+                q.processAllAvailable()
+                self.assertTrue(any(p.numInputRows > 0 for p in q.recentProgress))
+                self.assertIsNone(q.exception())
+            finally:
+                if q:
+                    q.stop()
+
     def test_pickling_error(self):
         class NoPickle:
             def __reduce__(self):