Databricks用foreachBatch流式处理Event Hub XML存CSV报Python Proxy错误

AI 小北

2026-8-28

问题背景

我编写代码用于从Azure Event Hub检索XML文件，提取部分标签内容后保存为CSV文件。由于需要将不同标签值存储为独立文件，我使用df.writeStream的foreachbatch方法实现分文件提取与存储逻辑，以下为运行环境、实现代码以及在Azure Databricks上执行时返回的错误信息。

运行环境

Databricks Runtime版本： 10.4 LTS
Apache Spark 3.2.1
Scala 2.12
Maven引入的Event Hubs库： com.microsoft.azure:azure-eventhubs-spark_2.12:2.3.22

问题实现代码

# Databricks notebook source
import lxml.etree as ET
import pyspark.sql.types as T
from os.path import dirname, join


# 定义XML文件中的命名空间，用于从默认（""）命名空间或指定命名空间选取元素
namespaces = {
    "": "http://www.fpml.org/FpML-5/reporting",
    "xsi": "http://www.w3.org/2001/XMLSchema-instance"
}

# trade date **********************************
trade_header = T.StructType([
    T.StructField("messageId", T.StringType(), False),
    T.StructField("tradeDate", T.StringType(), False)
])


def to_xml_message_trade_date(xml_string):
    root = ET.fromstring(xml_string)
    messageId = root.find(".//messageId", namespaces).text
    tradeDate = root.find(".//tradeDate", namespaces).text
    return [messageId, tradeDate]

extract_udf = udf(to_xml_message_trade_date, trade_header)
**********************************************

connectionString = "Endpoint=sb://xxxxxx.servicebus.windows.net/;SharedAccessKeyName=xxxx;SharedAccessKey=xxxxxxx;EntityPath=xxxxx"

ehConf = {
  'eventhubs.connectionString' : sc._jvm.org.apache.spark.eventhubs.EventHubsUtils.encrypt(connectionString)
}

stream_data = spark \
  .readStream \
  .format("eventhubs") \
  .options(**ehConf) \
  .option('multiLine', True) \
  .option('mode', 'PERMISSIVE') \
  .load()

df_str = stream_data.withColumn("data", stream_data["body"].cast("string"))


def write2csv(df, epoch_id):
     df.persist()
     df_tuples = df.select(extract_udf("data").alias("extracted_data"))
     df_parsed = df_tuples.select("extracted_data.*")

     df_parsed \
       .write \
       .format("csv") \
       .mode(SaveMode.Append) \
       .option("header", True) \
       .save("dbfs:/FileStore/Incoming/trade_date/")

     df.unpersist()
 

query = df_str \
    .writeStream \
    .outputMode("append") \
    .foreachBatch(write2csv) \
    .trigger(processingTime="1 seconds") \
    .start()

query.awaitTermination()

运行报错信息

StreamingQueryException: An exception was raised by the Python Proxy. Return Message: Traceback (most recent call last):

StreamingQueryException                   Traceback (most recent call last)
<command-1879221600357983> in <module>
      6     .start()
      7 
----> 8 query.awaitTermination()
      9 
     10 #     .format("csv") \

/databricks/spark/python/pyspark/sql/streaming.py in awaitTermination(self, timeout)
    101             return self._jsq.awaitTermination(int(timeout * 1000))
    102         else:
--> 103             return self._jsq.awaitTermination()
    104 
    105     @property

/databricks/spark/python/lib/py4j-0.10.9-src.zip/py4j/java_gateway.py in __call__(self, *args)
   1303         answer = self.gateway_client.send_command(command)
   1304         return_value = get_return_value(
-> 1305             answer, self.gateway_client, self.target_id, self.name)

正常运行的参考代码

以下代码可以正常流式读取并将标签值保存为单个文件，仅在使用foreachbatch实现分文件存储时出现上述问题：

df_trade_date \
  .writeStream \
  .format("csv") \
  .trigger(processingTime="30 seconds") \
  .option("checkpointLocation", "dbfs:/FileStore/checkpoint/") \
  .option("path", "dbfs:/FileStore/Incoming/trade_date/") \
  .option("header", True) \
  .outputMode("append") \
  .start() \
  .awaitTermination()

问题

当前实现遗漏了哪些配置？是否有对应的修复建议？

问题原因与修复方案

代码存在3个核心问题，逐个修复即可正常运行：

缺失流查询必填的checkpointLocation配置
所有结构化流作业，无论是否使用foreachBatch，都必须配置checkpoint目录存储消费偏移量、批次执行状态，否则作业启动直接失败。你贴的可正常运行的代码里已经配置了该参数，但foreachBatch版本中完全遗漏。
缺少必要依赖导入
代码中使用了SaveMode，但没有导入pyspark.sql.SaveMode，执行到写入步骤时会直接抛出NameError，这是触发Python Proxy报错的直接原因之一。另外原代码没有显式导入udf，部分Databricks版本虽然会自动导入，但显式导入可以避免环境差异导致的异常。
UDF调用位置错误触发序列化问题
驱动节点定义的UDF直接放在foreachBatch函数内给executor侧执行，很容易出现序列化失败。最优方案是提前在流DataFrame上完成所有解析转换，把解析完成的结构化DataFrame传给foreachBatch，不要在批次写入逻辑中做UDF计算。

修复后可运行代码示例

# 补全缺失导入
import lxml.etree as ET
import pyspark.sql.types as T
from pyspark.sql.functions import udf
from pyspark.sql import SaveMode

# 提前完成数据解析，不要在foreachBatch内部做UDF转换
namespaces = {
    "": "http://www.fpml.org/FpML-5/reporting",
    "xsi": "http://www.w3.org/2001/XMLSchema-instance"
}

trade_header = T.StructType([
    T.StructField("messageId", T.StringType(), False),
    T.StructField("tradeDate", T.StringType(), False)
])


def to_xml_message_trade_date(xml_string):
    root = ET.fromstring(xml_string)
    messageId = root.find(".//messageId", namespaces).text
    tradeDate = root.find(".//tradeDate", namespaces).text
    return [messageId, tradeDate]

extract_udf = udf(to_xml_message_trade_date, trade_header)

connectionString = "Endpoint=sb://xxxxxx.servicebus.windows.net/;SharedAccessKeyName=xxxx;SharedAccessKey=xxxxxxx;EntityPath=xxxxx"

ehConf = {
  'eventhubs.connectionString' : sc._jvm.org.apache.spark.eventhubs.EventHubsUtils.encrypt(connectionString)
}

stream_data = spark \
  .readStream \
  .format("eventhubs") \
  .options(**ehConf) \
  .option('multiLine', True) \
  .option('mode', 'PERMISSIVE') \
  .load()

# 提前完成解析，传入foreachBatch的是已经处理好的结构化DataFrame
df_parsed = stream_data \
    .withColumn("data", stream_data["body"].cast("string")) \
    .select(extract_udf("data").alias("extracted_data")) \
    .select("extracted_data.*")


def write2csv(df, epoch_id):
     # 批次内直接写入即可，不需要再做转换计算
     df.write \
       .format("csv") \
       .mode(SaveMode.Append) \
       .option("header", True) \
       .save("dbfs:/FileStore/Incoming/trade_date/")
 

query = df_parsed \
    .writeStream \
    .outputMode("append") \
    .foreachBatch(write2csv) \
    .option("checkpointLocation", "dbfs:/FileStore/checkpoint_trade_date/") \
    .trigger(processingTime="1 seconds") \
    .start()

query.awaitTermination()

额外说明：后续要实现不同标签值存储到独立文件，直接在write2csv函数内对传入的df做过滤拆分，分别写入不同路径即可，每个写入路径不需要单独配置checkpoint，只要流查询全局的checkpoint配置正确即可，注意每个独立的流查询要使用单独的checkpoint目录，不能复用。

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