3. 并行计算入门

特征生成，说白了就是跟数据量较劲。单机跑特征，数据量一上来，CPU 就一个核在忙，其他核全闲着。我见过不少团队，特征脚本跑一整天，一看监控，CPU 利用率不到 20%。

这太浪费了。Python 的多进程，就是解决这个问题的。

3.1 为什么是 multiprocessing，不是 threading？

很多人一上来就问：Python 不是有 threading 吗？为什么不用线程？

嗯，这里有个关键点。Python 的 GIL（全局解释器锁）让多线程在 CPU 密集型任务上基本失效。你开 8 个线程，实际上还是只有一个核在工作。特征生成是典型的 CPU 密集型——你要算均值、方差、聚合统计，全是数值运算。

多进程就不一样了。每个进程有独立的 Python 解释器，有自己的 GIL。8 个进程，就能用满 8 个核。

3.2 最简单的并行：Pool.map

我个人习惯，先从 Pool.map 入手。它跟 Python 内置的 map 函数用法几乎一样，只是自动把任务分到多个进程。

看个例子。假设我们要对 100 万条用户行为数据生成滑动窗口特征：

import pandas as pd
import numpy as np
from multiprocessing import Pool
import time

def generate_features(user_id):
    """
    为单个用户生成特征
    我在项目中遇到过：如果函数参数是 DataFrame 切片，
    一定要确保数据是独立的，不要共享引用。
    """
    # 模拟耗时计算
    time.sleep(0.01)
    # 假设这里有一堆特征逻辑
    features = {
        'user_id': user_id,
        'avg_amount': np.random.randn(),
        'max_amount': np.random.randn(),
        'count_transactions': np.random.randint(1, 100)
    }
    return features

if __name__ == '__main__':
    user_ids = list(range(10000))
    
    # 单进程版本
    start = time.time()
    single_results = [generate_features(uid) for uid in user_ids]
    print(f'单进程耗时: {time.time() - start:.2f}s')
    
    # 多进程版本
    start = time.time()
    with Pool(processes=4) as pool:
        multi_results = pool.map(generate_features, user_ids)
    print(f'多进程耗时: {time.time() - start:.2f}s')

你想想看，这个例子跑下来，多进程基本能快 3-4 倍。为什么不是 4 倍？因为进程间通信、数据序列化都有开销。

3.3 数据分块：别让进程饿着

直接用 pool.map 有个问题：如果每个任务耗时差异很大，有的进程早早干完，有的还在忙。这叫「负载不均」。

我建议用 chunksize 参数。它控制每个进程一次拿多少个任务。

# 自动计算合适的 chunksize
chunksize = len(user_ids) // (pool._processes * 4)
results = pool.map(generate_features, user_ids, chunksize=chunksize)

为什么乘以 4？这是经验值。让每个进程一次拿 4 批任务，既能减少调度开销，又能保持负载均衡。

我曾经在一个风控特征项目里，没设 chunksize，结果 16 核的机器，15 个进程在等最后一个进程算完。设了 chunksize 之后，总耗时从 47 分钟降到 12 分钟。

3.4 共享数据：别踩序列化的坑

多进程之间不共享内存。数据要传给子进程，必须序列化（pickle）。

这里有个大坑：不要传大对象。比如把整个 DataFrame 传给每个进程。序列化 1GB 的数据，光传输就卡死你。

正确的做法是：

def generate_features(user_id, data_path):
    # 每个进程自己读数据
    df = pd.read_parquet(data_path)
    user_data = df[df['user_id'] == user_id]
    # ... 特征逻辑
    return features

if __name__ == '__main__':
    data_path = '/data/features.parquet'
    user_ids = list(range(10000))
    
    # 用 partial 固定 data_path 参数
    from functools import partial
    func = partial(generate_features, data_path=data_path)
    
    with Pool(4) as pool:
        results = pool.map(func, user_ids)

3.5 并行特征生成的完整流程

下面这张图，是我在项目中总结的并行特征生成标准流程：

3.6 实战：并行生成用户统计特征

来看一个更真实的场景。假设我们有用户交易数据，要按用户分组生成统计特征。

import pandas as pd
import numpy as np
from multiprocessing import Pool
from functools import partial

def generate_user_stats(user_df):
    """
    为单个用户生成统计特征
    注意：user_df 是 DataFrame 切片，必须保证是副本
    """
    features = {
        'user_id': user_df['user_id'].iloc[0],
        'total_amount': user_df['amount'].sum(),
        'avg_amount': user_df['amount'].mean(),
        'std_amount': user_df['amount'].std(),
        'max_amount': user_df['amount'].max(),
        'min_amount': user_df['amount'].min(),
        'transaction_count': len(user_df),
        'unique_merchants': user_df['merchant_id'].nunique(),
        'amount_p99': user_df['amount'].quantile(0.99)
    }
    return features

def parallel_feature_generation(df, n_processes=4):
    """
    并行特征生成主函数
    """
    # 按用户分组
    user_groups = [group for _, group in df.groupby('user_id')]
    
    print(f'总用户数: {len(user_groups)}')
    print(f'使用进程数: {n_processes}')
    
    # 并行计算
    with Pool(n_processes) as pool:
        # 用 imap_unordered 可以实时获取结果
        results = []
        for result in pool.imap_unordered(generate_user_stats, user_groups):
            results.append(result)
    
    # 合并结果
    feature_df = pd.DataFrame(results)
    return feature_df

if __name__ == '__main__':
    # 模拟数据
    np.random.seed(42)
    n_users = 10000
    n_transactions = 500000
    
    df = pd.DataFrame({
        'user_id': np.random.randint(0, n_users, n_transactions),
        'amount': np.random.exponential(100, n_transactions),
        'merchant_id': np.random.randint(0, 500, n_transactions),
        'timestamp': pd.date_range('2024-01-01', periods=n_transactions, freq='T')
    })
    
    # 跑并行特征生成
    feature_df = parallel_feature_generation(df, n_processes=8)
    print(f'生成特征维度: {feature_df.shape}')
    print(feature_df.head())

3.7 避坑总结

最后，我把多进程特征生成中常见的坑列出来。这些都是我踩过的：

问题	表现	解决方案
内存爆炸	每个进程复制大 DataFrame	只传索引，子进程自己读数据
负载不均	部分进程空闲等待	设置合适的 chunksize
序列化失败	报错：Can't pickle local object	函数定义在模块顶层，不要用 lambda
死锁	程序卡住不退出	用 with 语句管理 Pool，确保资源释放
结果顺序错乱	特征与用户 ID 不匹配	用 imap_unordered 或手动排序

多进程不是银弹。数据量小的时候，进程间通信的开销可能比计算本身还大。我一般有个经验法则：数据量少于 10 万行，单进程就够了；超过 100 万行，多进程能省下几小时。