如何使用索引从XArray DataArray提取像素值到DataFrame列

问题背景 #

我目前在处理一项非常规栅格计算任务：手里有多幅基于NLCD的90亿像素栅格地图，需要提取其中约5亿个历史建成区像素对应的所有栅格取值，最初用于提取建成区索引的代码如下：

built_up_index = pandas.DataFrame(np.column_stack(np.where(unbuilt == 0)), columns = ["row", "column"]).sort_values(["row", "column"])

这段代码会生成一个两列的DataFrame，分别存储所有在任意一期NLCD栅格中被标记为建成区的像素的行索引、列索引——其中unbuilt是存储建成区标识的二值栅格，值为0时对应该位置是建成区。

目标是基于这份索引表，读取所有NLCD年度地图及其他关联栅格的取值，组装成每行对应一个像素、每列对应一个变量的结构化表，列结构依次为像素行号、列号、NLCD2001取值、NLCD2004取值、其他自定义计算指数等，目标结构示例：

初始方案的问题 #

一开始尝试用xarray的isel方法直接传入索引数组取值，代码如下：

test = sprawl_2001.isel({'y': np.array(built_up_frame.iloc[:,0]), 'x': np.array(built_up_frame.iloc[:,1])}, drop = True).to_dataset(name="var").to_dataframe()

取前10000条样本做子集测试时代码可以正常运行：

test = sprawl_2001.isel({'y': np.array(built_up_frame.iloc[0:10000,0]), 'x': np.array(built_up_frame.iloc[0:10000,1])}, drop = True).to_dataset(name="var").to_dataframe()

但返回结果完全不符合预期：返回的DataFrame长度是传入索引长度的平方。xarray的isel传入多维度索引数组时会默认对不同维度的索引做笛卡尔积，生成二维数组后再展平，而实际需要的是索引对一一对应的像素值，也就是和传入索引长度一致的一维取值向量。

逐像素循环遍历的方案虽然能实现逻辑，但面对5亿量级的数据运行效率极低，显然存在更高效的向量化实现路径。
后续排查确认直接传索引的方案走不通的核心原因：xarray的isel按上述方式调用时，会生成和原始数据集维度（约161000列、104000行）一致的、带大量缺失值的数组，根本无法直接输出需要的一维值向量。

可落地实现方案 #

最终改用np.extract方法实现需求，代码如下：

def src_to_frame(src, unbuilt, varname):
    return pd.DataFrame(np.extract(unbuilt == 0, src), columns=[varname])

函数参数说明：

• src：存储目标变量的栅格数组
• unbuilt：和src同尺寸的建成区标识栅格，值为0对应历史建成区像素
• varname：输出列的变量名

这个方案可以直接生成目标结构的结构化表，内存占用在常规工作站的RAM承载范围内，虽然不一定是理论最优解，但可以稳定运行处理全量数据。

可选优化方向：

• 提前把所有待提取的栅格堆叠为三维数组（维度为[变量数, 行, 列]），一次性做布尔索引提取后再转成DataFrame，减少多次数组遍历的开销
• 如果单设备内存不足，可以按行块分块读取栅格，逐块提取建成区像素值后追加写入磁盘存储的列式格式（比如parquet），避免全量数据加载进内存

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