Rechunking error when saving a queried raster element

[LucaMarconato]

As reported here, scverse/spatialdata-io#117 querying a raster element and trying to save it may lead to the following error:

ValueError: Attempt to save array to zarr with irregular chunking, please call `arr.rechunk(...)` first.

A solution is to rechunk the data, I will make a fix for this.

[josenimo]

Hey @LucaMarconato

I ran into this error just now, trying to save a spatial query.
I chunked the image before parsing at the dask image object type.
any clues on how to deal with this?

versions:

spatialdata: 0.2.5.post0
geopandas: 1.0.1
anndata: 0.11.1
napari_spatialdata: 0.5.3

Details

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[26], line 2
      1 datetime = time.strftime("%Y%m%d_%H%M")
----> 2 sdata_subset.write(f"../data/sdata/{datetime}_sdata_subset.zarr")

File /opt/homebrew/Caskroom/mambaforge/base/envs/spatialdata/lib/python3.10/site-packages/spatialdata/_core/spatialdata.py:1159, in SpatialData.write(self, file_path, overwrite, consolidate_metadata, format)
   1156 store.close()
   1158 for element_type, element_name, element in self.gen_elements():
-> 1159     self._write_element(
   1160         element=element,
   1161         zarr_container_path=file_path,
   1162         element_type=element_type,
   1163         element_name=element_name,
   1164         overwrite=False,
   1165         format=format,
   1166     )
   1168 if self.path != file_path:
   1169     old_path = self.path

File /opt/homebrew/Caskroom/mambaforge/base/envs/spatialdata/lib/python3.10/site-packages/spatialdata/_core/spatialdata.py:1203, in SpatialData._write_element(self, element, zarr_container_path, element_type, element_name, overwrite, format)
   1200 parsed = _parse_formats(formats=format)
   1202 if element_type == "images":
-> 1203     write_image(image=element, group=element_type_group, name=element_name, format=parsed["raster"])
   1204 elif element_type == "labels":
   1205     write_labels(labels=element, group=root_group, name=element_name, format=parsed["raster"])

File /opt/homebrew/Caskroom/mambaforge/base/envs/spatialdata/lib/python3.10/site-packages/spatialdata/_io/io_raster.py:239, in write_image(image, group, name, format, storage_options, **metadata)
    231 def write_image(
    232     image: Union[DataArray, DataTree],
    233     group: zarr.Group,
   (...)
    237     **metadata: Union[str, JSONDict, list[JSONDict]],
    238 ) -> None:
--> 239     _write_raster(
    240         raster_type="image",
    241         raster_data=image,
    242         group=group,
    243         name=name,
    244         format=format,
    245         storage_options=storage_options,
    246         **metadata,
    247     )

File /opt/homebrew/Caskroom/mambaforge/base/envs/spatialdata/lib/python3.10/site-packages/spatialdata/_io/io_raster.py:170, in _write_raster(raster_type, raster_data, group, name, format, storage_options, label_metadata, **metadata)
    166 # Scaler needs to be None since we are passing the data already downscaled for the multiscale case.
    167 # We need this because the argument of write_image_ngff is called image while the argument of
    168 # write_labels_ngff is called label.
    169 metadata[raster_type] = data
--> 170 write_single_scale_ngff(
    171     group=group_data,
    172     scaler=None,
    173     fmt=format,
    174     axes=parsed_axes,
    175     coordinate_transformations=None,
    176     storage_options=storage_options,
    177     **metadata,
    178 )
    179 assert transformations is not None
    180 overwrite_coordinate_transformations_raster(
    181     group=_get_group_for_writing_transformations(), transformations=transformations, axes=input_axes
    182 )

File /opt/homebrew/Caskroom/mambaforge/base/envs/spatialdata/lib/python3.10/site-packages/ome_zarr/writer.py:501, in write_image(image, group, scaler, chunks, fmt, axes, coordinate_transformations, storage_options, compute, **metadata)
    498 dask_delayed_jobs = []
    500 if isinstance(image, da.Array):
--> 501     dask_delayed_jobs = _write_dask_image(
    502         image,
    503         group,
    504         scaler,
    505         chunks=chunks,
    506         fmt=fmt,
    507         axes=axes,
    508         coordinate_transformations=coordinate_transformations,
    509         storage_options=storage_options,
    510         name=None,
    511         compute=compute,
    512         **metadata,
    513     )
    514 else:
    515     mip, axes = _create_mip(image, fmt, scaler, axes)

File /opt/homebrew/Caskroom/mambaforge/base/envs/spatialdata/lib/python3.10/site-packages/ome_zarr/writer.py:604, in _write_dask_image(image, group, scaler, chunks, fmt, axes, coordinate_transformations, storage_options, name, compute, **metadata)
    598     shapes.append(image.shape)
    600     LOGGER.debug(
    601         "write dask.array to_zarr shape: %s, dtype: %s", image.shape, image.dtype
    602     )
    603     delayed.append(
--> 604         da.to_zarr(
    605             arr=image,
    606             url=group.store,
    607             component=str(Path(group.path, str(path))),
    608             storage_options=options,
    609             compute=False,
    610             compressor=options.get("compressor", zarr.storage.default_compressor),
    611             dimension_separator=group._store._dimension_separator,
    612         )
    613     )
    614     datasets.append({"path": str(path)})
    616 # Computing delayed jobs if necessary

File /opt/homebrew/Caskroom/mambaforge/base/envs/spatialdata/lib/python3.10/site-packages/dask/array/core.py:3875, in to_zarr(arr, url, component, storage_options, overwrite, region, compute, return_stored, **kwargs)
   3872     raise ValueError("Cannot use `region` keyword when url is not a `zarr.Array`.")
   3874 if not _check_regular_chunks(arr.chunks):
-> 3875     raise ValueError(
   3876         "Attempt to save array to zarr with irregular "
   3877         "chunking, please call `arr.rechunk(...)` first."
   3878     )
   3880 storage_options = storage_options or {}
   3882 if storage_options:

ValueError: Attempt to save array to zarr with irregular chunking, please call `arr.rechunk(...)` first.

[wakelin-g]

Hey @LucaMarconato

I ran into this error just now, trying to save a spatial query. I chunked the image before parsing at the dask image object type. any clues on how to deal with this?

versions:

spatialdata: 0.2.5.post0
geopandas: 1.0.1
anndata: 0.11.1
napari_spatialdata: 0.5.3

You can manually rechunk the images before you save. @LucaMarconato had this in an issue somewhere I just can't find it to reference now. Make sure you have the newest version of dask.

def rechunk(msi):
    for scale in msi.keys();
        msi[scale]['image'] = msi[scale]['image'].chunk('auto')


sdata_query = spatialdata.polygon_query(...)
rechunk(sdata_query['morphology_mip'])
rechunk(sdata_query['morphology_focus'])
sdata_query.write('query.zarr')

[LucaMarconato]

Thanks for reporting. I have tried reproducing this bug in the latest main (incoming release) but I could not and the saving worked without problems. Please let me know if the bug is still open for you with the latest main (or please let me know after we make the new release).

If it is still open, it would be very helpful if you could share a minimal working example as a reproducible script. In doing that, it would be ideal if you could use the blobs() dataset:

from spatialdata.datasets import blobs

sdata = blobs()

Thank you!

[josenimo]

Hello friends,

So I tried doing what I was doing with the blobs dataset, and there seems to be an issue writing.

print(spatialdata.__version__)
from spatialdata.datasets import blobs
sdata = blobs()
sdata.write("/Volumes/Extreme SSD/P26_BCCN/P26E20_spatialdata/test_4.zarr")

0.3.0
/opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/anndata/_core/anndata.py:401: FutureWarning: The dtype argument is deprecated and will be removed in late 2024.
  warnings.warn(
/opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/models/models.py:1053: UserWarning: Converting `region_key: region` to categorical dtype.
  return convert_region_column_to_categorical(adata)

Traceback

---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
Cell In[7], line 4
      2 from spatialdata.datasets import blobs
      3 sdata = blobs()
----> 4 sdata.write("/Volumes/Extreme SSD/P26_BCCN/P26E20_spatialdata/test_4.zarr")

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py:1219, in SpatialData.write(self, file_path, overwrite, consolidate_metadata, format)
   1216 store.close()
   1218 for element_type, element_name, element in self.gen_elements():
-> 1219     self._write_element(
   1220         element=element,
   1221         zarr_container_path=file_path,
   1222         element_type=element_type,
   1223         element_name=element_name,
   1224         overwrite=False,
   1225         format=format,
   1226     )
   1228 if self.path != file_path:
   1229     old_path = self.path

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py:1271, in SpatialData._write_element(self, element, zarr_container_path, element_type, element_name, overwrite, format)
   1269     write_shapes(shapes=element, group=element_type_group, name=element_name, format=parsed["shapes"])
   1270 elif element_type == "tables":
-> 1271     write_table(table=element, group=element_type_group, name=element_name, format=parsed["tables"])
   1272 else:
   1273     raise ValueError(f"Unknown element type: {element_type}")

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/_io/io_table.py:89, in write_table(table, group, name, group_type, format)
     87 else:
     88     region, region_key, instance_key = (None, None, None)
---> 89 write_adata(group, name, table)  # creates group[name]
     90 tables_group = group[name]
     91 tables_group.attrs["spatialdata-encoding-type"] = group_type

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/anndata/_io/specs/registry.py:487, in write_elem(store, k, elem, dataset_kwargs)
    463 def write_elem(
    464     store: GroupStorageType,
    465     k: str,
   (...)
    468     dataset_kwargs: Mapping[str, Any] = MappingProxyType({}),
    469 ) -> None:
    470     """
    471     Write an element to a storage group using anndata encoding.
    472 
   (...)
    485         E.g. for zarr this would be `chunks`, `compressor`.
    486     """
--> 487     Writer(_REGISTRY).write_elem(store, k, elem, dataset_kwargs=dataset_kwargs)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/anndata/_io/utils.py:249, in report_write_key_on_error.<locals>.func_wrapper(*args, **kwargs)
    247     raise ValueError("No element found in args.")
    248 try:
--> 249     return func(*args, **kwargs)
    250 except Exception as e:
    251     path = _get_display_path(store)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/anndata/_io/specs/registry.py:349, in Writer.write_elem(self, store, k, elem, dataset_kwargs, modifiers)
    347     store.clear()
    348 elif k in store:
--> 349     del store[k]
    351 write_func = self.find_write_func(dest_type, elem, modifiers)
    353 if self.callback is None:

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/zarr/hierarchy.py:517, in Group.__delitem__(self, item)
    516 def __delitem__(self, item):
--> 517     return self._write_op(self._delitem_nosync, item)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/zarr/hierarchy.py:952, in Group._write_op(self, f, *args, **kwargs)
    949     lock = self._synchronizer[group_meta_key]
    951 with lock:
--> 952     return f(*args, **kwargs)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/zarr/hierarchy.py:524, in Group._delitem_nosync(self, item)
    520 path = self._item_path(item)
    521 if contains_array(self._store, path) or contains_group(
    522     self._store, path, explicit_only=False
    523 ):
--> 524     rmdir(self._store, path)
    525 else:
    526     raise KeyError(item)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/zarr/storage.py:212, in rmdir(store, path)
    209 store_version = getattr(store, "_store_version", 2)
    210 if hasattr(store, "rmdir") and store.is_erasable():  # type: ignore
    211     # pass through
--> 212     store.rmdir(path)
    213 else:
    214     # slow version, delete one key at a time
    215     if store_version == 2:

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/zarr/storage.py:1549, in FSStore.rmdir(self, path)
   1547 store_path = self.dir_path(path)
   1548 if self.fs.isdir(store_path):
-> 1549     self.fs.rm(store_path, recursive=True)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/fsspec/implementations/local.py:172, in LocalFileSystem.rm(self, path, recursive, maxdepth)
    170     if osp.abspath(p) == os.getcwd():
    171         raise ValueError("Cannot delete current working directory")
--> 172     shutil.rmtree(p)
    173 else:
    174     os.remove(p)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/shutil.py:759, in rmtree(path, ignore_errors, onerror, onexc, dir_fd)
    757 try:
    758     while stack:
--> 759         _rmtree_safe_fd(stack, onexc)
    760 finally:
    761     # Close any file descriptors still on the stack.
    762     while stack:

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/shutil.py:703, in _rmtree_safe_fd(stack, onexc)
    701 except OSError as err:
    702     err.filename = path
--> 703     onexc(func, path, err)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/shutil.py:700, in _rmtree_safe_fd(stack, onexc)
    698             os.unlink(entry.name, dir_fd=topfd)
    699         except OSError as err:
--> 700             onexc(os.unlink, fullname, err)
    701 except OSError as err:
    702     err.filename = path

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/shutil.py:698, in _rmtree_safe_fd(stack, onexc)
    696     pass
    697 try:
--> 698     os.unlink(entry.name, dir_fd=topfd)
    699 except OSError as err:
    700     onexc(os.unlink, fullname, err)

FileNotFoundError: [Errno 2] No such file or directory: '/Volumes/Extreme SSD/P26_BCCN/P26E20_spatialdata/test_4.zarr/tables/table'
Error raised while writing key 'table' of <class 'zarr.hierarchy.Group'> to /tables

I ran the same code, but writing to my desktop:

print(spatialdata.__version__)
from spatialdata.datasets import blobs
sdata = blobs()
sdata.write("/Users/jnimoca/Desktop/test.zarr")

0.3.0
INFO     The Zarr backing store has been changed from None the new file path: /Users/jnimoca/Desktop/test.zarr     
/opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/anndata/_core/anndata.py:401: FutureWarning: The dtype argument is deprecated and will be removed in late 2024.
  warnings.warn(
/opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/models/models.py:1053: UserWarning: Converting `region_key: region` to categorical dtype.
  return convert_region_column_to_categorical(adata)

I am not sure why, but the issue of writing to an external drive is blocking the sdata.write()

I am now trying to write my data as sdata into my desktop 🙃
EDIT: it failed

EDIT2:
somehow writing to SSD is working after a couple of kernel restarts...
this worked as expected

print(spatialdata.__version__)
from spatialdata.datasets import blobs
sdata = blobs()
sdata.write("/Volumes/Extreme SSD/P26_BCCN/P26E20_spatialdata/test_4.zarr")

[josenimo]

I tried out rechunking with the approach that @wakelin-g recommends.

It runs into this Value Error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[19], line 2
      1 path_to_store_zarr = "/Users/jnimoca/Desktop/spatialdata_autorechunk.zarr"
----> 2 sdata.write(path_to_store_zarr)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py:1219, in SpatialData.write(self, file_path, overwrite, consolidate_metadata, format)
   1216 store.close()
   1218 for element_type, element_name, element in self.gen_elements():
-> 1219     self._write_element(
   1220         element=element,
   1221         zarr_container_path=file_path,
   1222         element_type=element_type,
   1223         element_name=element_name,
   1224         overwrite=False,
   1225         format=format,
   1226     )
   1228 if self.path != file_path:
   1229     old_path = self.path

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py:1263, in SpatialData._write_element(self, element, zarr_container_path, element_type, element_name, overwrite, format)
   1260 parsed = _parse_formats(formats=format)
   1262 if element_type == "images":
-> 1263     write_image(image=element, group=element_type_group, name=element_name, format=parsed["raster"])
   1264 elif element_type == "labels":
   1265     write_labels(labels=element, group=root_group, name=element_name, format=parsed["raster"])

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/_io/io_raster.py:242, in write_image(image, group, name, format, storage_options, **metadata)
    234 def write_image(
    235     image: DataArray | DataTree,
    236     group: zarr.Group,
   (...)
    240     **metadata: str | JSONDict | list[JSONDict],
    241 ) -> None:
--> 242     _write_raster(
    243         raster_type="image",
    244         raster_data=image,
    245         group=group,
    246         name=name,
    247         format=format,
    248         storage_options=storage_options,
    249         **metadata,
    250     )

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/spatialdata/_io/io_raster.py:202, in _write_raster(raster_type, raster_data, group, name, format, storage_options, label_metadata, **metadata)
    200 parsed_axes = _get_valid_axes(axes=list(input_axes), fmt=format)
    201 storage_options = [{"chunks": chunk} for chunk in chunks]
--> 202 dask_delayed = write_multi_scale_ngff(
    203     pyramid=data,
    204     group=group_data,
    205     fmt=format,
    206     axes=parsed_axes,
    207     coordinate_transformations=None,
    208     storage_options=storage_options,
    209     **metadata,
    210     compute=False,
    211 )
    212 # Compute all pyramid levels at once to allow Dask to optimize the computational graph.
    213 da.compute(*dask_delayed)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/ome_zarr/writer.py:254, in write_multiscale(pyramid, group, chunks, fmt, axes, coordinate_transformations, storage_options, name, compute, **metadata)
    252     data = da.array(data).rechunk(chunks=chunks_opt)
    253     options["chunks"] = chunks_opt
--> 254 da_delayed = da.to_zarr(
    255     arr=data,
    256     url=group.store,
    257     component=str(Path(group.path, str(path))),
    258     storage_options=options,
    259     compressor=options.get("compressor", zarr.storage.default_compressor),
    260     dimension_separator=group._store._dimension_separator,
    261     compute=compute,
    262 )
    264 if not compute:
    265     dask_delayed.append(da_delayed)

File /opt/homebrew/Caskroom/mambaforge/base/envs/sdata_v3/lib/python3.12/site-packages/dask/array/core.py:3875, in to_zarr(arr, url, component, storage_options, overwrite, region, compute, return_stored, **kwargs)
   3872     raise ValueError("Cannot use `region` keyword when url is not a `zarr.Array`.")
   3874 if not _check_regular_chunks(arr.chunks):
-> 3875     raise ValueError(
   3876         "Attempt to save array to zarr with irregular "
   3877         "chunking, please call `arr.rechunk(...)` first."
   3878     )
   3880 storage_options = storage_options or {}
   3882 if storage_options:

ValueError: Attempt to save array to zarr with irregular chunking, please call `arr.rechunk(...)` first.

[josenimo]

I see many issues with dask so I am just letting you know what Dask I am using:

!mamba list | grep dask
dask                      2024.11.2                pypi_0    pypi
dask-expr                 1.1.19                   pypi_0    pypi
dask-image                2024.5.3           pyhd8ed1ab_0    conda-forge

[LucaMarconato]

Hi @josenimo, the chunk('auto') bug is now fixed, the solution is already in main, but you can also manually rechunk with the code used in the PR linked to this issue #821. Basically passing the chunks directly instead of 'auto'. The reason why this occurred is discussed here: dask/dask#11713. So, the bug is fixed with 2024.12.0, but we do not support it yet.

[LucaMarconato]

Regarding the problem that you had while writing to the SSD, it would be great if you could share more info on it (in a separate issue). It doesn't seem to be caused by spatialdata. It could be a problem of zarr or something in between (anndata). I'd try to reproduce it with spatialdata, and if you can, then with anndata, and finally try to reproduce it with zarr. In this way the best repository to discuss it is identified.

I will now close the current issue as it is addressed as detailed in the previous comment.