Add data_vars_only rechunking param; update test

- Add a data_vars_only parameter to chunk_dataset and
  update_dataset_chunk_encoding.

- Update test_append_time_slice to use this parameter in order to
  ensure compatibility with xarray 2024.3.0

Addresses Issue #958.

Author: pont-us

test/core/test_timeslice.py

@@ -38,15 +38,18 @@ def write_cube(self, start_date, num_days: int):
         cube = self.make_cube(start_date, num_days)
         cube.to_zarr(self.CUBE_PATH)
 
-    def make_cube(self, start_date, num_days: int) -> xr.Dataset:
+    @staticmethod
+    def make_cube(start_date, num_days: int) -> xr.Dataset:
         cube = new_cube(
             time_periods=num_days,
             time_freq="1D",
             time_start=start_date,
             variables=dict(precipitation=0.1, temperature=270.5, soil_moisture=0.2),
         )
         chunk_sizes = dict(time=1, lat=90, lon=90)
-        cube = chunk_dataset(cube, chunk_sizes, format_name="zarr")
+        cube = chunk_dataset(
+            cube, chunk_sizes, format_name="zarr", data_vars_only=True
+        )
         return cube
 
     def test_find_time_slice(self):

xcube/core/chunk.py

@@ -8,7 +8,10 @@
 
 
 def chunk_dataset(
-    dataset: xr.Dataset, chunk_sizes: Dict[str, int] = None, format_name: str = None
+    dataset: xr.Dataset,
+    chunk_sizes: Dict[str, int] = None,
+    format_name: str = None,
+    data_vars_only: bool = False,
 ) -> xr.Dataset:
     """Chunk *dataset* using *chunk_sizes* and optionally
     update encodings for given *format_name*.
@@ -17,14 +20,23 @@ def chunk_dataset(
         dataset: input dataset
         chunk_sizes: mapping from dimension name to new chunk size
         format_name: optional format, e.g. "zarr" or "netcdf4"
+        data_vars_only: only chunk data variables, not coordinates
 
     Returns:
         the (re)chunked dataset
     """
-    dataset = dataset.chunk(chunks=chunk_sizes)
+
+    if data_vars_only:
+        for variable in dataset.data_vars:
+            dataset[variable] = dataset[variable].chunk(chunk_sizes)
+    else:
+        dataset = dataset.chunk(chunks=chunk_sizes)
     if format_name:
         dataset = update_dataset_chunk_encoding(
-            dataset, chunk_sizes=chunk_sizes, format_name=format_name
+            dataset,
+            chunk_sizes=chunk_sizes,
+            format_name=format_name,
+            data_vars_only=data_vars_only,
         )
     return dataset
 

xcube/core/update.py

@@ -234,6 +234,7 @@ def update_dataset_chunk_encoding(
     chunk_sizes: Dict[str, int] = None,
     format_name: str = None,
     in_place: bool = False,
+    data_vars_only: bool = False
 ) -> xr.Dataset:
     """Update each variable's encoding in *dataset* with respect to *chunk_sizes*
     so *dataset* is written in chunks for given *format_name*.
@@ -245,6 +246,7 @@ def update_dataset_chunk_encoding(
         format_name: format name, e.g. "zarr" or "netcdf4".
         in_place: If ``True``, *dataset* will be modified in place and
             returned.
+        data_vars_only: only chunk data variables, not coordinates
     """
     if format_name == FORMAT_NAME_ZARR:
         chunk_sizes_attr_name = "chunks"
@@ -254,7 +256,7 @@ def update_dataset_chunk_encoding(
         return dataset
     if not in_place:
         dataset = dataset.copy()
-    for var_name in dataset.variables:
+    for var_name in dataset.data_vars if data_vars_only else dataset.variables:
         var = dataset[var_name]
         if chunk_sizes is not None: