Dask中set_index(col, compute=True)的compute参数作用是什么？

现象原理 #

set_index的compute参数目前仅在shuffle='disk'模式下生效，其他shuffle模式未适配该参数，这是Dask当前版本的实现逻辑决定的：

• 当使用默认的shuffle='tasks'模式时，set_index全程只做任务图构建，代码路径中没有处理compute参数的逻辑，不管该参数设为True还是False，都只会返回懒加载的Dask DataFrame，不会触发实际计算。
• 当使用shuffle='disk'模式时，底层走磁盘重排逻辑，该逻辑中会判断compute参数：如果设为True，会立即触发shuffle任务计算，把重排后的分区数据写入本地临时磁盘文件，最终返回的Dask DataFrame分区直接指向这些已落地的文件，不需要再重复执行shuffle流程。

现有官方文档对compute参数的说明未明确标注其仅适配磁盘shuffle场景，属于文档描述疏漏。

验证compute=True已执行的方法 #

你可以通过以下几种方式确认操作已生效：

1. 查看任务执行记录：本地运行Dask时，访问默认的调度面板可看到shuffle相关任务已经执行完成，无待执行任务。
2. 对比任务图结构：对返回的Dask DataFrame执行.visualize()，如果compute=True生效，任务图仅包含读取磁盘文件的节点，不存在shuffle计算相关的任务链；未生效的情况下会展示完整的shuffle重排任务图。
3. 耗时对比：通过时间统计可以直观看到差异，示例代码如下：

import time
import dask.datasets

df = dask.datasets.timeseries()

# 开启compute=True的场景
start = time.time()
df_disk = df.set_index(
    'name', 
    divisions=('Alice', 'Michael', 'Zelda'), 
    shuffle='disk',
    compute=True
)
print(f"set_index执行耗时：{time.time() - start:.2f}s") # 输出明显的计算耗时

# 后续聚合操作耗时
start = time.time()
print(df_disk.count().compute())
print(f"后续聚合操作耗时：{time.time() - start:.2f}s") # 耗时极短，远低于set_index耗时

4. 查看临时磁盘文件：Dask磁盘shuffle默认会在系统临时目录下生成前缀为dask-scratch-space的临时文件夹，执行compute=True时该目录下会生成对应分区的数据文件，这些文件会在DataFrame对象被回收或进程退出后自动清理。

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