diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 5ff3ff20b6a..667e75de0b4 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -44,6 +44,9 @@ New Features By `Thomas Hirtz `_. - allow passing a function to ``combine_attrs`` (:pull:`4896`). By `Justus Magin `_. +- The values stored in an dataarray can now be referenced in a call to + py:func:`xarray.DataArray.query` via 'self' (:issue:`5492`, :pull:`5493`). + By `Tom Nicholas `_. Breaking changes ~~~~~~~~~~~~~~~~ @@ -73,6 +76,9 @@ Bug fixes - Fix the ``repr`` of :py:class:`Variable` objects with ``display_expand_data=True`` (:pull:`5406`) By `Justus Magin `_. +- Fixed :py:func:`xarray.DataArray.query` to not fail with an unnamed dataarray + (:issue:`5492`, :pull:`5493`). + By `Tom Nicholas `_. Documentation diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index eab4413d5ce..ed0dcd45500 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -4437,6 +4437,8 @@ def query( dimension(s), where the indexers are given as strings containing Python expressions to be evaluated against the values in the array. + The values stored in dataarrays can also be referenced in queries as 'self'. + Parameters ---------- queries : dict, optional @@ -4488,17 +4490,43 @@ def query( array([3, 4]) Dimensions without coordinates: x + + >>> da = xr.DataArray(np.arange(0, 5, 1), dims="x", name=None) + >>> da + + array([0, 1, 2, 3, 4]) + Dimensions without coordinates: x + >>> da.query(x="self > 2") + + array([3, 4]) + Dimensions without coordinates: x """ - ds = self._to_dataset_whole(shallow_copy=True) - ds = ds.query( + if self.name is None: + # Naming unnamed dataarrays as 'self' allows querying their values still + name = "self" + else: + # For consistency allow named datarrays to be referred to as 'self' also + name = self.name + queries = either_dict_or_kwargs(queries, queries_kwargs, "query") + queries = { + d: (q.replace("self", name) if isinstance(q, str) else q) + for d, q in queries.items() + } + queries_kwargs = {} + + ds = self._to_dataset_whole(name=name).query( queries=queries, parser=parser, engine=engine, missing_dims=missing_dims, **queries_kwargs, ) - return ds[self.name] + + da = ds[name] + if name == "self": + da.name = None + return da def curvefit( self, diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 95b6036712c..a0b7afc8786 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -4659,6 +4659,8 @@ def test_query(self, backend, engine, parser): bb = DataArray(data=b, dims=["x"], name="b") cc = DataArray(data=c, dims=["y"], name="c") dd = DataArray(data=d, dims=["z"], name="d") + nn = DataArray(data=a, dims=["x"]) + nn.name = None elif backend == "dask": import dask.array as da @@ -4667,6 +4669,8 @@ def test_query(self, backend, engine, parser): bb = DataArray(data=da.from_array(b, chunks=3), dims=["x"], name="b") cc = DataArray(data=da.from_array(c, chunks=7), dims=["y"], name="c") dd = DataArray(data=da.from_array(d, chunks=12), dims=["z"], name="d") + nn = DataArray(data=da.from_array(a, chunks=3), dims=["x"]) + nn.name = None # query single dim, single variable actual = aa.query(x="a > 5", engine=engine, parser=parser) @@ -4704,6 +4708,16 @@ def test_query(self, backend, engine, parser): with pytest.raises(UndefinedVariableError): aa.query(x="spam > 50") # name not present + # test with nameless dataarray (GH issue 5492) + actual = nn.query(x="self > 5", engine=engine, parser=parser) + expect = nn.isel(x=(nn > 5)) + assert_identical(expect, actual) + + # test referring to named dataarray as self + actual = aa.query(x="self > 5", engine=engine, parser=parser) + expect = aa.isel(x=(aa > 5)) + assert_identical(expect, actual) + @requires_scipy @pytest.mark.parametrize("use_dask", [True, False]) def test_curvefit(self, use_dask):