Databricks Connect本地IDE运行foreachBatch流式代码报错求助

问题描述 #

本地IDE通过Databricks Connect运行Spark Streaming代码时抛出错误：

pyspark.errors.exceptions.connect.SparkException: No PYTHON_UID found for session (some uid)

• 代码在Databricks Notebook或集群上可正常运行
• 其他非Streaming场景下Databricks Connect工作正常
• 本地环境版本：databricks-connect 14.3.1、databricks-sdk 0.26.0、pyspark 3.5.1、Python 3.11.4

错误原因 #

foreachBatch函数的逻辑会在集群端执行，但本地Databricks Connect的Python会话上下文无法传递到集群，导致集群端Python进程无法找到关联的PYTHON_UID会话标识。

解决方法 #

方法1：在foreachBatch函数内使用传入DataFrame的SparkSession #

修改foreach_batch_function，显式使用传入DataFrame关联的集群SparkSession，避免依赖本地的DatabricksSession：

def foreach_batch_function(df, epoch_id):
    # 获取集群端的SparkSession，而非本地会话
    cluster_spark = df.sparkSession
    # 执行写入操作，此时会使用集群的会话上下文
    df.write.format("delta").mode("overwrite").saveAsTable(
        "sample.mvp_from_foreach_batch"
    )

方法2：修改DatabricksSession配置，禁用Worker守护进程 #

在创建DatabricksSession时添加配置，强制集群端Python进程直接关联本地会话：

from databricks.connect import DatabricksSession

spark = DatabricksSession.builder \
    .config("spark.databricks.connect.worker.daemon.enabled", "false") \
    .getOrCreate()

方法3：确保foreachBatch函数无外部依赖 #

避免在foreach_batch_function内直接引用本地的spark对象或其他非序列化外部变量，所有操作仅通过传入的df和epoch_id完成，保证函数可在集群端独立执行。

验证修改 #

重新运行代码后，通过assert spark.table("sample.mvp_from_foreach_batch").count() == 3确认表记录数符合预期，错误不再抛出。

内容的提问来源于stack exchange，提问作者Thomas W.