Databricks中UDF使用线程刷新API令牌失效问题排查

问题背景 #

有一个需要逐行调用API的DataFrame，通过UDF实现处理逻辑：

processDf = FilesDf.withColumn("processed", row_udf(col("input_id"))).select("processed.*").cache()

但处理耗时超过1小时，而API请求头中的令牌有效期仅1小时。尝试用后台线程定期刷新全局headers变量：

def get_new_token():
    app = TokenApplication(...)
    result = app.acquire_token_for_client(...)
    ...
    headers = {"Authorization": token}
    return headers

headers = get_new_token()

stop_event = threading.Event()

def refresh_token_periodically():
    global headers
    while not stop_event.is_set():
        time.sleep(55 * 60)  # 每55分钟刷新一次
        headers = get_new_token()

token_refresh_thread = threading.Thread(target=refresh_token_periodically)
token_refresh_thread.daemon = True
token_refresh_thread.start()

测试时令牌能正常刷新，但UDF中使用的仍是失效的旧令牌。想确认：UDF是否提前加载所有行并使用初始令牌？该方案是否可行？

核心原因分析 #

你的方案无效的根本原因是Spark的分布式执行模型，而非UDF提前加载行：

1. Spark的Driver节点负责分发任务，Executor节点负责实际执行UDF。
2. 你定义的全局headers变量存储在Driver端，当UDF被分发到Executor时，这个变量会被序列化后传递给Executor，Executor端的UDF使用的是这个序列化副本，而非Driver端的原变量。
3. Driver端的后台线程更新的是Driver本地的headers，这个更新无法同步到Executor端的副本，因此UDF始终使用初始令牌，后续调用自然失效。

可行解决方案 #

方案1：使用广播变量（Broadcast Variable） #

广播变量是Spark用来在集群节点间共享只读数据的机制，通过可变对象包装令牌，可实现Driver端更新后Executor端获取最新值：

from pyspark.sql import SparkSession
import threading
import time

spark = SparkSession.builder.appName("TokenRefresh").getOrCreate()

def get_new_token():
    # 替换为实际的令牌获取逻辑
    app = TokenApplication(...)
    result = app.acquire_token_for_client(...)
    token = result.get("access_token")
    return {"Authorization": f"Bearer {token}"}

# 用可变列表包装headers，广播后可更新内部值
headers_wrapper = [get_new_token()]
headers_broadcast = spark.sparkContext.broadcast(headers_wrapper)

stop_event = threading.Event()

def refresh_token_periodically():
    while not stop_event.is_set():
        time.sleep(55 * 60)
        # 更新包装类中的最新令牌
        new_headers = get_new_token()
        headers_wrapper[0] = new_headers

# 启动刷新线程
token_refresh_thread = threading.Thread(target=refresh_token_periodically)
token_refresh_thread.daemon = True
token_refresh_thread.start()

# 定义UDF，使用广播变量的最新值
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StructType, StructField, StringType

# 根据实际API返回结构定义Schema
result_schema = StructType([
    StructField("output_id", StringType(), True),
    StructField("data", StringType(), True)
])

def call_api(input_id):
    # 获取广播变量中的最新headers
    current_headers = headers_broadcast.value[0]
    # 替换为实际API调用逻辑
    import requests
    response = requests.get(f"https://api.example.com/process/{input_id}", headers=current_headers)
    response_data = response.json()
    return (response_data.get("output_id"), response_data.get("data"))

row_udf = udf(call_api, result_schema)

# 处理DataFrame
processDf = FilesDf.withColumn("processed", row_udf(col("input_id"))).select("processed.*").cache()
processDf.show()

# 程序结束时停止线程
stop_event.set()
token_refresh_thread.join()

原理：广播的是可变列表的引用，Driver端更新列表内部值后，Executor端通过广播变量的value属性能直接获取最新内容。

方案2：UDF内部本地刷新令牌 #

在UDF内部维护令牌的本地副本，定期检查并刷新，适合API令牌限制宽松的场景：

import time
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StructType, StructField, StringType

# UDF内部维护的令牌状态
last_refresh_time = 0
current_headers = None

def get_new_token():
    # 替换为实际令牌获取逻辑
    app = TokenApplication(...)
    result = app.acquire_token_for_client(...)
    token = result.get("access_token")
    return {"Authorization": f"Bearer {token}"}

result_schema = StructType([
    StructField("output_id", StringType(), True),
    StructField("data", StringType(), True)
])

def call_api(input_id):
    global last_refresh_time, current_headers
    # 每55分钟刷新一次令牌
    if time.time() - last_refresh_time > 55 * 60 or current_headers is None:
        current_headers = get_new_token()
        last_refresh_time = time.time()
    # 实际API调用逻辑
    import requests
    response = requests.get(f"https://api.example.com/process/{input_id}", headers=current_headers)
    response_data = response.json()
    return (response_data.get("output_id"), response_data.get("data"))

row_udf = udf(call_api, result_schema)
processDf = FilesDf.withColumn("processed", row_udf(col("input_id"))).select("processed.*").cache()

注意：每个Executor的每个任务都会维护独立的令牌状态，可能会产生较多令牌刷新请求，需评估API的调用频率限制。

总结 #

• 你的原方案不可行，因为分布式环境下Driver端的全局变量无法同步到Executor端的UDF副本。
• UDF并非提前加载所有行使用初始令牌，而是Executor端的UDF始终使用任务分发时的变量副本。
• 推荐优先使用广播变量方案，能在集群内高效共享最新令牌，减少重复刷新请求。

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