PySpark Pandas UDF无法返回拟合Tokenizer的解决方案咨询

你的现有代码存在关键问题：每个Pandas UDF的分区会独立拟合自己的Tokenizer，导致不同分区的词汇表不一致，且Pandas UDF无法直接返回复杂的Tokenizer对象。以下是两种可行解决方案：

方案一：小数据量场景 - Driver端全局拟合后广播 #

适合数据量不大的情况，先将所有文本收集到Driver端拟合Tokenizer，再广播到所有Executor复用：

from pyspark.sql.functions import pandas_udf, broadcast
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
import pandas as pd
import numpy as np

def tokenize_and_get_tokenizer(df, text_col, maxlen, padding_type):
    # 收集所有文本到Driver，全局拟合Tokenizer
    all_texts = df.select(text_col).rdd.flatMap(lambda row: row).collect()
    tokenizer = Tokenizer(num_words=None, char_level=True, oov_token='UNK')
    tokenizer.fit_on_texts(all_texts)
    
    # 广播拟合好的Tokenizer到所有Executor
    broadcast_tokenizer = broadcast(tokenizer)
    
    # 定义Pandas UDF处理文本
    @pandas_udf('array<decimal>')
    def process_text(text_series):
        sequences = broadcast_tokenizer.value.texts_to_sequences(text_series)
        padded_data = pad_sequences(sequences, maxlen=maxlen, padding=padding_type, truncating=padding_type)
        return pd.Series(padded_data.tolist())
    
    # 添加分词后的列
    df_tokenized = df.withColumn(f"{text_col}_tokenized", process_text(text_col))
    return df_tokenized, tokenizer

调用该函数后，可同时获取处理后的DataFrame和拟合完成的Tokenizer，直接用于后续测试数据处理。

方案二：大数据量场景 - 分布式统计词汇表后构建Tokenizer #

如果数据量过大无法收集到Driver，可通过Spark分布式统计全局字符频次，再构建Tokenizer：

from pyspark.sql.functions import explode, split, count, col
from tensorflow.keras.preprocessing.text import Tokenizer

def build_global_tokenizer(df, text_col):
    # 拆分文本为单个字符，统计全局频次
    char_freq = df.select(explode(split(col(text_col), '')).alias('char')) \
                  .groupBy('char').agg(count('*').alias('freq')) \
                  .orderBy(col('freq').desc())
    
    # 获取排序后的字符列表，构建Tokenizer词汇表
    chars = char_freq.select('char').rdd.flatMap(lambda row: row).collect()
    tokenizer = Tokenizer(num_words=None, char_level=True, oov_token='UNK')
    
    # 手动设置词汇映射，确保和全局频次一致
    tokenizer.word_index = {char: idx + 1 for idx, char in enumerate(chars)}
    tokenizer.index_word = {idx + 1: char for idx, char in enumerate(chars)}
    tokenizer.index_word[0] = ''  # 0对应padding
    
    return tokenizer

# 构建Tokenizer后广播使用
tokenizer = build_global_tokenizer(df, 'text')
broadcast_tokenizer = broadcast(tokenizer)

# 后续处理UDF同方案一

关键注意事项 #

• 禁止在Pandas UDF内部拟合Tokenizer：每个分区会生成独立的词汇表，导致训练和测试数据的分词规则不一致，结果完全不可靠。
• 广播机制确保一致性：通过Spark的广播变量，所有Executor会复用同一个拟合好的Tokenizer，保证全局分词规则统一。

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