如何使用dask和xarray加载和处理zarr文件

Question

我在s3中有每月的zarr文件，其中包含网格温度数据。我想为一个经度/经度拉取多个月的数据，并创建该时间序列的数据帧。一些伪代码：

datasets=[]
for file in files:
    s3 = s3fs.S3FileSystem()        
    zarr_store = s3fs.S3Map(file, s3=s3)
    zarr = xr.open_zarr(store=zarr_store, consolidated=True)
    ds = zarr.sel(latitude=lat,
                  longitude=long,
                  time=slice(start_date.strftime("%Y-%m-%d"),
                             end_date.strftime("%Y-%m-%d"))
                       )
    datasets.append(ds)

con = xr.concat(datasets, dim='time')
df = con.to_dataframe()

所以这段代码可以工作，但速度非常慢。我希望用dask来加快速度。我的计划是将方法更改为一次处理一个文件，然后返回一个数据帧。然后，我将调用client.map()并生成所有的dfs，最后将它们连接在一起。所以我最终得到了类似的结果：

def load(file, lat: float, long: float, start_date, end_date):

    s3 = s3fs.S3FileSystem()
    s3_path = file['s3_bucket'] + '/' + file['zarr_s3_key']
    zarr_store = s3fs.S3Map(s3_path, s3=s3)
    zarr = xr.open_zarr(store=zarr_store, consolidated=True)

    ds = zarr.sel(latitude=lat,
                  longitude=long,
                  time=slice(start_date.strftime("%Y-%m-%d"),
                             end_date.strftime("%Y-%m-%d"))
                 )

    tmp = x.result().to_array().values
    df_time = zarr.coords['time'].sel(time=slice(start_date.strftime("%Y-%m-%d"), end_date.strftime("%Y-%m-%d"))).values
    df = pd.DataFrame({'time': df_time, 'lat': lat, 'long': long, 'dat': tmp})
    df.set_index(['time', 'lat', 'long'], inplace=True)

    return df

if __name__ == '__main__':
    client = Client('tcp://xxx')

    start_date = date(2000, 1, 7)
    end_date = date(2000, 10, 20)
    lat = 2
    lon = 10

    # get the s3 locations of the zarr files from the db
    files = get_files()

    # try just running with one file
    res = client.submit(load, files[0], lat, lon, start_date, end_date) 

    # run them all
    future = client.map(load, files,
                        repeat(lat), repeat(lon),
                        repeat(start_date), repeat(end_date))
    x = client.gather(future)

当我仅将客户端连接到本地计算机时，此代码运行良好。但是，当我尝试连接到远程群集时，我在xr.open_zarr调用中得到以下错误：

KeyError：'XXX/data.zarr/.zmetadata‘

我尝试更改代码，并在方法调用之外加载zarr并传递它们，但结果只得到了nans。我是不是漏掉了什么？这不是解决我想要做的事情的正确方法吗？

Answer 1

如果只想在某一时刻提取一个时间序列，可以只创建一个Dask客户机，然后让xarray并行完成这项工作。在下面的示例中，我们只有一个zarr数据集，但只要工作人员忙于处理每个Zarr文件中的块，您就不会从并行解析Zarr文件中获得任何好处。

import xarray as xr
import fsspec
import hvplot.xarray

from dask.distributed import Client

url = 's3://mur-sst/zarr'  # Amazon Public Data

ds = xr.open_zarr(fsspec.get_mapper(url, anon=True), consolidated=True)

timeseries = ds['analysed_sst'].sel(time=slice('2015-01-01','2020-01-01'),
                                    lat=43, 
                                    lon=-70).persist()

timeseries.hvplot()

产生：

​

​

这是Full Jupyter Notebook

Answer 2

能够解决这个问题，但将把它留下来，以防将来对其他人有帮助。

所以这被证明是权限问题。工作人员无法访问S3存储桶，这就是我获取KeyError的原因。

不过，如果这是批量加载/处理zarr的最佳方式，我仍然愿意听取人们的意见。