PySpark+sentence-transformers文本嵌入遇序列化错误求助

问题根源 #

你遇到的TypeError: cannot pickle '_thread.RLock' object错误，是因为Spark在序列化UDF关联的text_embedder实例时，会尝试序列化整个对象——哪怕你用@property延迟加载模型，类实例的内部结构（或SentenceTransformer库的隐含依赖）包含了无法被pickle序列化的线程锁对象，导致序列化失败。

需求可行性 #

完全可以实现PySpark结合Sentence-Transformers的大规模文本嵌入，核心思路是在每个Worker节点上延迟加载模型，且每个Worker进程仅加载一次，避免模型序列化传递的问题。

正确实现方式 #

方案1：普通UDF + 模块级延迟加载模型 #

from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, FloatType
from sentence_transformers import SentenceTransformer

# 模块级变量，每个Worker进程仅初始化一次模型
_model = None

def get_model():
    global _model
    if _model is None:
        _model = SentenceTransformer('dangvantuan/sentence-camembert-large')
    return _model

def embed_single_text(text):
    if not text:
        return []
    model = get_model()
    return model.encode(text).tolist()

# 注册UDF
embed_text_udf = F.udf(embed_single_text, returnType=ArrayType(FloatType()))

# 生成嵌入列
df_with_embeds = df.withColumn("text_embedded", embed_text_udf(F.col("text")))

方案2：Pandas UDF（推荐，性能更优） #

利用Pandas UDF的批量处理能力，配合Sentence-Transformers的批量编码接口，大幅提升处理效率：

from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, FloatType
from sentence_transformers import SentenceTransformer
import pandas as pd

_model = None

def get_model():
    global _model
    if _model is None:
        _model = SentenceTransformer('dangvantuan/sentence-camembert-large')
    return _model

@F.pandas_udf(ArrayType(FloatType()))
def embed_text_batch(text_series: pd.Series) -> pd.Series:
    # 过滤空文本
    valid_texts = text_series.dropna().tolist()
    if not valid_texts:
        return pd.Series([[]] * len(text_series))
    
    model = get_model()
    embeddings = model.encode(valid_texts, batch_size=32)  # 可根据内存调整batch_size
    
    # 映射回原序列，保持顺序
    result = []
    idx = 0
    for text in text_series:
        if pd.isna(text):
            result.append([])
        else:
            result.append(embeddings[idx].tolist())
            idx += 1
    return pd.Series(result)

# 生成嵌入列
df_with_embeds = df.withColumn("text_embedded", embed_text_batch(F.col("text")))

性能表现分析 #

采用Pandas UDF的方案性能非常理想，原因如下：

• 模型复用：每个Worker进程仅加载一次模型，避免重复加载的巨大开销
• 批量处理：Sentence-Transformers的encode方法对批量输入优化极佳，单条处理的效率远低于批量处理
• Arrow优化：Pandas UDF基于Apache Arrow传递数据，减少了Spark与Python之间的序列化/反序列化开销

性能优化建议 #

• 确保Worker节点有足够内存加载sentence-camembert-large大模型（建议每个Worker分配至少8G内存）
• 根据集群CPU核心数调整Spark并行度参数（如spark.executor.cores、spark.sql.shuffle.partitions），充分利用集群资源
• 调整batch_size参数，平衡内存占用与处理速度

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