Python中如何使用外部变量查询xarray Dataset？

问题重现 #

你尝试在xarray的query方法中使用外部变量替代硬编码值，但直接引用变量会触发错误：

正常运行的硬编码查询示例 #

import numpy as np
import xarray as xr

a = np.arange(0, 5, 1)
b = np.linspace(0, 1, 5)
ds = xr.Dataset({"a": ("x", a), "b": ("x", b)})
# 硬编码数值查询正常执行
ds.query(x="a > 2")

外部变量引用失败案例 #

z = 2
# 直接引用变量报错
ds.query(x="a > z")

UndefinedVariableError: name 'z' is not defined

尝试沿用Pandas的@前缀语法同样失效：

ds.query(x="a > @z")

SyntaxError: The '@' prefix is not allowed in top-level eval calls.
please refer to your variables by name without the '@' prefix.

可行解决方案 #

1. 字符串格式化嵌入变量 #

用Python的f-string或str.format()将外部变量的值直接嵌入查询字符串：

z = 2
ds.query(x=f"a > {z}")

该方式会将查询字符串解析为"a > 2"，和硬编码效果完全一致，适合简单场景。

2. 利用query的local_dict参数 #

xarray的query方法支持通过local_dict参数传递外部变量字典，直接在查询表达式中使用变量名：

z = 2
ds.query(x="a > z", local_dict={"z": z})

如果需要传递多个外部变量，只需在字典中添加对应键值对即可，可读性更强。

3. 用where方法替代query #

若无需字符串表达式的查询语法，推荐使用xarray原生的where方法直接进行布尔索引，更贴合xarray的使用习惯：

z = 2
ds.where(ds.a > z, drop=True)

drop=True会移除不满足条件的维度条目，行为和query一致。

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