如何将 GRIB 文件目录加载到 Dask 数组中

Question

假设我有一个包含数千个 GRIB 文件的目录。我想将这些文件加载到一个 dask 数组中，以便我可以查询它们。我该怎么做呢？下面的尝试似乎可行，但它需要打开每个 GRIB 文件，并且运行和我所有的记忆都需要很长时间。一定有更好的办法。

我的尝试：

import dask.array as da
from dask import delayed
import gdal
import glob
import os


def load(filedir):
    files = sorted(glob.glob(os.path.join(filedir, '*.grb')))
    data = [da.from_array(gdal.Open(f).ReadAsArray(), chunks=[500,500,500], name=f) for f in files]
    return da.stack(data, axis=0)

file_dir = ...
array = load(file_dir)

Answer 1

最好的方法是使用 dask.delayed. In this case, you'd create a delayed function to read the array, and then compose a dask array from those delayed objects using the da.from_delayed 函数。大致如下：

# This function isn't run until compute time
@dask.delayed(pure=True)
def load(file):
    return gdal.Open(file).ReadAsArray()


# Create several delayed objects, then turn each into a dask
# array. Note that you need to know the shape and dtype of each
# file
data = [da.from_delayed(load(f), shape=shape_of_f, dtype=dtype_of_f)
        for f in files]

x = da.stack(data, axis=0)

请注意，这使得加载每个文件成为一个任务。如果单个文件很大，您可能希望自己在 load 函数中将它们分块。我对 gdal 不熟悉，但从 ReadAsArray 方法的简要介绍来看，这可能适用于 xoff/yoff/xsize/ysize参数（不确定）。您必须自己编写此代码，但对于大文件来说它的性能可能更高。

或者您可以使用上面的代码，然后调用 rechunk 重新分块为更小的块。这仍然会导致在单个任务中读取每个文件，但后续步骤可以处理更小的块。这是否值得取决于您个人文件的大小。

x = x.rechunk((500, 500, 500))  # or whatever chunks you want

如何将 GRIB 文件目录加载到 Dask 数组中

How to load a directory of GRIB files into a Dask array

gdal

grib

dask