Spark Structured Streaming foreachBatch触发STREAMING_CONNECT_SERIALIZATION_ERROR

在Databricks中使用Spark Structured Streaming时，自定义Reader类的read_events函数单独执行可正常返回结果，但在foreachBatch操作中调用该函数处理微批数据时，触发了STREAMING_CONNECT_SERIALIZATION_ERROR错误。

Reader类代码（含转换函数） #

from pyspark.sql import DataFrame
from pyspark.sql.functions import col

class Reader():

    def __init__(self, spark):
        self.spark = spark
        self.events_table = "events_table"

    def read_events(self, transactions_df):
        events_df = self.spark.read.table(self.events_table)

        result_df = (
            transactions_df.alias("T")
            .join(
                events_df.alias("E"),
                (col("E.col1") == col("T.col1"))
            )
            .select(
                col("E.col1"),
                col("E.col2"),
            )
        )

        return result_df

单独执行验证代码 #

input_df = spark.read.table("transactions_table")
Reader(spark).read_events(input_df).display()

原Streaming类代码 #

from pyspark.sql import DataFrame

class Streaming():

    def __init__(self, spark):
        self.spark = spark
        self.transactions_table= "transactions_table"
        self.transactions_table_stream_checkpoint= "checkpoint_path"
        self.reader= Reader(spark)

    def process_batch_of_messages(self, df, batch_id):
        result_df = self.reader.read_events(df)
        print(f"For batch {batch_id} we have {result_df.count()} records.")

    def launch(self):
        (
            self.spark.readStream.format("delta")
            .option("skipChangeCommits", "true")
            .table(self.transactions_table)
            .writeStream.option("checkpointLocation", self.transactions_table_stream_checkpoint)
            .foreachBatch(
                lambda transactions, batch_id: self.process_batch_of_messages(
                    df=transactions, batch_id=batch_id
                )
            )
            .start()
        )


def entrypoint():
    stream = Streaming(spark)
    stream.launch()


if __name__ == "__main__":
    entrypoint()

触发的错误信息 #

STREAMING_CONNECT_SERIALIZATION_ERROR] Cannot serialize the function `foreachBatch`. If you accessed the Spark session, or a DataFrame defined outside of the function, or any object that contains a Spark session, please be aware that they are not allowed in Spark Connect. For `foreachBatch`, please access the Spark session using `df.sparkSession`, where `df` is the first parameter in your `foreachBatch` function. For `StreamingQueryListener`, please access the Spark session using `self.spark`. For details please check out the PySpark doc for `foreachBatch` and `StreamingQueryListener`.

根因分析 #

• foreachBatch的处理逻辑需要被序列化后发送到Executor节点执行，但原代码中Reader实例持有了初始化时传入的Spark Session对象，而Spark Session无法被序列化（在Spark Connect模式下明确禁止外部传入Session）。
• 单独执行时逻辑在Driver端直接运行，不需要序列化分发，因此可以正常执行；但流式处理中foreachBatch的函数会被序列化，持有Session的Reader实例无法完成序列化，导致报错。

解决方案 #

修改要点 #

1. 移除Reader类初始化时的Spark Session参数，改为在read_events方法中通过传入的DataFrame获取Session（transactions_df.sparkSession），这是Spark官方推荐的在foreachBatch中获取Session的方式。
2. 调整Streaming类中Reader的初始化逻辑，不再传入Spark Session。

修改后的完整代码 #

修改后的Reader类

from pyspark.sql import DataFrame
from pyspark.sql.functions import col

class Reader():

    def __init__(self):
        self.events_table = "events_table"

    def read_events(self, transactions_df):
        # 从当前批次的DataFrame获取Spark Session
        spark = transactions_df.sparkSession
        events_df = spark.read.table(self.events_table)

        result_df = (
            transactions_df.alias("T")
            .join(
                events_df.alias("E"),
                (col("E.col1") == col("T.col1"))
            )
            .select(
                col("E.col1"),
                col("E.col2"),
            )
        )

        return result_df

修改后的Streaming类

from pyspark.sql import DataFrame

class Streaming():

    def __init__(self, spark):
        self.spark = spark
        self.transactions_table= "transactions_table"
        self.transactions_table_stream_checkpoint= "checkpoint_path"
        # 初始化Reader时不再传入Spark Session
        self.reader= Reader()

    def process_batch_of_messages(self, df, batch_id):
        result_df = self.reader.read_events(df)
        print(f"For batch {batch_id} we have {result_df.count()} records.")

    def launch(self):
        (
            self.spark.readStream.format("delta")
            .option("skipChangeCommits", "true")
            .table(self.transactions_table)
            .writeStream.option("checkpointLocation", self.transactions_table_stream_checkpoint)
            .foreachBatch(
                lambda transactions, batch_id: self.process_batch_of_messages(
                    df=transactions, batch_id=batch_id
                )
            )
            .start()
        )


def entrypoint():
    stream = Streaming(spark)
    stream.launch()


if __name__ == "__main__":
    entrypoint()

内容的提问来源于stack exchange，提问作者wbaeckelmans