如何实现带进程级单次初始化代码的PySpark UDF

Python端完全可以实现和Scala单例对象初始化UDF对等的效果，核心利用PySpark Python worker的进程隔离特性：每个executor上启动的Python worker是独立进程，进程内的全局状态在整个进程生命周期内留存，只会初始化一次，不会随逐行数据处理重复执行。

核心逻辑说明 #

你对PySpark执行模型的理解是准确的：每个executor会启动独立Python进程处理分配到的UDF计算任务，进程内全局状态不跨进程、跨executor共享。只要把重量级资源（大模型、数据库连接、超大配置字典等）的初始化逻辑放在UDF逐行执行逻辑的外层，就能实现每个进程仅跑一次初始化的效果，不需要严格保证恰好执行一次，也不会出现逐行重复初始化的性能问题。

这种逻辑和你给出的Scala实现完全对应：Scala单例object MyUDF里的代码是每个JVM（executor）执行一次，Python的模块级全局代码是每个Python worker进程执行一次，逐行调用的UDF逻辑直接复用提前初始化好的对象即可。

常用实现方案 #

• 方案1：模块级预初始化
  该方案行为和Scala单例最接近，Python worker加载脚本时就会执行初始化逻辑，进程生命周期内全程复用资源，逻辑最简单不易出错。

  # 以下初始化代码每个Python进程加载脚本时仅执行一次
  # 替换成你自己的重量级初始化逻辑，比如加载模型、建立连接、加载大字典
  def init_heavy_resource():
      global heavy_model
      heavy_model = load_your_large_geocoding_model()
  
  heavy_model = None
  init_heavy_resource()
  
  from pyspark.sql.functions import udf
  from pyspark.sql.types import StringType
  
  @udf(returnType=StringType())
  def my_geocode_udf(lat: float, long: float) -> str:
      # 以下为逐行执行的逻辑，直接复用已经初始化好的全局对象，无重复初始化开销
      return heavy_model.search(lat, long)
  

  适用场景：确定UDF一定会被调用、初始化逻辑稳定无额外分支的场景。

• 方案2：惰性单次初始化
  该方案不会在脚本加载时就初始化资源，只有第一次实际调用UDF时才会触发初始化，整个进程仅执行一次，避免不必要的资源开销。

  from pyspark.sql.functions import udf
  from pyspark.sql.types import StringType
  
  heavy_model = None
  def _get_initialized_model():
      global heavy_model
      if heavy_model is None:
          # 第一次调用时才执行初始化，后续调用直接返回已初始化对象
          heavy_model = load_your_large_geocoding_model()
      return heavy_model
  
  @udf(returnType=StringType())
  def my_geocode_udf(lat: float, long: float) -> str:
      model = _get_initialized_model()
      return model.search(lat, long)
  

  适用场景：同一个脚本包含多个UDF、部分UDF可能不会被实际调用的场景，可以节省内存和初始化耗时。如果使用Pandas UDF（向量化UDF），该方案同样适用，逻辑不需要调整。

• 方案3：mapPartitions分区级初始化（备选）
  如果使用mapPartitions类算子而非普通UDF，可以在分区迭代逻辑的外层写初始化代码，这种方式是每个分区初始化一次。如果单个worker进程需要处理多个分区，会重复执行初始化，性能略低于全局变量方案，适合分区数和worker进程数基本匹配的场景。

避坑提示 #

• 不要把重量级初始化逻辑写在UDF函数体内部，否则每处理一行数据就会触发一次初始化，性能会出现数量级下降
• 不要在Driver端初始化重对象后直接传入UDF使用，这种方式会把Driver端的对象序列化后通过网络传输到每个任务，不仅传输、反序列化开销极大，还可能遇到对象不可序列化的报错
• 本地初始化的重对象不需要支持跨进程序列化，因为对象是每个worker进程在本地生成的，不存在跨进程传输问题

内容的提问来源于stack exchange，提问作者Or Bar Yaacov