make dask foldby 创建多个分区

Question

我试图了解 dask.foldby 的工作原理。 考虑以下代码。

tasks = [{"task_group": i // 10, "numbers": list(range(i, i + 10))} for i in range(1, 1000, 10)]
tb = db.from_sequence(tasks)

我创建了一个包含 100 件物品的 dask 包。

def aggregate_task(task):
    return np.array([sum(task['numbers'])] * 10000) # A relatively big result

def add(x, y):
    x = aggregate_task(x) if isinstance(x, dict) else x
    y = aggregate_task(y) if isinstance(y, dict) else y
    return x + y

res = tb.foldby(lambda task: task['task_group'] % 5, add, split_every=25)

然后我按某个分组 function 折叠项目，计算 5 个总和。

res.compute()
[(0, array([96100, 96100, 96100, ..., 96100, 96100, 96100])),
 (1, array([98100, 98100, 98100, ..., 98100, 98100, 98100])),
 (2, array([100100, 100100, 100100, ..., 100100, 100100, 100100])),
 (3, array([102100, 102100, 102100, ..., 102100, 102100, 102100])),
 (4, array([104100, 104100, 104100, ..., 104100, 104100, 104100]))]

当我查看 res 任务图时，我看到：

折叠的结果似乎存储在单个工作人员/单个分区中（我不太确定 1 个分区 = 1 个工作人员）。 我可以确认res只有一个分区。

>> res.npartitions
1

我的问题是：是否可以要求 dask 将每个 foldby 组的结果保存在单独的分区中（以及因此单独的工作人员）？ 我问的原因是，每个结果数组都很庞大，出于 memory 的原因，我想将它们放在不同的工作人员处。 它还可以将这些结果并行写入磁盘中的不同文件。

Answer 1

使用当前的 API，我认为这不可能在一个 go 中完成。 一种技巧/解决方法是创建过滤袋：

bag1 = tb.filter(lambda task: task['task_group'].isin([1,2,3]))
bag2 = tb.filter(lambda task: task['task_group'].isin([4,5]))

然后代码继续在过滤袋上应用.foldby 。