Bugfix: Set equal chunking for shapes and dataset

rsciio/_hierarchical.py

@@ -263,12 +263,12 @@ def _read_array(group, dataset_key):
             key = "ragged_shapes"
         if key in group:
             ragged_shape = group[key]
-            # if the data is chunked saved array we must first
-            # cast to a numpy array to avoid multiple calls to
-            # _decode_chunk in zarr (or h5py)
+            # Use same chunks as data so that apply_gufunc doesn't rechunk
+            # Reduces the transfer of data between workers which
+            # significantly improves performance for distributed loading
             data = da.from_array(data, chunks=data.chunks)
-            shape = da.from_array(ragged_shape, chunks=ragged_shape.chunks)
-            shape = shape.rechunk(data.chunks)
+            shape = da.from_array(ragged_shape, chunks=data.chunks)
+
             data = da.apply_gufunc(unflatten_data, "(),()->()", data, shape)
         return data
 
@@ -764,15 +764,15 @@ def overwrite_dataset(
                     shapes[i] = np.array(data[i].shape)
 
             shape_dset = cls._get_object_dset(
-                group, shapes, f"_ragged_shapes_{key}", shapes.shape, **kwds
+                group, shapes, f"_ragged_shapes_{key}", chunks, **kwds
             )
 
             cls._store_data(
                 (new_data, shapes),
                 (dset, shape_dset),
                 group,
                 (key, f"_ragged_shapes_{key}"),
-                (chunks, shapes.shape),
+                (chunks, chunks),
                 show_progressbar,
             )
         else:

upcoming_changes/211.bugfix.rst

@@ -0,0 +1 @@
+Fix saving ragged arrays of vectors from/to a chunked ``hspy`` and ``zspy`` store.  Greatly increases the speed of saving and loading ragged arrays from chunked datasets.