Xarray布尔数组需compute()：open_mfdataset用where报KeyError求助

问题原因 #

• open_mfdataset默认采用Dask延迟加载数据，所有操作均为惰性执行，ds["d2fd"].max()返回的是未计算的Dask数组
• 新版本Xarray（2023.7.0）收紧了索引规则：使用形状未知的Dask布尔数组做索引会触发错误，因为无法确定结果数据集的结构
• open_dataset直接将数据加载到内存，所有操作立即计算，因此不会出现该问题，旧版本Xarray未限制此类操作

解决方法 #

方法1：先计算最大值标量再索引

仅计算最大值并转为本地标量，适合大数据场景，内存占用低：

import xarray as xr
ds = xr.open_mfdataset('my_file.nc')

# 计算最大值并转为内存中的标量
max_d2fd = ds["d2fd"].max().compute()
# 使用已计算的标量执行where操作
result = ds.where(ds["d2fd"] == max_d2fd, drop=True)

方法2：加载全部数据到内存

若数据集规模较小，直接将整个数据集加载到内存，后续操作与open_dataset一致：

import xarray as xr
ds = xr.open_mfdataset('my_file.nc').load()
# 此时max()会立即计算为内存中的标量，可直接执行where操作
result = ds.where(ds["d2fd"] == ds["d2fd"].max(), drop=True)

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