当然可以实现这个需求！你想要的是基于a数组里的有限值（非NaN）同步过滤整个Dataset中的a和b，把二维数据压缩成一维，同时保证对应位置的数据匹配。下面给你几种更贴合xarray风格的实现方式，比手动处理每个数组更高效：

方法1：使用where + stack + dropna #

这种方法利用xarray的广播和维度操作特性，一步完成掩码应用和过滤：

import xarray as xr
import numpy as np

n = np.nan
a = np.array([[n,n,1], [1,2,2], [2,n,n]], dtype='float32')
b = np.random.rand(3,3)
xr_dataset = xr.Dataset({'a': xr.DataArray(a, dims=['x', 'y']), 'b': xr.DataArray(b, dims=['x', 'y'])})

# 基于a的有限值创建掩码，对整个Dataset应用掩码
mask = np.isfinite(xr_dataset['a'])
# 合并x、y维度为point，然后删除含NaN的条目
xr_flat = xr_dataset.where(mask).stack(point=('x', 'y')).dropna(dim='point')
# 可选：移除point索引，得到纯一维数组结构
xr_flat = xr_flat.reset_index('point').drop_vars('point')

步骤解释： #

• where(mask)：将Dataset中所有不满足a为有限值的位置设为NaN
• stack(point=('x', 'y'))：把二维的x、y维度合并成一个名为point的一维索引
• dropna(dim='point')：删除所有包含NaN的point条目，留下有效数据
• 最后重置并删除point索引，就得到和你示例中完全一致的xr_flat

方法2：直接使用布尔索引 #

如果你习惯numpy的索引方式，也可以直接用拉平后的掩码过滤整个Dataset：

# 生成拉平后的布尔掩码
flat_mask = np.isfinite(xr_dataset['a']).values.ravel()
# 合并维度后用掩码过滤，再清理索引
xr_flat = xr_dataset.stack(point=('x', 'y'))[flat_mask].reset_index('point').drop_vars('point')

验证结果 #

你可以打印xr_flat['a'].values和手动处理的a_flat对比，会发现两者完全一致，同时b的数据也同步对应上了。

内容的提问来源于stack exchange，提问作者Loïc Dutrieux