Handle extra indexes for zarr region writes

xarray/backends/api.py

@@ -1562,7 +1562,9 @@ def _auto_detect_regions(ds, region, open_kwargs):
     return region
 
 
-def _validate_and_autodetect_region(ds, region, mode, open_kwargs) -> dict[str, slice]:
+def _validate_and_autodetect_region(
+    ds, region, mode, open_kwargs
+) -> tuple[dict[str, slice], Dataset]:
     if region == "auto":
         region = {dim: "auto" for dim in ds.dims}
 
@@ -1592,7 +1594,10 @@ def _validate_and_autodetect_region(ds, region, mode, open_kwargs) -> dict[str,
                 "step on all slices in ``region`` must be 1 or None, got "
                 f"region={region}"
             )
+    # Need to initially keep indexes for auto-detecting regions, then safe to drop
+    ds = ds.drop_vars(ds.indexes)
 
+    # Raise for any non-index variables not corresponding to the region write
     non_matching_vars = [
         k for k, v in ds.variables.items() if not set(region).intersection(v.dims)
     ]
@@ -1607,7 +1612,7 @@ def _validate_and_autodetect_region(ds, region, mode, open_kwargs) -> dict[str,
             f".drop_vars({non_matching_vars!r})"
         )
 
-    return region
+    return region, ds
 
 
 def _validate_datatypes_for_zarr_append(zstore, dataset):
@@ -1779,9 +1784,10 @@ def to_zarr(
             storage_options=storage_options,
             zarr_version=zarr_version,
         )
-        region = _validate_and_autodetect_region(dataset, region, mode, open_kwargs)
-        # can't modify indexed with region writes
-        dataset = dataset.drop_vars(dataset.indexes)
+        region, dataset = _validate_and_autodetect_region(
+            dataset, region, mode, open_kwargs
+        )
+
         if append_dim is not None and append_dim in region:
             raise ValueError(
                 f"cannot list the same dimension in both ``append_dim`` and "

xarray/tests/test_backends.py

@@ -2818,11 +2818,15 @@ def test_write_preexisting_override_metadata(self) -> None:
     def test_write_region_errors(self) -> None:
         data = Dataset({"u": (("x",), np.arange(5))})
         data2 = Dataset({"u": (("x",), np.array([10, 11]))})
+        data3 = Dataset(
+            {"u": (("x",), np.array([10, 11])), "v": (("y",), np.arange(3))},
+            coords={"y": ("y", [1, 2, 3])},
+        )
 
         @contextlib.contextmanager
         def setup_and_verify_store(expected=data):
             with self.create_zarr_target() as store:
-                data.to_zarr(store, **self.version_kwargs)
+                expected.to_zarr(store, **self.version_kwargs)
                 yield store
                 with self.open(store) as actual:
                     assert_identical(actual, expected)
@@ -2891,6 +2895,20 @@ def setup_and_verify_store(expected=data):
             ):
                 data2.to_zarr(store, region={"x": slice(3)}, **self.version_kwargs)
 
+        # region writes should ignore indexes that don't overlap the region
+        expected3 = Dataset(
+            {"u": (("x",), np.array([10, 11, 2, 3, 4]))}, coords={"y": ("y", [1, 2, 3])}
+        )
+        with setup_and_verify_store(expected3) as store:
+            data3.drop_vars("v").to_zarr(
+                store, region={"x": slice(2)}, **self.version_kwargs
+            )
+
+        # but fail on data vars that don't overlap the region
+        with setup_and_verify_store() as store:
+            with pytest.raises(ValueError, match=r"when setting `region` explicitly"):
+                data3.to_zarr(store, region={"x": slice(2)}, **self.version_kwargs)
+
     @requires_dask
     def test_encoding_chunksizes(self) -> None:
         # regression test for GH2278