Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@ Bug fixes
By `Stan West <https://github.com/stanwest>`_.
- Fix bug in :py:func:`where` when passing non-xarray objects with ``keep_attrs=True``. (:issue:`6444`, :pull:`6461`)
By `Sam Levang <https://github.com/slevang>`_.
- Allow passing both ``other`` and ``drop=True`` arguments to ``xr.DataArray.where``
and ``xr.Dataset.where`` (:pull:`6466`, :pull:`6467`).
By `Michael Delgado <https://github.com/delgadom>`_.

Documentation
~~~~~~~~~~~~~
Expand Down
14 changes: 9 additions & 5 deletions xarray/core/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -1197,8 +1197,7 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
By default, these locations filled with NA.
drop : bool, optional
If True, coordinate labels that only correspond to False values of
the condition are dropped from the result. Mutually exclusive with
``other``.
the condition are dropped from the result.

Returns
-------
Expand Down Expand Up @@ -1251,6 +1250,14 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
[15., nan, nan, nan]])
Dimensions without coordinates: x, y

>>> a.where(a.x + a.y < 4, -1, drop=True)
<xarray.DataArray (x: 4, y: 4)>
array([[ 0, 1, 2, 3],
[ 5, 6, 7, -1],
[10, 11, -1, -1],
[15, -1, -1, -1]])
Dimensions without coordinates: x, y

See Also
--------
numpy.where : corresponding numpy function
Expand All @@ -1264,9 +1271,6 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
cond = cond(self)

if drop:
if other is not dtypes.NA:
raise ValueError("cannot set `other` if drop=True")

if not isinstance(cond, (Dataset, DataArray)):
raise TypeError(
f"cond argument is {cond!r} but must be a {Dataset!r} or {DataArray!r}"
Expand Down
7 changes: 5 additions & 2 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4548,8 +4548,11 @@ def test_where_other(self):
actual = ds.where(lambda x: x > 1, -1)
assert_equal(expected, actual)

with pytest.raises(ValueError, match=r"cannot set"):
ds.where(ds > 1, other=0, drop=True)
actual = ds.where(ds > 1, other=-1, drop=True)
expected_nodrop = ds.where(ds > 1, -1)
_, expected = xr.align(actual, expected_nodrop, join="left")
assert_equal(actual, expected)
assert actual.a.dtype == int

with pytest.raises(ValueError, match=r"cannot align .* are not equal"):
ds.where(ds > 1, ds.isel(x=slice(3)))
Expand Down