PySpark大数据量下匹配词向量表计算文本平均向量的优化方案

原代码核心问题 #

原代码存在两个致命的性能问题，无法适配大数据量场景：

1. 调用collect()将全量文本数据拉取到Driver节点单点处理，数据量超过Driver内存就会直接OOM
2. 循环遍历每个词时都单独触发过滤、collect操作，会产生大量的Spark作业调度开销，执行效率极低

优化实现方案 #

全程使用Spark原生分布式算子实现，所有计算都在Executor端执行，无数据拉取到Driver节点，支持TB级数据量运行。

完整实现代码 #

import pyspark.sql.functions as F
from pyspark.sql.types import ArrayType, DoubleType
from pyspark.sql.functions import broadcast

# 第一步：给原始文本每行添加唯一标识ID，同时预处理空值、空字符串为合法空数组
df_with_id = df.withColumn("row_id", F.monotonically_increasing_id()) \
                 .withColumn("text", F.when(F.col("text").isNull(), F.array()) \
                                       .when(F.col("text") == "", F.array()) \
                                       .otherwise(F.col("text")))

# 第二步：将每行的分词数组炸开，一行转多行，每个词单独为一行
df_exploded = df_with_id.select("row_id", F.explode("text").alias("word"))

# 第三步：和词向量对照表左关联匹配每个词对应的向量
# 词向量表属于小表的场景下，加broadcast广播到所有Executor，大幅提升关联性能
df_joined = df_exploded.join(broadcast(df_vec.withColumn("word", F.trim(F.col("word")))), on="word", how="left")

# 第四步：定义UDF计算同一行所有有效向量的平均值
calculate_avg_vec = F.udf(
    lambda vec_list: [0.0]*3 if not [v for v in vec_list if v] \
    else [sum(dim)/len([v for v in vec_list if v]) for dim in zip(*[v for v in vec_list if v])],
    ArrayType(DoubleType())
)

# 第五步：按行ID分组聚合，计算每行的平均向量
df_result = df_joined.groupBy("row_id") \
                     .agg(F.collect_list("vector").alias("all_vecs")) \
                     .withColumn("avg_vector", calculate_avg_vec("all_vecs"))

# 关联回原始文本列，得到最终结果
df_final = df_result.join(df_with_id, on="row_id", how="left").select("text", "avg_vector")

# 查看结果
df_final.show(truncate=False)

代码说明 #

• 加了broadcast优化是因为词向量对照表通常数据量不大，广播到所有Executor后无需Shuffle即可完成关联，性能提升非常明显
• 代码里自动对df_vec的word列做了trim处理，避免示例中"could "带空格导致匹配失败的问题
• 没有匹配到任何有效词的行、空行都会返回[0.0, 0.0, 0.0]，和原逻辑保持一致
• 向量维度如果有变化，只需要修改UDF返回的默认空向量长度即可，无需修改其他逻辑

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