多进程计算PySpark数据各pt分组Faiss索引无输出问题求助

问题场景 #

现有如下PySpark DataFrame：

id    emb          pt
1    [5,0.8..]     'h'
2    [0.7, 0.8..]  'd'
3    [1, 3, ..]    'h'
4    [0.3, 0.8..]  'f'

需求是按pt分组对emb列生成Faiss索引，代码逻辑为：

1. 获取唯一pt值：

pt = [val.pt for val in df.select('pt').distinct().collect()]

2. 编写多进程函数处理单个pt：

def multiprocess(pt):    
    item_emb = df.filter(f.col('pt') == pt).select('emb').rdd.flatMap(lambda x: np.float32(x)).collect()
    item_id = df.filter(f.col('pt') == pt).select('id').rdd.flatMap(lambda x: x).collect()
    item_ids_np = np.array(item_id, dtype=np.int64)
    item_embs_np = np.array(item_emb)
    pool_size = item_embs_np.shape[0]
    emb_size = item_embs_np.shape[1]
    
    nlist = int(math.sqrt(pool_size))
    quantizer = faiss.IndexFlatL2(emb_size)
    
    index = faiss.IndexIVFFlat(quantizer, emb_size, nlist)
    index.train(item_embs_np)
    
    index.add_with_ids(item_embs_np, item_ids_np)
    faiss.write_index(index, 'Faiss_Index/train_'+pt+'.index')

3. 启动多进程执行：

pool = mp.Pool(processes = (mp.cpu_count() - 1))
pool.map(multiprocess, pt)

但代码持续运行无输出文件，需排查原因。

核心问题与解决方案 #

1. 多进程无法共享PySpark上下文 #

PySpark的DataFrame依赖Spark上下文（SparkContext），Python多进程的子进程无法继承主进程的Spark上下文，导致子进程中df.filter()操作无法执行，要么卡住要么静默失败。

解决方法：
在主进程中提前将每个pt对应的id和emb提取为本地数据，再传递给子进程处理，避免子进程直接操作Spark DataFrame。

2. RDD转换逻辑错误 #

原代码中flatMap(lambda x: np.float32(x))会将每个emb数组的元素拆分，导致item_emb变成一维列表，后续转numpy数组时形状错误（无法获取shape[1]），但子进程的异常不会主动抛出。

解决方法：
使用map替代flatMap，直接提取emb数组：

item_embs = [np.float32(row.emb) for row in rows]

3. 目标目录不存在 #

faiss.write_index不会自动创建目录，如果Faiss_Index目录不存在，会触发IO错误但异常被多进程吞掉。

解决方法：
在函数开头添加目录创建逻辑：

os.makedirs('Faiss_Index', exist_ok=True)

4. 缺失异常捕获 #

子进程中的异常无法直接在主进程显示，导致无法定位错误。

解决方法：
在处理函数中添加try-except块，打印错误信息。

修改后的完整代码 #

import os
import math
import numpy as np
import faiss
import multiprocessing as mp
from pyspark.sql import functions as f

# 主进程提前提取各pt对应的本地数据
pt_data = {}
for val in df.select('pt').distinct().collect():
    pt_val = val.pt
    # 一次性获取当前pt的所有id和emb
    rows = df.filter(f.col('pt') == pt_val).select('id', 'emb').collect()
    item_ids = [row.id for row in rows]
    item_embs = [np.float32(row.emb) for row in rows]
    pt_data[pt_val] = (item_ids, item_embs)

def process_pt(pt_item):
    pt_val, (item_ids, item_embs) = pt_item
    # 确保输出目录存在
    os.makedirs('Faiss_Index', exist_ok=True)
    
    try:
        item_ids_np = np.array(item_ids, dtype=np.int64)
        item_embs_np = np.array(item_embs)
        pool_size = item_embs_np.shape[0]
        
        # 跳过空数据的pt
        if pool_size == 0:
            print(f"Warning: pt {pt_val} has no data to process")
            return
        
        emb_size = item_embs_np.shape[1]
        nlist = int(math.sqrt(pool_size))
        quantizer = faiss.IndexFlatL2(emb_size)
        
        index = faiss.IndexIVFFlat(quantizer, emb_size, nlist)
        index.train(item_embs_np)
        index.add_with_ids(item_embs_np, item_ids_np)
        
        # 保存索引
        index_path = f'Faiss_Index/train_{pt_val}.index'
        faiss.write_index(index, index_path)
        print(f"Successfully saved index to {index_path}")
    except Exception as e:
        print(f"Error processing pt {pt_val}: {str(e)}")

# 多进程执行（必须加if __name__ == '__main__'避免多进程重复初始化）
if __name__ == '__main__':
    pool = mp.Pool(processes=mp.cpu_count()-1)
    pool.map(process_pt, pt_data.items())
    pool.close()
    pool.join()

额外注意事项 #

• 必须添加if __name__ == '__main__'：Python多进程在Windows系统下需要该判断，避免子进程重复执行主代码逻辑；Linux/macOS下虽非强制，但添加后更规范。
• 若数据量极大，主进程提前提取所有数据可能占用过多内存，此时建议改用Spark的groupBy('pt')结合mapGroups或pandas_udf处理，避免Python多进程的内存瓶颈。

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