Databricks UDF中执行where查询报错：无法序列化_thread.RLock对象

错误原因 #

你遇到的TypeError: cannot pickle '_thread.RLock' object错误，核心问题是Spark UDF内部不能直接引用Spark DataFrame（或关联的SparkSession、SparkContext等带锁对象）。Spark会将UDF序列化后分发到集群各个executor执行，但DataFrame依赖的_thread.RLock属于不可序列化对象，无法跨节点传输，因此触发报错。

解决方案 #

不要用UDF做逐行查询，改用Spark原生的分布式关联操作，以下是两种可行方案：

方案一：自连接+窗口函数（推荐，适合大数据量） #

利用Spark的分布式关联和窗口函数实现需求，性能远优于UDF逐行查询：

from pyspark.sql import Window
from pyspark.sql.functions import col, first, when

# 1. 自连接原表，匹配doc相同且alm不同的记录
joined_df = mytab.alias("main")
    .join(
        mytab.alias("lookup"),
        (col("main.Doc_mat") == col("lookup.doc")) & (col("main.Alm") != col("lookup.alm")),
        how="left"
    )

# 2. 定义窗口：按主表的关键字段分组，确保每个主表行只取第一条匹配的lookup记录
window_spec = Window.partitionBy("main.Alm", "main.Proveedor", "main.Doc_mat")
    .orderBy(col("lookup.doc").asc())  # 可根据业务需求调整排序规则

# 3. 生成目标字段：有匹配则取第一条alm，无匹配则返回原Proveedor值
result_df = joined_df.withColumn(
    "alm_origen",
    when(
        first(col("lookup.alm"), ignorenulls=True).over(window_spec).isNull(),
        col("main.Proveedor")
    ).otherwise(
        first(col("lookup.alm"), ignorenulls=True).over(window_spec)
    )
)

# 保留原表所有字段+新增的alm_origen
result_df = result_df.select("main.*", "alm_origen")

方案二：本地字典+广播UDF（仅适合小数据集） #

如果原表数据量较小，可将数据拉取到driver节点生成字典，再通过广播变量传递给UDF：

from pyspark.sql.functions import udf, broadcast
from pyspark.sql.types import StringType

# 将原表的doc和alm映射关系收集到本地字典
lookup_dict = {}
for row in mytab.select("doc", "alm").collect():
    if row.doc not in lookup_dict:
        lookup_dict[row.doc] = []
    lookup_dict[row.doc].append(row.alm)

# 广播字典到所有executor，优化传输效率
broadcast_lookup = spark.sparkContext.broadcast(lookup_dict)

def buscarx(Alm_r, Pro, Data_mat):
    # 从广播字典中获取当前doc对应的alm列表
    alms = broadcast_lookup.value.get(Data_mat, [])
    # 返回第一个不等于Alm_r的alm
    for alm in alms:
        if alm != Alm_r:
            return alm
    # 无匹配时返回Pro
    return Pro

# 注册UDF并生成新字段
buscarx_udf = udf(buscarx, StringType())
df_temp = mytab.withColumn("alm_origen", buscarx_udf(mytab.Alm, mytab.Proveedor, mytab.Doc_mat))

注意事项 #

• 优先选择方案一：完全基于Spark分布式计算，避免数据拉取到driver节点，适合大规模数据集
• 方案二仅用于小数据场景：collect()会将全表数据加载到driver内存，数据量大时会导致内存溢出

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