Propagate indexes in DataArray binary operations.

Works by propagating indexes in DataArray._replace.

xref pydata#2227. Tests pass!

Author: dcherian

xarray/core/dataarray.py

@@ -386,14 +386,17 @@ def _replace(
         variable: Variable = None,
         coords=None,
         name: Union[Hashable, None, Default] = _default,
+        indexes=None,
     ) -> "DataArray":
         if variable is None:
             variable = self.variable
         if coords is None:
             coords = self._coords
         if name is _default:
             name = self.name
-        return type(self)(variable, coords, name=name, fastpath=True)
+        # if indexes is None:
+        #     indexes = self.indexes
+        return type(self)(variable, coords, name=name, fastpath=True, indexes=indexes)
 
     def _replace_maybe_drop_dims(
         self, variable: Variable, name: Union[Hashable, None, Default] = _default
@@ -440,7 +443,8 @@ def _from_temp_dataset(
     ) -> "DataArray":
         variable = dataset._variables.pop(_THIS_ARRAY)
         coords = dataset._variables
-        return self._replace(variable, coords, name)
+        indexes = dataset._indexes
+        return self._replace(variable, coords, name, indexes=indexes)
 
     def _to_dataset_split(self, dim: Hashable) -> Dataset:
         def subset(dim, label):
@@ -2506,7 +2510,7 @@ def func(self, other):
             coords, indexes = self.coords._merge_raw(other_coords)
             name = self._result_name(other)
 
-            return self._replace(variable, coords, name)
+            return self._replace(variable, coords, name, indexes=indexes)
 
         return func
 

xarray/core/groupby.py

@@ -529,6 +529,7 @@ def _maybe_unstack(self, obj):
             for dim in self._inserted_dims:
                 if dim in obj.coords:
                     del obj.coords[dim]
+                    del obj.indexes[dim]
         return obj
 
     def fillna(self, value):

xarray/core/indexes.py

@@ -35,6 +35,9 @@ def __contains__(self, key):
     def __getitem__(self, key):
         return self._indexes[key]
 
+    def __delitem__(self, key):
+        del self._indexes[key]
+
     def __repr__(self):
         return formatting.indexes_repr(self)
 

xarray/tests/test_dataarray.py

@@ -3953,6 +3953,17 @@ def test_matmul(self):
         expected = da.dot(da)
         assert_identical(result, expected)
 
+    def test_binary_op_propagate_indexes(self):
+        # regression test for GH2227
+        self.dv["x"] = np.arange(self.dv.sizes["x"])
+        expected = self.dv.indexes["x"]
+
+        actual = (self.dv * 10).indexes["x"]
+        assert expected is actual
+
+        actual = (self.dv > 10).indexes["x"]
+        assert expected is actual
+
     def test_binary_op_join_setting(self):
         dim = "x"
         align_type = "outer"

xarray/tests/test_dataset.py

@@ -4951,6 +4951,14 @@ def test_filter_by_attrs(self):
         )
         assert not bool(new_ds.data_vars)
 
+    def test_binary_op_propagate_indexes(self):
+        ds = Dataset(
+            {"d1": DataArray([1, 2, 3], dims=["x"], coords={"x": [10, 20, 30]})}
+        )
+        expected = ds.indexes["x"]
+        actual = (ds * 2).indexes["x"]
+        assert expected is actual
+
     def test_binary_op_join_setting(self):
         # arithmetic_join applies to data array coordinates
         missing_2 = xr.Dataset({"x": [0, 1]})