使用xarray.reindex补全时间序列后写入NetCDF内存分配失败

尝试将覆盖美国本土的单日网格化雷达降雨GRIB2文件转为NetCDF格式。当GRIB2存在时间步缺失时，先用xarray.Dataset.reindex填充NA值，再执行xarray.Dataset.to_netcdf会触发内存分配错误；跳过reindex则转换成功。数据集块设置为(70, 3500, 7000)，但调用to_netcdf时尝试加载(210, 3500, 7000)的块。

可复现资源 #

可获取对应的测试代码与数据。

代码 #

#%% 导入库
import time
start_time = time.time()
import xarray as xr
import cfgrib
from glob import glob
import pandas as pd
import dask
dask.config.set(**{'array.slicing.split_large_chunks': False}) # 关闭大切片加载内存警告
dask.config.set(scheduler='synchronous') # 强制单线程计算（NetCDF只能串行写入）
#%% 参数设置
chnk_sz = "7000MB"
fl_out_nc = "out_netcdfs/20010101.nc"
fldr_in_grib = "in_gribs/20010101.grib2"

#%% 加载并导出数据集
ds = xr.open_dataset(fldr_in_grib, engine="cfgrib", chunks={"time":chnk_sz},
                    backend_kwargs={'indexpath': ''})

# 重索引补全时间步
start_date = pd.to_datetime('2001-01-01')
tstep = pd.Timedelta('0 days 00:05:00')
new_index = pd.date_range(start=start_date, end=start_date + pd.Timedelta(1, "day"),\
                                    freq=tstep, inclusive='left')

ds = ds.reindex(indexers={"time":new_index})
ds = ds.unify_chunks()
ds = ds.chunk(chunks={'time':chnk_sz})

print("######## 写入NetCDF前的数据集信息 ########")
print(ds)
print(' ')
print("######## 错误信息 ########")
ds.to_netcdf(fl_out_nc, encoding= {"unknown":{"zlib":True}})

输出信息 #

######## 写入NetCDF前的数据集信息 ########
<xarray.Dataset>
Dimensions:     (time: 288, latitude: 3500, longitude: 7000)
Coordinates:
  * time        (time) datetime64[ns] 2001-01-01 ... 2001-01-01T23:55:00
  * latitude    (latitude) float64 54.99 54.98 54.98 54.97 ... 20.03 20.02 20.01
  * longitude   (longitude) float64 230.0 230.0 230.0 ... 300.0 300.0 300.0
    step        timedelta64[ns] ...
    surface     float64 ...
    valid_time  (time) datetime64[ns] dask.array<chunksize=(288,), meta=np.ndarray>
Data variables:
    unknown     (time, latitude, longitude) float32 dask.array<chunksize=(70, 3500, 7000), meta=np.ndarray>
Attributes:
    GRIB_edition:            2
    GRIB_centre:             161
    GRIB_centreDescription:  161
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             161
    history:                 2022-09-10T14:50 GRIB to CDM+CF via cfgrib-0.9.1...
 
######## 错误信息 ########
输出超出大小限制，请在文本编辑器中查看完整输出
---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
d:\Dropbox\_Sharing\reprex\2022-9-9_writing_ncdf_fails\reprex\exporting_netcdfs_reduced.py in <cell line: 22>()
     160 print(' ')
     161 print("######## 错误信息 ########")
---> 162 ds.to_netcdf(fl_out_nc, encoding= {"unknown":{"zlib":True}})

File c:\Users\Daniel\anaconda3\envs\weather_gen_3\lib\site-packages\xarray\core\dataset.py:1882, in Dataset.to_netcdf(self, path, mode, format, group, engine, encoding, unlimited_dims, compute, invalid_netcdf)
   1879     encoding = {}
   1880 from ..backends.api import to_netcdf
-> 1882 return to_netcdf(  # type: ignore  # mypy cannot resolve the overloads:(
   1883     self,
   1884     path,
   1885     mode=mode,
   1886     format=format,
   1887     group=group,
   1888     engine=engine,
   1889     encoding=encoding,
   1890     unlimited_dims=unlimited_dims,
   1891     compute=compute,
   1892     multifile=False,
   1893     invalid_netcdf=invalid_netcdf,
   1894 )

File c:\Users\xxxxx\anaconda3\envs\weather_gen_3\lib\site-packages\xarray\backends\api.py:1219, in to_netcdf(dataset, path_or_file, mode, format, group, engine, encoding, unlimited_dims, compute, multifile, invalid_netcdf)
...
    121     return arg

File <__array_function__ internals>:180, in where(*args, **kwargs)

MemoryError: 无法分配19.2 GiB内存用于形状为(210, 3500, 7000)、数据类型为float32的数组

环境配置 #

Windows 11 Home
xarray 2022.3.0
cfgrib 0.9.10.1
dask 2022.7.0

内容的提问来源于stack exchange，提问作者Daniel Lassiter