Dask延迟总和被杀死但有足够的资源

Question

我正在创建一个读取整个文件夹的函数，创建一个 Dask 数据帧，然后处理该数据帧的分区并对结果求和，如下所示：

import dask.dataframe as dd
from dask import delayed, compute

def partitions_func(folder):
    df = dd.read_csv(f'{folder}/*.csv')
    partial_results = []
    for partition in df.partitions:
        partial = another_function(partition)
        partial_results.append(partial)
    total = delayed(sum)(partial_results)
    return total

partitions_func (another_function) 中调用的函数也有延迟。

@delayed
def another_function(partition):
    # Partition processing
    return result

我检查了，处理过程中创建的变量都很小，所以它们应该不会引起任何问题。分区可以很大，但不能超过可用 RAM。

当我执行partitions_func(folder) 时，进程被杀死。起初，我认为问题在于有两个delayed，一个在another_function，一个在delayed(sum)。

从 another_function 中删除 delayed 装饰器会导致问题，因为参数是 Dask 数据帧，您不能执行像 tolist() 这样的操作。我尝试从sum 中删除delayed，因为我认为这可能是并行化和可用资源的问题，但进程也会被终止。

但是，我知道有 5 个分区。如果我从 partitions_func 中删除语句 total = delayed(sum)(partial_results) 并改为“手动”计算总和，一切都会按预期工作：

total = partial_results[0].compute() + partial_results[1].compute() + partial_results[2].compute() \
        + partial_results[3].compute() + partial_results[4].compute()

谢谢！

Answer 1

Dask 数据帧会创建一系列延迟对象，因此当您调用延迟函数 another_function 时，它会变成嵌套延迟，dask.compute 将无法处理它。一种选择是使用.map_partitions()，典型的例子是df.map_partitions(len).compute()，它将计算每个分区的长度。因此，如果您可以重写another_function 以接受pandas 数据帧，并删除延迟的装饰器，那么您的代码将大致如下所示：

df = dd.read_csv(f'{folder}/*.csv')
total = df.map_partitions(another_function)

现在total 是一个延迟对象，您可以将其传递给dask.compute（或简单地运行total = df.map_partitions(another_function).compute()）。