如何将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.在这种情况下,您将创建一个延迟函数来读取数组,然后delayed使用该da.from_delayed函数从这些对象组成一个dask数组.有点像:

# 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参数(不确定).您必须自己编写此代码,但对于大型文件可能更有效.

或者,您可以使用上面的代码,然后调用rechunkrechunk到更小的块.这仍然会导致在单个任务中读取每个文件,但后续步骤可以使用较小的块.这是否值得,取决于您单个文件的大小.

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