Dask map_partitions未利用全部Worker：Dask read_csv与Pandas转Dask的差异及优化

问题描述 #

我有一项CPU密集型任务，希望尽可能利用Dask的所有Worker，但遇到两种不同行为：

• 使用Dask的read_csv读取CSV后，通过map_partitions处理时仅用1个Worker
• 使用Pandas的read_csv读取后转成Dask DataFrame，所有8个Worker都会被充分利用

需要解释两者差异，同时实现无需转换步骤、直接用Dask read_csv即可充分利用Worker的方法。

代码示例 #

import dask as d
import pandas as pd

def fWrapper(x):
            p = doSomething(x.ADDRESS, param)
            return(pd.DataFrame(p, columns=["ADDRESS", "DATA","TOKEN", "CLASS"]))

# 仅使用1个Worker，而非可用的8个
dask_df = d.dataframe('path\\to\\file')
dask_df.set_index(UID, npartitions = 8,   drop = False)
ddf2 = dask_df.map_partitions(fWrapper, meta={"ADDRESS" : object, "DATA" : object, "TOKEN" : object, "CLASS" : object}).compute() 

# 使用全部8个Worker
df = pd.read_csv('path\\to\\file')
df.set_index('UID', drop=False)
dask_df2 =d.dataframe.from_pandas(df, npartitions=dask_params['df_npartitions'], sort=True)
ddf3 = dask_df2.map_partitions(fWrapper, meta={"ADDRESS" : object, "DATA" : object, "TOKEN" : object, "CLASS" : object}).compute() 

行为差异原因 #

• Dask read_csv的分区与赋值问题
  Dask read_csv默认根据文件大小自动分区，若文件是单个小文件，只会创建1个分区。后续调用set_index时，你指定了npartitions=8但未将结果赋值回变量，原dask_df仍保持1个分区。map_partitions按分区并行执行，1个分区自然只会占用1个Worker。
• Pandas转Dask的分区逻辑
  Pandas读取整个文件后，from_pandas通过npartitions=8直接将数据切分为8个有序分区，map_partitions处理8个分区时，可同时调度到8个Worker，充分利用资源。

直接用Dask read_csv实现多Worker利用的方案 #

核心是修复分区操作的赋值问题，同时可在read_csv阶段主动控制分区：

优化后的代码 #

import dask.dataframe as dd
import pandas as pd

def fWrapper(x):
    p = doSomething(x.ADDRESS, param)
    return pd.DataFrame(p, columns=["ADDRESS", "DATA","TOKEN", "CLASS"])

# 方案1：read_csv时指定分区，再设置索引
dask_df = dd.read_csv('path\\to\\file', npartitions=8)  # 直接指定分区数
dask_df = dask_df.set_index('UID', drop=False)  # 必须将重新分区的结果赋值回变量
ddf2 = dask_df.map_partitions(fWrapper, meta={"ADDRESS": object, "DATA": object, "TOKEN": object, "CLASS": object}).compute()

# 方案2：无需排序索引时用repartition更高效
dask_df = dd.read_csv('path\\to\\file')
dask_df = dask_df.repartition(npartitions=8)  # 快速拆分分区，避免全局排序开销
ddf2 = dask_df.map_partitions(fWrapper, meta={"ADDRESS": object, "DATA": object, "TOKEN": object, "CLASS": object}).compute()

关键注意事项 #

1. 保留分区操作结果：Dask的DataFrame操作是惰性的且返回新对象，原对象不会被修改，必须将set_index或repartition的结果赋值给变量，后续操作才会基于新分区。
2. read_csv阶段主动分区：大文件可通过blocksize参数让Dask自动划分多分区；小文件直接指定npartitions强制多分区，CPU密集型任务下这点开销可忽略。
3. 选择合适的分区方式：若任务不需要基于UID的有序索引，优先用repartition，避免set_index的全局排序开销。

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