如何为Dask DataFrame创建自定义计算图并实现高效并行化？

其实你没必要绕这么大弯子——Dask DataFrame本身就内置了并行化优化，你的问题主要是对它的任务调度和操作模型理解有点偏差。咱们一步步拆解，找到最简洁高效的方案：

为什么你的初始代码看起来串行？ #

你最开始的代码：

import dask.dataframe as dd
ddf = dd.read_parquet('dataframe')
ddf['a'] = ddf['a'] + 1
ddf['b'] = ddf['b'] + 2
ddf.visualize()

从可视化的任务图看好像是串行，但实际执行时，Dask会自动并行处理独立的任务：

• 不同分区的a+1和b+2操作会被并行调度到不同的工作进程/线程。
• 同一分区内的两个列操作，Dask的优化器其实可以合并为一个任务（或者在执行时并行计算，因为它们没有依赖关系）。
  可视化图的“串行”只是任务的逻辑依赖展示，不代表实际执行的顺序。

为什么手动用delayed会出问题？ #

你尝试用delayed包装函数的思路是对的，但用错了场景：

• delayed是用来包装普通Python函数的，而Dask DataFrame本身已经是延迟执行的并行对象，用delayed包装后会把它转换成Delayed类型，失去DataFrame的原生优化方法，导致需要多次compute才能拿到结果。
• 当你把ddf[['a']]传给delayed函数时，Dask会自动将分区的Pandas DataFrame传入函数（这是Dask的分区调度逻辑），所以你看到的输入是Pandas DataFrame而非Dask DataFrame，这是正常的，但手动包装反而增加了复杂度。

最优并行化方案 #

方案1：用assign批量执行列操作（最简洁） #

dd.assign可以同时定义多个列的转换，Dask会自动优化任务图，确保独立操作并行执行：

import dask.dataframe as dd

ddf = dd.read_parquet('dataframe')
# 同时对a和b进行增量操作
ddf = ddf.assign(
    a=ddf['a'] + 1,
    b=ddf['b'] + 2
)
ddf.visualize()
ddf.compute()

这个方法代码简洁，而且Dask会自动处理并行调度，不需要手动干预。

方案2：用map_partitions（适合复杂分区操作） #

如果你的列操作更复杂（比如需要对整个分区做多个逻辑处理），用map_partitions可以将操作打包到每个分区的单个任务中，减少任务数量，提升效率：

import dask.dataframe as dd

def process_partition(df):
    # 同一分区内的所有操作在这里完成
    df['a'] = df['a'] + 1
    df['b'] = df['b'] + 2
    # 还可以添加其他复杂操作
    return df

ddf = dd.read_parquet('dataframe')
ddf = ddf.map_partitions(process_partition)
ddf.compute()

这种方式下，每个分区只生成一个任务，不同分区的任务会被并行执行，既保证了并行性，又减少了任务调度的开销。

总结最佳实践 #

• 优先使用Dask DataFrame的原生操作（赋值、assign、map_partitions），不要手动用delayed包装DataFrame对象——Dask已经帮你做好了并行优化。
• 任务图的可视化只是逻辑依赖展示，实际执行时Dask会自动调度并行任务，你可以通过观察CPU使用率来确认并行效果。
• 如果需要对多个列做独立的简单操作，assign是最简洁的选择；如果是复杂的分区级操作，map_partitions效率更高。

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