4、并行计算进阶：Joblib与Dask分布式特征工程实战

单机并行做到极致后，你会发现瓶颈不在CPU，而在内存。我遇到过好几次，特征工程跑了一整天，最后OOM崩了——那种感觉，真不好受。这时候就需要分布式计算上场了。

这一章，我们聊两个工具：Joblib 和 Dask。前者轻量，适合单机多核；后者能上集群，处理TB级数据。说白了，一个是你随身带的瑞士军刀，另一个是工程队用的挖掘机。

4.1 Joblib：轻量级并行利器

Joblib 是 scikit-learn 生态里的老熟人。很多人只用它做模型持久化，其实它的并行能力被严重低估了。

核心优势：零学习成本、自动内存管理、适合中等规模特征生成。

4.1.1 基础用法：Parallel 与 delayed

先看一个典型场景：对100个特征列做窗口聚合。单循环太慢，我们上并行。

from joblib import Parallel, delayed
import pandas as pd
import numpy as np

def compute_rolling_feature(df, col, window):
    """计算单个特征的滚动统计量"""
    return df[col].rolling(window=window).mean()

# 特征列列表
feature_cols = [f'feature_{i}' for i in range(100)]
windows = [5, 10, 20]

# 并行执行
results = Parallel(n_jobs=-1, verbose=10)(
    delayed(compute_rolling_feature)(df, col, w)
    for col in feature_cols
    for w in windows
)

这里有个坑：n_jobs=-1 表示用满所有CPU核心。但如果你在共享服务器上跑，我建议设成 n_jobs=os.cpu_count() - 1，留一个核心给系统，不然别人会找你麻烦。

4.1.2 内存管理：batch_size 与 pre_dispatch

Joblib 默认会把所有任务一次性提交。如果特征数量上万，内存会瞬间爆炸。我曾经吃过这个亏——32GB内存的机器，直接卡死。

解决方案：

results = Parallel(
    n_jobs=8,
    batch_size=10,          # 每批只提交10个任务
    pre_dispatch='2*n_jobs', # 预调度任务数
    verbose=10
)(
    delayed(compute_rolling_feature)(df, col, w)
    for col in feature_cols
    for w in windows
)

我的经验：batch_size 设成 2~4 倍的核心数，效果最好。pre_dispatch 控制预调度数量，避免内存被任务描述撑爆。

4.1.3 缓存机制：Memorize 加速重复计算

特征工程经常要反复调试。每次重跑都重新计算，太浪费了。Joblib 的 Memory 模块可以缓存中间结果。

from joblib import Memory

# 指定缓存目录
memory = Memory(location='./cache', verbose=0)

@memory.cache
def expensive_feature_transform(df, params):
    """耗时特征变换，结果会被缓存到磁盘"""
    # 模拟复杂计算
    result = df.groupby('user_id').rolling(
        window=params['window']
    ).agg(params['agg_func'])
    return result

# 第一次运行：实际计算
result1 = expensive_feature_transform(df, {'window': 10, 'agg_func': 'mean'})

# 第二次运行：直接从缓存加载（秒出）
result2 = expensive_feature_transform(df, {'window': 10, 'agg_func': 'mean'})

注意：缓存依赖输入参数的哈希值。如果 DataFrame 内容变了但变量名没变，缓存不会自动失效。我建议在参数里加一个版本号，或者定期清理缓存目录。

4.2 Dask：从单机到集群的平滑过渡

当数据量超过单机内存时，Joblib 就力不从心了。Dask 的设计理念很聪明——它把 Pandas/NumPy 的 API 几乎原样搬到了分布式环境。你不需要学一套新语法。

4.2.1 Dask DataFrame：懒加载与任务图

Dask 的核心是懒执行。你写的操作不会立即计算，而是构建一个任务图。只有调用 .compute() 时才会真正执行。

import dask.dataframe as dd

# 读取多个CSV文件（懒加载）
ddf = dd.read_csv('data/part_*.csv')

# 定义特征生成逻辑（还是懒加载）
ddf['rolling_mean'] = ddf.groupby('user_id')['amount'].transform(
    lambda x: x.rolling(10).mean()
)

# 添加更多特征
ddf['lag_1'] = ddf.groupby('user_id')['amount'].shift(1)
ddf['lag_7'] = ddf.groupby('user_id')['amount'].shift(7)

# 真正触发计算
result = ddf.compute()

为什么会这样？因为 Dask 把每个分区当成一个 Pandas DataFrame，操作在分区内并行。你想想看，如果数据有100个分区，每个分区独立计算，速度自然快。

4.2.2 分区策略：决定性能的关键

分区数直接影响并行度。分区太少，资源闲置；分区太多，调度开销大。

分区策略 适用场景 注意事项
按文件分区 数据已按逻辑分片 文件大小不均会导致倾斜
按列重分区 需要按某列聚合 shuffle 开销大，慎用
按行数重分区 均匀分布数据 每个分区建议 100~200MB
自动分区 快速原型 性能不一定最优 
# 手动设置分区数
ddf = ddf.repartition(npartitions=20)

# 或者按分区大小控制
ddf = ddf.repartition(partition_size='100MB')

避坑指南：我曾经把一个10GB的数据集分成500个分区，结果调度器比计算还慢。后来改成50个分区，速度提升了3倍。经验公式：分区数 ≈ 核心数 × 2~4。

4.2.3 分布式特征工程实战

下面是一个完整的分布式特征工程流程。假设我们有用户行为日志，需要生成用户画像特征。

import dask.dataframe as dd
from dask.distributed import Client

# 启动分布式集群（本地模式）
client = Client(n_workers=4, threads_per_worker=2, memory_limit='4GB')

# 读取数据
ddf = dd.read_parquet('user_behavior/*.parquet')

# 特征1：用户行为统计
user_stats = ddf.groupby('user_id').agg({
    'amount': ['sum', 'mean', 'std', 'count'],
    'category': lambda x: x.nunique(),
    'timestamp': ['min', 'max']
}).compute()

# 特征2：时间窗口特征（使用 map_partitions）
def window_features(partition):
    """对单个分区计算窗口特征"""
    partition = partition.sort_values('timestamp')
    partition['amount_rolling_7d'] = (
        partition.groupby('user_id')['amount']
        .rolling('7D', min_periods=1)
        .mean()
        .values
    )
    return partition

ddf = ddf.map_partitions(window_features)

# 特征3：跨分区特征（需要 shuffle）
ddf['user_total_amount'] = ddf.groupby('user_id')['amount'].transform('sum')

# 触发计算
result = ddf.compute()

4.2.4 集群部署与资源调优

单机跑通了，怎么上集群？Dask 支持多种调度器：

  • LocalCluster：单机多进程，适合开发和测试
  • SLURMCluster：适配 SLURM 调度系统，常见于学术集群
  • KubernetesCluster：容器化部署，弹性伸缩
  • YARNCluster：Hadoop 生态集成
# 集群模式示例（SLURM）
from dask_jobqueue import SLURMCluster

cluster = SLURMCluster(
    queue='gpu',
    cores=8,
    memory='32GB',
    walltime='02:00:00'
)
cluster.scale(jobs=10)  # 申请10个节点

client = Client(cluster)

# 现在所有计算都在集群上执行
ddf = dd.read_parquet('/shared/data/*.parquet')
result = ddf.groupby('user_id').agg(...).compute()

性能陷阱：Dask 的 shuffle 操作（如 groupby、join）会触发数据重分布，这是最耗时的环节。我建议：
1. 尽量在分区内完成计算，避免 shuffle
2. 必须 shuffle 时，提前按 key 分区
3. 使用 set_index() 预排序，减少后续 shuffle 开销

4.3 Joblib vs Dask：如何选择？

很多学员问我：到底用哪个？我的回答是：看数据量。

维度 Joblib Dask
数据规模 单机内存可容纳 TB级，超内存
学习成本 极低，几分钟上手 中等，需理解懒执行
调度开销 几乎为零 有额外通信成本
缓存机制 内置 Memory 需配合外部存储
集群支持 不支持 原生支持
适用场景 特征工程调试、小批量 生产流水线、大规模

我个人习惯：原型阶段用 Joblib，快速验证特征有效性。上线前迁移到 Dask，保证生产环境的稳定性和扩展性。

4.4 本章知识体系

下面这张图总结了 Joblib 和 Dask 在特征工程中的定位与协作关系。

特征工程并行计算体系 原始数据 数据量 < 内存？ 是 否 Joblib 并行 Dask 分布式 Parallel + delayed Memory 缓存 batch_size 控制 懒执行 + 任务图 分区策略与 shuffle 集群部署（SLURM/K8s） 生成特征矩阵 核心原则：数据不动代码动，分区内计算优先

嗯，到这里你应该能感受到：Joblib 和 Dask 不是二选一的关系，而是互补的。小规模快速迭代用 Joblib，大规模生产部署上 Dask。掌握了这两个工具，大部分特征工程的性能问题都能迎刃而解。

← 上一章 📖 返回目录 下一章 →