使用 Dask 和 Xarray 改进 GroupBy.map

本文最初发表于 Xarray 博客。

使用由 Dask 数组 支持的 Xarray 运行大规模 GroupBy-Map 模式，是许多典型地理空间工作负载的重要组成部分。去趋势化是需要此模式的一个非常常见的操作。

在本文中，我们将探讨过去为什么以及如何给 Xarray 用户带来了如此多的陷阱，以及我们如何通过对 Dask 子集选择数据方式的一些更改来提高性能和可伸缩性。

什么是 GroupBy.map？

GroupBy.map 允许您应用接受并返回 Xarray 对象的用户定义函数 (UDF)。UDF 将接收包含与单个组对应的 Dask 数组的 Xarray 对象（可以是 Dataset 或 DataArray）。Groupby.reduce 与之非常相似，因为它也应用 UDF，但在这种情况下，UDF 将接收底层 Dask 数组，而不是 Xarray 对象。

应用

考虑一个典型的需要应用去趋势化步骤的工作流程。您可能希望通过去除随时间变化的趋势来平滑数据。这是气候科学中常见的操作，通常看起来大致如下：

def detrending_step(arr: DataArray) -> DataArray:
    # important: the rolling operation is applied within a group
    return arr - arr.rolling(time=30, min_periods=1).mean()

data.groupby("time.dayofyear").map(detrending_step)

我们按一年中的某一天进行分组，然后计算特定日期的 30 年窗口的滚动平均值。

我们的示例将在一个 1 TiB 的数组上运行，包含 64 年的数据，结构如下：

数组不算太大，数据块大小也合理。

问题

这个通用应用看起来很简单。按一年中的某一天进行分组，并对每个组应用一个 UDF。这个应用中有一些陷阱，可能导致操作结果无法使用。我们的数组是按时间排序的，这意味着我们必须从数组中的许多不同区域中挑选条目来创建一个组（对应于一年中的某一天）。挑选每年的同一天本质上是一个步长为 365 的切片操作。

我们的示例在时间轴上包含一年的数据在一个数据块中。对于任何需要随机访问数据条目的工作负载，都存在这个问题。这种特定的访问模式意味着我们必须从每个数据块中挑选一个值，这效率很低。右侧显示了我们正在操作的各个组。

这种模式的主要问题之一是 Dask 将为每个时间条目创建一个输出数据块，例如，每个组将包含与年份数量一样多的数据块。

这导致数据块数量大幅增加

这个简单的操作将数据块数量从 5000 增加到接近 200 万。每个数据块只有几百千字节的数据。这太糟糕了！

Dask 计算通常与您拥有的数据块数量成正比。将数据块数量增加如此大的倍数是灾难性的。每个后续操作，例如简单的 a-b，将创建 200 万个额外的任务。

用户唯一的解决方法是事后重新分块到更合理的大小，但这仍然将极其昂贵的索引操作保留在计算图中。

请注意，这是 [flox](https://xarray.dev/blog/flox) 通过使用并行原生算法解决的底层问题，适用于诸如 .mean() 之类的聚合操作，以避免对每个组进行索引的开销。

数据选择算法的改进

Dask 选择数据的方法客观上相当糟糕。对底层算法的重写使我们能够获得更稳健的结果。新算法在如何从每个数据块中挑选值方面更加智能，但最重要的是，它会尽量保持输入数据块大小不变。

对于我们的初始示例，它将把每个组放入一个数据块中。这意味着沿时间轴的数据块数量将等于组的数量，即 365。

该算法将数据块数量从 200 万减少到大约 3 万，这是一个巨大的改进，也是 Dask 轻松可以处理的规模。计算图现在小得多，后续操作也将运行得快得多。

这一改进将帮助我们上面列出的所有操作，并使规模比以前更加可靠。该算法在 Dask 和 Xarray 中得到了广泛应用，因此影响了许多方法。

接下来是什么？

Xarray 为 groupby(...).map(...) 一次选择一个组，即每个组需要一个操作。如果数据集包含非常多的组，这将损害可伸缩性，因为计算会创建一个非常昂贵的图。当前正在努力实现基于混洗的替代 API 来规避这个问题。当前的拉取请求可在此处获取：此处。

通过索引导致输出数据块的碎片化会损害所有以随机模式选择数据的工作流程。这也包括：

• 如果您没有显式使用切片，则包括 .sel
• .isel
• .sortby
• groupby(...).quantile()
• 等等。

我们预计所有这些工作负载现在都将得到显著改进。

此外，Dask 改进了很多 与增加数据块大小或工作负载周期中的数据块碎片化相关的内容，未来还会有更多改进。这将帮助很多用户获得更好、更可靠的性能。