diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 6d56c410cc4..8fbce607ed2 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -170,6 +170,9 @@ Deprecations Bug Fixes ~~~~~~~~~ +- :py:meth:`Dataset.query` now resolves coordinates in datasets with no data + variables (:issue:`6449`). + By `Liang, Jingyi `_. - Results of the string accessor methods ``strip``, ``pad``, ``center``, ``ljust``, ``rjust``, ``zfill`` and ``replace`` on dask arrays are no longer truncated to a single character, and ``repeat`` no longer fails on dask arrays diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index ccec59b403d..3e2cfbea086 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -10025,9 +10025,12 @@ def query( msg = f"expr for dim {dim} must be a string to be evaluated, {type(expr)} given" raise ValueError(msg) - # evaluate the queries to create the indexers + # Include coords so pandas considers the resolvers nonempty even when + # there are no data variables. indexers = { - dim: pd.eval(expr, resolvers=[self], parser=parser, engine=engine) + dim: pd.eval( + expr, resolvers=[self, self.coords], parser=parser, engine=engine + ) for dim, expr in queries.items() } diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 21a26e1e34e..16a641653e0 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -37,7 +37,7 @@ from xarray.core.common import duck_array_ops, full_like from xarray.core.coordinates import Coordinates, DatasetCoordinates from xarray.core.indexes import Index, PandasIndex -from xarray.core.types import ArrayLike +from xarray.core.types import ArrayLike, QueryEngineOptions, QueryParserOptions from xarray.core.utils import is_scalar from xarray.groupers import SeasonResampler, TimeResampler from xarray.namedarray.pycompat import array_type, integer_types @@ -7933,11 +7933,18 @@ def test_astype_attrs(self) -> None: assert not data.astype(float, keep_attrs=False).attrs assert not data.astype(float, keep_attrs=False).var1.attrs + @pytest.mark.parametrize("coords_only", [False, True]) @pytest.mark.parametrize("parser", ["pandas", "python"]) @pytest.mark.parametrize( "engine", ["python", None, pytest.param("numexpr", marks=[requires_numexpr])] ) - def test_query(self, use_dask: bool, engine, parser) -> None: + def test_query( + self, + use_dask: bool, + engine: QueryEngineOptions, + parser: QueryParserOptions, + coords_only: bool, + ) -> None: """Test querying a dataset.""" # setup test data np.random.seed(42) @@ -7992,6 +7999,9 @@ def test_query(self, use_dask: bool, engine, parser) -> None: }, ) + if coords_only: + ds = ds.set_coords(list(ds.data_vars)) + # query single dim, single variable with raise_if_dask_computes(): actual = ds.query(x="a2 > 5", engine=engine, parser=parser) @@ -8084,6 +8094,17 @@ def test_query(self, use_dask: bool, engine, parser) -> None: expect = ds.isel(dict(x=(a > 5), y=(c < 0.5), z=(d == "bar"))) assert_identical(expect, actual) + # query a dimension without a coordinate + with raise_if_dask_computes(): + actual = ds.query(x="x > 5", engine=engine, parser=parser) + assert_identical(ds.isel(x=slice(6, None)), actual) + + # query a dimension coordinate + ds = ds.assign_coords(x=a + 10) + with raise_if_dask_computes(): + actual = ds.query(x="x > 15", engine=engine, parser=parser) + assert_identical(ds.isel(x=slice(6, None)), actual) + # test error handling with pytest.raises(ValueError): ds.query("a > 5") # type: ignore[arg-type] # must be dict or kwargs