Dask数组循环填充耗时递增问题排查及优化方案咨询

我拥有约1000个尺寸一致（约4000×4000）的小数组，以及一个通过dask.array.zeros创建的大数组（尺寸约1e5×1e5），其分块大小与小数组量级相当。需要将所有小数组合并到大数组中，具体方式是将每个小数组求和到大数组对应索引区域，且这些索引区域与分块边界不重合。

采用循环遍历小数组、逐个执行索引相加操作的方式时，最初几次迭代速度尚可，但很快进程停滞，迭代时间呈二次方增长。需注意此时尚未发生实际数据移动，所有操作都是惰性的，小数组从磁盘惰性加载，大数组也是惰性零数组。

以下代码可复现该问题（复用同一个小数组，不影响现象）：

import dask.array
import time
import numpy as np

bgsize = 2**18 # big array size
chsize = 2**13 # chunk size
N = 4096 # small array size

da = dask.array.zeros((bgsize, bgsize), chunks=(chsize, chsize), dtype="float16")
arr = dask.array.random.random((N, N))


t0 = time.time()
for i in range(2000):
    tprev = time.time()
    ip0 = np.random.randint(0, bgsize-N)
    ip1 = np.random.randint(0, bgsize-N)
    da[slice(ip0, ip0+N), slice(ip1, ip1+N)] += arr
    t = time.time()
    print(i, "done", t - tprev, t - t0)

初始条件是磁盘上有一个存储所有小数组的多维Zarr数组，最终目标是尽快将包含所有（可能重叠的）小数组贡献的大数组写入磁盘，本质是将磁盘上的多维数据转换为二维数据，但小数组间存在任意偏移，无法直接堆叠。此外，每次循环迭代将单个小数组写入磁盘不可行，因为求和操作需要先从磁盘读取目标区域的现有数据。

性能下降的核心原因 #

• 计算图指数级膨胀：Dask是惰性计算框架，每次执行da[slice] += arr时，不会立即执行计算，而是在内部计算图中新增一个加法任务。由于切片与分块边界不重合，每个新任务都需要依赖大数组对应区域的所有历史任务节点。随着循环次数增加，计算图的节点数呈**O(n²)**增长，每次更新计算图时需要遍历和处理的节点数量越来越多，直接导致迭代时间二次方上升。
• 跨分块依赖加剧复杂度：小数组尺寸与大数组分块量级相当，但切片偏移不与分块对齐，每个加法操作会跨越多块大数组分块，进一步增加了任务间的依赖关系，放大了计算图构建的开销。

解决方案 #

1. 批量收集任务，延迟构建计算图 #

不在循环中逐次修改大数组，先收集所有小数组的偏移信息，最后一次性构建加法操作的计算图，避免计算图膨胀：

import dask.array as da
import numpy as np
from dask.utils import reduce

bgsize = 2**18
chsize = 2**13
N = 4096
num_arrays = 2000

# 初始化大数组
big_da = da.zeros((bgsize, bgsize), chunks=(chsize, chsize), dtype="float16")
# 模拟从Zarr加载的小数组
small_da = da.random.random((N, N))

# 收集所有偏移信息
offsets = []
for _ in range(num_arrays):
    ip0 = np.random.randint(0, bgsize - N)
    ip1 = np.random.randint(0, bgsize - N)
    offsets.append((ip0, ip1))

# 定义单个小数组的累加逻辑
def add_small_to_big(big, offset):
    ip0, ip1 = offset
    return big.at[ip0:ip0+N, ip1:ip1+N].add(small_da)

# 批量合并所有累加操作
final_da = reduce(lambda acc, off: add_small_to_big(acc, off), offsets, initial=big_da)

# 执行计算并写入Zarr
final_da.to_zarr("output.zarr", overwrite=True)

• 优势：reduce以线性方式合并任务，计算图节点数从O(n²)降至O(n)，彻底解决二次方增长问题。

2. 分组处理任务，控制计算图规模 #

如果小数组数量极大，可将任务分成若干组，每组构建一次计算图，避免单次任务过多导致的内存和构建开销：

import dask.array as da
import numpy as np

bgsize = 2**18
chsize = 2**13
N = 4096
num_arrays = 2000
group_size = 200  # 每组处理200个小数组

big_da = da.zeros((bgsize, bgsize), chunks=(chsize, chsize), dtype="float16")
small_da = da.random.random((N, N))

# 生成所有偏移
offsets = [
    (np.random.randint(0, bgsize-N), np.random.randint(0, bgsize-N))
    for _ in range(num_arrays)
]

# 分组执行累加
for i in range(0, num_arrays, group_size):
    group_offsets = offsets[i:i+group_size]
    for off in group_offsets:
        ip0, ip1 = off
        big_da = big_da.at[ip0:ip0+N, ip1:ip1+N].add(small_da)
    # 可选：persist当前计算图，避免后续构建压力
    # big_da = big_da.persist()

big_da.to_zarr("output.zarr", overwrite=True)

• 优势：通过分组平衡了计算图规模和构建速度，适合超大规模的小数组集合。

3. 分块级聚合（偏移可映射场景） #

如果能将小数组的偏移映射到大数组的分块上，可先按分块聚合所有需要累加的小数组切片，再一次性写入对应分块：

1. 遍历所有小数组，记录每个大数组分块需要接收的小数组切片（即小数组在该分块内的部分）
2. 对每个分块，将所有对应的小数组切片求和，再写入大数组的分块

• 优势：彻底避免跨分块的任务依赖，计算图复杂度最低，适合偏移分布均匀的场景。

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