如何在PySpark UDF中复用初始化耗时的类实例？

方法一：利用Pandas UDF结合全局变量 #

PySpark的Pandas UDF运行在Worker进程的Python环境中，每个Worker进程的全局变量只会初始化一次。我们可以借助这一点，在UDF内部判断全局实例是否存在，不存在则初始化：

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType  # 根据实际返回类型调整
from third_party_lib import Foo  # 导入第三方库的Foo类

# 定义全局变量，存储Worker进程级别的Foo实例
_foo_instance = None

def init_foo():
    global _foo_instance
    if _foo_instance is None:
        # 仅在进程首次调用时初始化一次
        _foo_instance = Foo()
    return _foo_instance

@pandas_udf(IntegerType())
def calculate_with_foo(input_series):
    # 确保当前进程已初始化Foo实例
    foo = init_foo()
    # 对Series中的每个元素调用calculate方法
    return input_series.apply(lambda x: foo.calculate(x))

# 示例使用
if __name__ == "__main__":
    spark = SparkSession.builder.appName("FooUDFExample").getOrCreate()
    df = spark.createDataFrame([(1,), (2,), (3,)], ["input_col"])
    result_df = df.withColumn("result", calculate_with_foo(df["input_col"]))
    result_df.show()
    spark.stop()

原理说明 #

每个Worker进程在第一次执行calculate_with_foo时，会触发init_foo初始化_foo_instance，后续该进程处理的所有数据都会复用这个实例，避免重复初始化。即使同一个Worker进程处理多个分区，也只会初始化一次。

方法二：使用模块级初始化分发 #

将Foo的初始化逻辑放在单独的Python模块中，通过Spark的addPyFile将模块分发到所有Worker，每个Worker进程加载模块时会自动初始化一次实例：

1. 创建名为foo_singleton.py的模块文件：

from third_party_lib import Foo

# 模块加载时自动初始化，每个Worker进程仅执行一次
foo_instance = Foo()

2. 在主程序中分发模块并定义UDF：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType

# 示例使用
if __name__ == "__main__":
    spark = SparkSession.builder.appName("FooModuleExample").getOrCreate()
    # 将自定义模块分发到所有Worker节点
    spark.sparkContext.addPyFile("foo_singleton.py")
    
    # 导入模块中的实例，需在addPyFile之后执行
    from foo_singleton import foo_instance
    
    def calculate_udf(input_val):
        return foo_instance.calculate(input_val)
    
    # 注册UDF
    calculate_udf_spark = udf(calculate_udf, IntegerType())
    
    df = spark.createDataFrame([(1,), (2,), (3,)], ["input_col"])
    result_df = df.withColumn("result", calculate_udf_spark(df["input_col"]))
    result_df.show()
    spark.stop()

原理说明 #

当Worker进程启动时，会加载foo_singleton.py模块，模块级的foo_instance = Foo()会执行一次，后续该进程中所有UDF调用都会复用这个实例。这种方式适合需要在多个UDF中共享同一个Foo实例的场景。

注意事项 #

• 如果Spark开启了动态资源分配，当Worker进程被销毁并重新创建时，会重新初始化Foo实例，这是正常的资源回收行为。
• 若Executor的spark.executor.cores设置大于1，每个Core对应的Worker进程都会独立初始化一次Foo实例（因为每个进程是独立的Python解释器），这符合Spark的多进程模型预期。
• 确保第三方库Foo已经安装在所有Worker节点上，否则会出现导入错误。

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