如何用Dask实现带分块的易并行工作负载并应用自定义分块函数

你可以使用dask.array.map_blocks实现按分块调用自定义函数，该方法会将你的自定义函数作用在输入数组的每一组对应分块上，既不会逐元素调用产生额外开销，也不会触发全量数据加载，可以完美替代原有数组运算。

完整实现代码 #

from dask.diagnostics import ProgressBar
import zarr
import dask.array as da
from time import sleep

# 你的自定义分块处理函数
def func(a, b):
    sleep(0.001 + len(a) * 1e-5)
    return a + b + 1

d1 = da.ones((1024, 1024, 1024, 128), chunks=64)
d2 = da.ones((1024, 1024, 1024, 128), chunks=64)

# 用map_blocks按分块应用自定义函数
d3 = da.map_blocks(func, d1, d2, dtype=d1.dtype, chunks=d1.chunks)
                                                                                                                                                                                                                      
with ProgressBar():
    d3.to_zarr("three.zarr")

关键参数说明 #

• func：你自定义的分块处理函数，输入为两个NumPy数组（对应d1、d2的同位置分块），输出为处理后的NumPy数组
• dtype：显式声明输出数组的类型，Dask需要该信息构建计算图
• chunks：声明输出数组的分块结构，这里和输入数组保持一致即可，如果你的函数会改变分块形状，按需调整即可

方案优势 #

• 每个分块仅调用一次自定义函数，完全避免逐元素调用的额外开销，适配低计算量的自定义函数
• 全程使用Dask数组的懒加载机制，仅在计算时加载对应分块到内存，不会触发全量数据加载，避免OOM
• 不需要转换为其他数据结构（如Bag），没有额外的序列化、类型转换开销，性能更优

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