Spark DataFrame应用KeyBERT外部函数替代RDD及多列处理方法

1. 不转换RDD直接基于DataFrame实现同等逻辑 #

不需要将DataFrame转为RDD，直接使用Pandas UDF即可实现逐行调用KeyBERT模型，且性能比普通Python UDF、RDD map更高（基于Apache Arrow做批量数据序列化，减少跨进程传输开销，且支持executor端懒加载模型避免重复初始化）。

实现步骤如下：

• 首先封装模型懒加载逻辑，避免KeyBERT模型无法序列化的报错，保证每个executor节点只初始化一次模型
• 定义符合返回结构的Pandas UDF，封装关键词提取逻辑
• 直接通过withColumn给DataFrame新增关键词列即可

完整代码示例：

from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, StructType, StructField, StringType, FloatType
from keybert import KeyBERT
import pandas as pd

# 模型懒加载：每个executor工作节点仅初始化一次KeyBERT实例
def get_keybert_model():
    if not hasattr(get_keybert_model, "model"):
        get_keybert_model.model = KeyBERT()
    return get_keybert_model.model

# 定义关键词返回结构：每个关键词对应匹配权重
keyword_result_schema = ArrayType(
    StructType([
        StructField("keyword", StringType(), nullable=False),
        StructField("similarity_score", FloatType(), nullable=False)
    ])
)

# 定义单文本关键词提取Pandas UDF
@F.pandas_udf(returnType=keyword_result_schema)
def extract_keywords(text_col: pd.Series) -> pd.Series:
    kw_model = get_keybert_model()
    processed = []
    for text in text_col:
        # extract_keywords的参数可以根据自身需求调整，和原有RDD逻辑保持一致即可
        raw_res = kw_model.extract_keywords(text)
        processed.append([
            {"keyword": kw, "similarity_score": score} 
            for kw, score in raw_res
        ])
    return pd.Series(processed)

# 直接在初始DataFrame上调用，无需转RDD
final_df = init_df.withColumn("var_1_keywords", extract_keywords(F.col("var_1")))

注意：不要在Driver端直接初始化KeyBERT实例后在UDF中引用，会触发模型序列化失败的报错，上述懒加载方式是Spark调用第三方Python模型的标准实践。

2. 同时对var_1、var_2两列应用KeyBERT模型 #

完全支持，有两种可选实现方式：

• 方式一：复用单文本UDF分别处理两列，代码简单易维护，适合中小数据量场景

final_df_two_cols = init_df.withColumn("var_1_keywords", extract_keywords(F.col("var_1")))\
                           .withColumn("var_2_keywords", extract_keywords(F.col("var_2")))

• 方式二：定义多输入Pandas UDF，一次处理两列文本，减少模型重复调用开销，适合大数据量场景

# 定义两列输入的UDF返回结构
two_col_result_schema = StructType([
    StructField("var_1_keywords", keyword_result_schema),
    StructField("var_2_keywords", keyword_result_schema)
])

@F.pandas_udf(returnType=two_col_result_schema)
def extract_keywords_two_col(var1_series: pd.Series, var2_series: pd.Series) -> pd.Series:
    kw_model = get_keybert_model()
    processed = []
    for v1_text, v2_text in zip(var1_series, var2_series):
        v1_kws = [{"keyword": kw, "similarity_score": score} for kw, score in kw_model.extract_keywords(v1_text)]
        v2_kws = [{"keyword": kw, "similarity_score": score} for kw, score in kw_model.extract_keywords(v2_text)]
        processed.append({"var_1_keywords": v1_kws, "var_2_keywords": v2_kws})
    return pd.Series(processed)

# 调用方式
final_df_two_cols_opt = init_df.withColumn(
        "kw_result", 
        extract_keywords_two_col(F.col("var_1"), F.col("var_2"))
    ).select(
        "id", "var_1", "var_2",
        "kw_result.var_1_keywords",
        "kw_result.var_2_keywords"
    )

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