Unable to save subsetted Xenium `sdata` to `.zarr`

[christinedien]

The bug can be reproduced using one of the spatialdata example datasets: I had used the Xenium Aligned Breast Cancer dataset

Describe the bug
After subsetting a spatialdata object using sdata.query.bounding_box(...), attempting to save the new subset as a .zarr sometimes results in an irregular chunking error. This occurs after a sdata.write(...) command

To Reproduce
Code to reproduce the behavior:

import spatialdata as sd
import spatialdata_io
import squidpy as sq
import spatialdata_plot

sdata = sd.read_zarr("/path/to/data_aligned.zarr/")

The following subset calls DO NOT result in an error

subset = sdata.query.bounding_box(
 axes=["x", "y"],
 min_coordinate=[0,0],
 max_coordinate=[15000,13000],
 target_coordinate_system="global",
)

subset.write(f'../data/xenium/subset_test/subset.zarr', overwrite=True)

subset3 = sdata.query.bounding_box(
 axes=["x", "y"],
 min_coordinate=[0,0],
 max_coordinate=[20000,13000],
 target_coordinate_system="global",
)

subset3.write(f'../data/xenium/subset_test/subset3.zarr', overwrite=True)

The following subsetting examples DO result in an error:

subset2 = sdata.query.bounding_box(
 axes=["x", "y"],
 min_coordinate=[15000,0],
 max_coordinate=[35000,13000],
 target_coordinate_system="global",
)

subset2.write(f'../data/xenium/subset_test/subset2.zarr', overwrite=True)

subset4 = sdata.query.bounding_box(
 axes=["x", "y"],
 min_coordinate=[10000,0],
 max_coordinate=[30000,13000],
 target_coordinate_system="global",
)
subset4.write(f'../data/xenium/subset_test/subset4.zarr', overwrite=True)

I initially thought maybe the subset size may have something to do with the irregular chunking. In this case the query size (20,000 x 13,000) in subset2 and subset4 resulted in an error but subset3 has the same query size and was able to be saved.

In my private dataset, only 1/12 of my queries is savable

Expected behavior
I would have expected the subsetted spatialdata object to be savable to a .zarr

Screenshots

sdata

subset3 (no error)

subset4 (error)

Desktop (optional):
OS: Red Hat Enterprise Linux release 8.8 (Ootpa)

Additional context
This bug was briefly discussed in a thread of this PR

The bug results in the following traceback:

---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[48], line 1
----> 1 subset4.write(f'../data/xenium/subset_test/subset4.zarr', overwrite=True)
<conda_env_loc>
File [<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py:1186](<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py#line=1185), in SpatialData.write(self, file_path, overwrite, consolidate_metadata, format)
 1183 store.close()
 1185 for element_type, element_name, element in self.gen_elements():
-> 1186 self._write_element(
 1187 element=element,
 1188 zarr_container_path=file_path,
 1189 element_type=element_type,
 1190 element_name=element_name,
 1191 overwrite=False,
 1192 format=format,
 1193 )
 1195 if self.path != file_path:
 1196 old_path = self.path

File [<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py:1230](<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_core/spatialdata.py#line=1229), in SpatialData._write_element(self, element, zarr_container_path, element_type, element_name, overwrite, format)
 1227 parsed = _parse_formats(formats=format)
 1229 if element_type == "images":
-> 1230 write_image(image=element, group=element_type_group, name=element_name, format=parsed["raster"])
 1231 elif element_type == "labels":
 1232 write_labels(labels=element, group=root_group, name=element_name, format=parsed["raster"])

File [<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_io/io_raster.py:242](<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_io/io_raster.py#line=241), 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 [<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_io/io_raster.py:202](<conda_env_loc>/lib/python3.12/site-packages/spatialdata/_io/io_raster.py#line=201), 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 [<conda_env_loc>/lib/python3.12/site-packages/ome_zarr/writer.py:253](<conda_env_loc>/lib/python3.12/site-packages/ome_zarr/writer.py#line=252), in write_multiscale(pyramid, group, chunks, fmt, axes, coordinate_transformations, storage_options, name, compute, **metadata)
 251 data = da.array(data).rechunk(chunks=chunks_opt)
 252 options["chunks"] = chunks_opt
--> 253 da_delayed = da.to_zarr(
 254 arr=data,
 255 url=group.store,
 256 component=str(Path(group.path, str(path))),
 257 storage_options=options,
 258 compressor=options.get("compressor", zarr.storage.default_compressor),
 259 dimension_separator=group._store._dimension_separator,
 260 compute=compute,
 261 )
 263 if not compute:
 264 dask_delayed.append(da_delayed)

File [<conda_env_loc>/lib/python3.12/site-packages/dask/array/core.py:3875](<conda_env_loc>/lib/python3.12/site-packages/dask/array/core.py#line=3874), 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.

Click to add a cell.

[LucaMarconato]

Thank you for the bug report. I will investigate this soon.

[austinv11]

Hi I am encountering the same issue, has there been any update?

[LucaMarconato]

Thanks @austinv11 for joining the conversation and thanks again @christinedien for the comprehensive bug report! It was very helpful to investigate the problem.

I dug into this and the bug seems to be related to this Dask bug: dask/dask#11544. The bug should have been fixed, but I managed to reproduce the problem you reported with pure dask code with Dask 2025.1.0, as I reported here: dask/dask#11713.

Even if the bug is fixed upstream, unfortunately we cannot install Dask versions >= 2025.1.0 due to this other Dask bug: dask/dask#11146. I will therefore implement a solution that works also with older versions of Dask. As the example in the linked issue shows, the bug seems to affect only the 'auto' option of rechunk(). Passing a fixed chunk size is a workaround, but it would have been nice to let chunk() (from xarray) or rechunk() (from dask) infer the chunk size automatically. Anyway if you need a specific chunk size you can always rechunk the data manually after the query.

[austinv11]

@LucaMarconato I appreciate the swift response! All of that makes sense.

I am not super familiar with dask and xarray, do you have guidance on how to select an appropriate chunk size? Specifically with regard to spatialdata queries?

[LucaMarconato]

@austinv11 you are welcome. Please find a workaround implemented here: #861.

Please let me know if you get to try the code and if the solution fixes your problem.

[lkmartin90]

Hi, I'm having the same issue and really struggling to follow the chain of github issues to the solution. I can't update dask as it doesn't work with the rest of the environment I'm using, and I can't find how to rechunk a spatial data object - is there an easy fix?
Thanks!

[christinedien]

Hi @lkmartin90, sorry to hear you are experiencing the same issue! @LucaMarconato's implenetation in #861 solved the issue for me (Thank you Luca by the way!!). It was merged into the main GitHub branch back in February, so if you install the latest version from the github (instead of through pip) it should have the fix! Good luck!

It doesn't require updating dask which isn't an option, as mentioned above:

Even if the bug is fixed upstream, unfortunately we cannot install Dask versions >= 2025.1.0 due to this other Dask bug: dask/dask#11146. I will therefore implement a solution that works also with older versions of Dask. As the example in the linked issue shows, the bug seems to affect only the 'auto' option of rechunk(). Passing a fixed chunk size is a workaround, but it would have been nice to let chunk() (from xarray) or rechunk() (from dask) infer the chunk size automatically. Anyway if you need a specific chunk size you can always rechunk the data manually after the query.

[lkmartin90]

Ah perfect, that's worked instantly, thanks!
However, I can't open the new file with napari's spatialdata plugin without it throwing an error (@LucaMarconato if the change was going to be permanent)