如何在Spark所有worker节点预初始化gensim对象解决UDF加载慢问题

核心问题根源：你将词典、语料、相似度索引的加载逻辑写在了UDF内部，每次调用UDF处理行都会重复读取磁盘加载资源，这是耗时高的核心原因。可通过进程级惰性初始化的方案，让每个worker进程仅加载一次资源，所有行处理复用已加载的对象。

优化后代码示例 #

import numpy as np
from gensim import corpora, models
from nltk.tokenize import word_tokenize
from pyspark import SparkFiles
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 模块级全局变量，常驻worker进程内存，仅初始化一次
_dictionary = None
_corpus = None
_tf_idf = None
_sims = None
_file_docs = None # 请自行补充file_docs的加载逻辑

def init_resources():
    global _dictionary, _corpus, _tf_idf, _sims, _file_docs
    # *仅当资源未加载时执行加载逻辑*
    if _dictionary is None:
        path_index = SparkFiles.get("corpus_final_production.index")
        path_dictionary = SparkFiles.get('dictionary_production.gensim')
        path_corpus = SparkFiles.get("corpus_final_production")
        _dictionary = corpora.Dictionary.load(path_dictionary)
        _corpus = corpora.MmCorpus(path_corpus)
        _tf_idf = models.TfidfModel(_corpus)
        _sims = models.similarities.Similarity(
            path_index,
            _tf_idf[_corpus],
            num_features=len(_dictionary)
        )
        # 此处补充你的file_docs加载逻辑，和其他资源一起加载
        # _file_docs = 你的加载代码

def test(string):
    # 先触发资源加载，第一次调用执行加载，后续直接复用
    init_resources()
    query_doc = word_tokenize(string.lower())
    query_doc_bow = _dictionary.doc2bow(query_doc)
    query_doc_tf_idf = _tf_idf[query_doc_bow]
    max_sims_origin = _file_docs[np.argmax(_sims[query_doc_tf_idf])]
    return max_sims_origin

test_udf = udf(test, StringType())
df_new = garuda.withColumn('max_sim_origin', test_udf(garuda.text))

方案原理 #

PySpark的Executor进程是常驻的，每个进程会处理多批次的行数据。模块级的全局变量会在进程生命周期内保留，init_resources仅会在每个worker进程第一次调用UDF的时候执行一次，后续所有行处理都直接复用已加载的资源，完全避免了重复加载的开销。

额外优化建议 #

• 可提前将预训练好的TfidfModel和Similarity索引序列化存储，加载时直接读取序列化后的文件，无需加载原始语料重新计算，进一步降低初始化耗时。
• 若数据量较大，可使用pandas_udf替代普通UDF，以批处理的方式处理数据，进一步提升执行效率。

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