Databricks调用自定义货币转换UDF时出现序列化错误求助

问题描述 #

在Databricks中编写货币转换函数并封装为UDF时执行报错，错误提示为序列化失败，无法在Worker节点引用SparkContext。

原代码 #

from pyspark.sql.functions import lit
from pyspark.sql.functions import when, col
from pyspark.sql.functions import udf
from pyspark.sql.types import FloatType

df_curr1 = spark.read.table("table")

def currency(from_currency,to_currency,date,rate_type):
    
    df_select = df_curr1.select(df_curr1["col0"].alias("BaseCurrency"), df_curr1["col1"].alias("TargetCurrency"),df_curr1["col2"].alias("EffectiveDate"),df_curr1["col3"].alias("ExpressInBaseCurrency"),df_curr1["col4"].alias("Ratetype"),df_curr1["col5"].alias("Rate"))

    #display(df_select)

    df_select = df_select.withColumn(
        "Rate",
        
        when(
            (df_select["ExpressInBaseCurrency"] == 2),

            1/df_select["Rate"]
        ).otherwise(df_select["Rate"]) )
    display(df_select)
    
    if from_currency == to_currency :
        rate = 1.0

    elif df_select.where(col("BaseCurrency") == from_currency).count() > 0:
        if df_select.where((df_select.BaseCurrency == from_currency ) & (df_select.TargetCurrency == to_currency ) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type)).count() > 0:
            
            selected_df = df_select.filter( (df_select.BaseCurrency == from_currency ) & (df_select.TargetCurrency == to_currency ) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type) )
            selected_df = selected_df.orderBy(selected_df["EffectiveDate"].desc())
            selected_df = selected_df.limit(1)
            rate = selected_df.select("Rate").head()[0]
        else:
            selected_df = df_select.filter( (df_select.BaseCurrency == from_currency ) & (df_select.TargetCurrency == to_currency ) & (df_select.EffectiveDate <= date) )
            selected_df = selected_df.orderBy(selected_df["EffectiveDate"].desc())
            selected_df = selected_df.limit(1)
            rate = selected_df.select("Rate").head()[0]
    elif  df_select.where(col("BaseCurrency") == to_currency ).count() > 0:
        
        if df_select.where((df_select.BaseCurrency == to_currency ) & (df_select.TargetCurrency == from_currency ) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type)).count() > 0:
            selected_df = df_select.filter( (df_select.BaseCurrency == to_currency ) & (df_select.TargetCurrency == from_currency) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type) )

            selected_df = selected_df.limit(1)
            display(selected_df)
            rate = selected_df.select("Rate").head()[0]
            rate = 1/rate

        else:
            if df_select.where((df_select.BaseCurrency == to_currency ) & (df_select.TargetCurrency == from_currency) & (df_select.EffectiveDate <= date)).count()>0:
                selected_df = df_select.filter( (df_select.BaseCurrency == to_currency ) & (df_select.TargetCurrency == from_currency) & (df_select.EffectiveDate <= date) )
                selected_df = selected_df.limit(1)
                display(selected_df)
                rate = selected_df.select("Rate").head()[0]
                rate = 1/rate
            else:
                selected_df = df_select.filter( (df_select.BaseCurrency == to_currency ) & (df_select.TargetCurrency == from_currency)  )
                selected_df = selected_df.limit(1)
                display(selected_df)
                rate = selected_df.select("Rate").head()[0]
                rate = 1/rate


    else:
        if df_select.where((df_select.TargetCurrency == from_currency ) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type)).count() > 0:
            selected_df = df_select.filter((df_select.TargetCurrency == from_currency ) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type))
            selected_df = selected_df.orderBy(selected_df["EffectiveDate"].desc())
            selected_df = selected_df.limit(1)
            first_value = selected_df.select("Rate").head()[0]
        else:
            selected_df = df_select.filter((df_select.TargetCurrency == from_currency ) & (df_select.EffectiveDate <= date))
            selected_df = selected_df.orderBy(selected_df["EffectiveDate"].desc())
            selected_df = selected_df.limit(1)
            first_value = selected_df.select("Rate").head()[0]
        
        if df_select.where((df_select.TargetCurrency == to_currency) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type)).count() > 0:
            selected_df2 = df_select.filter((df_select.TargetCurrency == to_currency) & (df_select.EffectiveDate <= date) & (df_select.Ratetype == rate_type))
            selected_df2 = selected_df2.orderBy(selected_df2["EffectiveDate"].desc())
            selected_df2 = selected_df2.limit(1)
            second_value = selected_df2.select("Rate").head()[0]
        else:
            selected_df2 = df_select.filter((df_select.TargetCurrency == to_currency) & (df_select.EffectiveDate <= date) )
            selected_df2 = selected_df2.orderBy(selected_df2["EffectiveDate"].desc())
            selected_df2 = selected_df2.limit(1)
            second_value = selected_df2.select("Rate").head()[0]

        rate = second_value / first_value


    return rate

currency_udf = udf(currency, FloatType())
invoice_currency_col = df_fact2["TransactionCurrency"]
currency1_col = df_fact2["Currency1"]
document_date_col = df_fact2["DateofFirstContact"]
rate_type_col = df_fact2["OpportunityType"]

# Apply the currency function using withColumn
df_fact2 = df_fact2.withColumn("rate1", currency_udf(invoice_currency_col, currency1_col, document_date_col, rate_type_col))

报错信息 #

PicklingError: Could not serialize object: RuntimeError: It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.

输入数据 #

df_fact2样本 #

df_curr1样本 #

预期输出 #

核心问题分析 #

• UDF内执行Spark分布式操作：currency函数内部直接操作df_curr1（Spark DataFrame），包括select、filter、count等，但UDF运行在Worker节点，无法访问Driver端的SparkContext，直接导致序列化失败。
• Driver端操作混入UDF：UDF内部调用display、head()等仅能在Driver执行的操作，这类操作无法分发到Worker节点。
• 性能逻辑错误：即使序列化成功，UDF内循环查询DataFrame的方式会引发大量重复计算，完全违背Spark分布式计算的设计理念。

正确解决方案：用Spark原生JOIN+窗口函数替代UDF #

Spark的核心是分布式数据处理，应通过DataFrame的JOIN和窗口函数实现需求，而非UDF。具体步骤如下：

1. 预处理汇率表df_curr1 #

统一汇率计算规则，生成标准化汇率表：

from pyspark.sql import Window
from pyspark.sql.functions import col, when, row_number, coalesce

# 预处理汇率表，统一为Base->Target的汇率
df_curr_clean = df_curr1.select(
    col("col0").alias("BaseCurrency"),
    col("col1").alias("TargetCurrency"),
    col("col2").alias("EffectiveDate"),
    col("col3").alias("ExpressInBaseCurrency"),
    col("col4").alias("Ratetype"),
    col("col5").alias("Rate")
).withColumn(
    "StandardRate",
    when(col("ExpressInBaseCurrency") == 2, 1/col("Rate")).otherwise(col("Rate"))
)

# 为每组(Base, Target, Ratetype)按生效日期倒序排序，取最新有效汇率
window_spec = Window.partitionBy("BaseCurrency", "TargetCurrency", "Ratetype")\
                   .orderBy(col("EffectiveDate").desc())

df_curr_ranked = df_curr_clean.withColumn(
    "rn", row_number().over(window_spec)
).filter(col("rn") == 1).drop("rn", "ExpressInBaseCurrency", "Rate")

# 生成反向汇率（Target->Base），覆盖反向转换场景
df_curr_reverse = df_curr_ranked.select(
    col("TargetCurrency").alias("BaseCurrency"),
    col("BaseCurrency").alias("TargetCurrency"),
    col("EffectiveDate"),
    col("Ratetype"),
    (1/col("StandardRate")).alias("StandardRate")
)

# 合并正向与反向汇率表
df_curr_final = df_curr_ranked.unionByName(df_curr_reverse)

2. 关联业务表并计算汇率 #

通过JOIN匹配汇率，处理所有边界场景：

# 第一步：匹配完全符合条件的汇率（货币对、Ratetype、生效日期）
df_fact_with_rate = df_fact2.join(
    df_curr_final,
    (col("TransactionCurrency") == df_curr_final["BaseCurrency"]) &
    (col("Currency1") == df_curr_final["TargetCurrency"]) &
    (col("OpportunityType") == df_curr_final["Ratetype"]) &
    (col("DateofFirstContact") >= df_curr_final["EffectiveDate"]),
    how="left"
).withColumnRenamed("StandardRate", "matched_rate")

# 第二步：无对应Ratetype时， fallback到同货币对的最新汇率
window_fallback = Window.partitionBy("TransactionCurrency", "Currency1")\
                       .orderBy(col("EffectiveDate").desc())

df_fact_fallback = df_fact_with_rate.withColumn(
    "fallback_rate",
    coalesce(col("matched_rate"), first("StandardRate").over(window_fallback))
)

# 第三步：处理同货币场景，直接设为1.0
df_final = df_fact_fallback.withColumn(
    "rate1",
    when(col("TransactionCurrency") == col("Currency1"), 1.0).otherwise(col("fallback_rate"))
).drop("BaseCurrency", "TargetCurrency", "EffectiveDate", "Ratetype", "matched_rate", "fallback_rate")

display(df_final)

方案优势 #

• 完全基于Spark分布式操作，彻底避免UDF的序列化问题与性能瓶颈
• 逻辑清晰，符合Spark计算范式
• 覆盖所有边界场景（同货币、无匹配Ratetype、反向汇率）

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