Dask map_partitions函数出现额外调用异常，寻求解决方法

将Pandas DataFrame转换为1分区的Dask DataFrame后，调用map_partitions()时，函数会被调用两次；如果设置为5个分区，则被调用6次——总是会额外多调用一次，传入无法识别的记录，且这些记录不会出现在最终输出里。首次额外调用传入的分区包含2条记录，这引发了其他问题。

环境详情 #

python 3.9.18
dask 2024.8.0
pandas 2.0.3

测试代码 #

import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame({
    'a': list(range(100))
})
ddf = dd.from_pandas(df, npartitions=1)

def some_func(df):
    print(df.shape)
    print(df.head())
    return df

ddf = ddf.map_partitions(some_func)
print(ddf.compute().shape)

运行输出 #

(2, 1)
  a
0 1
1 1
(100, 1)
  a
0 0
1 1
2 2
3 3
4 4
(100, 1)

解决方案 #

这个额外调用是Dask的元数据推断机制导致的：Dask默认会生成一个小样本DataFrame（你看到的两条记录），传入你的函数来获取返回值的结构（列名、数据类型等），以此确定最终输出的元数据，所以会多调用一次函数。

要避免这个问题，推荐两种方法：

方法1：显式指定meta参数 #

直接告诉Dask输出的元数据结构，这样它就不需要调用函数推断了。比如用原Pandas DataFrame的结构作为meta：

ddf = ddf.map_partitions(some_func, meta=df)

如果需要更严谨的类型控制，可以手动构造meta：

meta = pd.DataFrame(columns=['a'], dtype='int64')
ddf = ddf.map_partitions(some_func, meta=meta)

方法2：禁用全局元数据推断 #

通过Dask配置关闭查询规划（元数据推断的一部分），但这是全局设置，可能影响其他Dask操作，谨慎使用：

from dask import config
config.set({"dataframe.query-planning": False})

显式指定meta是最稳妥的方式，既能避免额外调用，也能保证Dask对输出结构的认知准确。

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