如何用PyArrow迭代器配合Gensim训练Word2Vec模型？

核心思路是自定义一个可重复迭代的类，封装PyArrow分批读取Parquet的逻辑，让Gensim的Word2Vec能完成两次数据扫描（第一次构建词汇表，第二次训练模型）。生成器只能遍历一次，不符合要求，而自定义类可以每次调用迭代器时重新从文件头开始读取。

步骤1：自定义可迭代类 #

这个类会保存Parquet文件路径、目标列名、分词函数等参数，每次迭代时重新打开文件并分批读取数据，返回分词后的文本序列：

import pyarrow.parquet as pq

class ParquetCorpus:
    def __init__(self, parquet_path, text_col='text', batch_size=10000, tokenizer=None):
        self.parquet_path = parquet_path
        self.text_col = text_col
        self.batch_size = batch_size
        self.tokenizer = tokenizer or self.default_tokenizer
    
    @staticmethod
    def default_tokenizer(text):
        # 替换成你的分词逻辑，比如空格分割、中文分词等
        return text.strip().split()
    
    def __iter__(self):
        # 每次迭代都重新打开Parquet文件，保证从头读取
        parquet_file = pq.ParquetFile(self.parquet_path)
        for batch in parquet_file.iter_batches(batch_size=self.batch_size):
            # 将批次数据转为Pandas DataFrame（或直接处理PyArrow Table）
            df = batch.to_pandas()
            # 遍历每一行文本，分词后返回
            for text in df[self.text_col].dropna():
                tokens = self.tokenizer(text)
                # 过滤过短的序列，避免影响训练
                if len(tokens) >= 2:
                    yield tokens

步骤2：用自定义类训练Word2Vec #

直接将ParquetCorpus实例传给Gensim的Word2Vec，它会自动完成两次扫描：

from gensim.models import Word2Vec

# 初始化语料对象
corpus = ParquetCorpus(
    parquet_path="your_large_file.parquet",
    text_col="content",  # 替换成你的文本列名
    batch_size=20000,    # 根据内存调整批次大小
    tokenizer=your_custom_tokenizer  # 可选，替换成你的分词函数
)

# 训练Word2Vec模型
model = Word2Vec(
    sentences=corpus,
    vector_size=100,
    window=5,
    min_count=5,
    workers=4,
    epochs=5
)

# 保存模型
model.save("word2vec_model.model")

关键注意事项 #

• 批次大小调整：根据你的内存情况设置batch_size，确保单个批次的数据能载入内存。
• 分词逻辑：一定要替换成适合你文本的分词方式（比如中文用jieba、英文用nltk），这直接影响模型效果。
• 空值过滤：代码中用dropna()过滤空文本，避免训练时报错。
• 多分区Parquet：如果你的Parquet是分区存储（比如按日期分区），可以修改__iter__方法，用pq.ParquetDataset代替pq.ParquetFile来读取所有分区：

  dataset = pq.ParquetDataset(self.parquet_path)
  for batch in dataset.iter_batches(batch_size=self.batch_size):
      # 后续处理逻辑不变
  

内容的提问来源于stack exchange，提问作者ℕʘʘḆḽḘ