Fix Dataset.where with drop=True and mixed dims

doc/whats-new.rst

@@ -34,6 +34,9 @@ Bug fixes
   allowing the ``encoding`` and ``unlimited_dims`` options with ``save_mfdataset``.
   (:issue:`6684`)
   By `Travis A. O'Brien <https://github.com/taobrienlbl>`_.
+- :py:meth:`Dataset.where` with ``drop=True`` now behaves correctly with mixed dimensions.
+  (:issue:`6227`, :pull:`6690`)
+  By `Michael Niklas <https://github.com/headtr1ck>`_.
 
 Documentation
 ~~~~~~~~~~~~~

xarray/core/common.py

@@ -1362,18 +1362,23 @@ def where(
                     f"cond argument is {cond!r} but must be a {Dataset!r} or {DataArray!r}"
                 )
 
-            # align so we can use integer indexing
             self, cond = align(self, cond)  # type: ignore[assignment]
 
-            # get cond with the minimal size needed for the Dataset
-            if isinstance(cond, Dataset):
-                clipcond = cond.to_array().any("variable")
-            else:
-                clipcond = cond
+            def _dataarray_indexer(dim: Hashable) -> DataArray:
+                return cond.any(dim=(d for d in cond.dims if d != dim))
+
+            def _dataset_indexer(dim: Hashable) -> DataArray:
+                cond_wdim = cond.drop(var for var in cond if dim not in cond[var].dims)
+                keepany = cond_wdim.any(dim=(d for d in cond.dims.keys() if d != dim))
+                return keepany.to_array().any("variable")
+
+            _get_indexer = (
+                _dataarray_indexer if isinstance(cond, DataArray) else _dataset_indexer
+            )
 
-            # clip the data corresponding to coordinate dims that are not used
-            nonzeros = zip(clipcond.dims, np.nonzero(clipcond.values))
-            indexers = {k: np.unique(v) for k, v in nonzeros}
+            indexers = {}
+            for dim in cond.sizes.keys():
+                indexers[dim] = _get_indexer(dim)
 
             self = self.isel(**indexers)
             cond = cond.isel(**indexers)

xarray/tests/test_dataset.py

@@ -4715,6 +4715,25 @@ def test_where_drop(self) -> None:
         actual8 = ds.where(ds > 0, drop=True)
         assert_identical(expected8, actual8)
 
+        # mixed dimensions: PR#6690, Issue#6227
+        ds = xr.Dataset(
+            {
+                "a": ("x", [1, 2, 3]),
+                "b": ("y", [2, 3, 4]),
+                "c": (("x", "y"), np.arange(9).reshape((3, 3))),
+            }
+        )
+        expected9 = xr.Dataset(
+            {
+                "a": ("x", [np.nan, 3]),
+                "b": ("y", [np.nan, 3, 4]),
+                "c": (("x", "y"), np.arange(3.0, 9.0).reshape((2, 3))),
+            }
+        )
+        actual9 = ds.where(ds > 2, drop=True)
+        assert actual9.sizes["x"] == 2
+        assert_identical(expected9, actual9)
+
     def test_where_drop_empty(self) -> None:
         # regression test for GH1341
         array = DataArray(np.random.rand(100, 10), dims=["nCells", "nVertLevels"])