Repository: scverse/spatialdata

Issue:
Error (Codec does not support buffers > 2GB) when trying to save zarr directory
Hi,

i have experienced an error when trying to load the output of 10x xenium into a spatialdata object, attaching a HE image, and trying to save it as an zarr files. I iterated over many samples with the same loop and it works for most of them, so the code per se works. However, for some samples i get the error 

ValueError: Codec does not support buffers of > 2147483647 bytes

when i try to save the sdata object.

Has anyone else experienced this problem before? How can i fix it so i get the zarr files for all my samples?

Thanks in adance,
Korbinian


I used this script to produce the error:

```
import spatialdata_io as sdio

sdata = sdio.xenium("spatial_data/outs")

he_image_datatree = sdio.xenium_aligned_image(
    image_path="HE.tif",
    alignment_file=None,
    dims=None  
)

sdata.images['he_image'] = he_image_datatree

sdata.write("test.zarr")
```

Hints from issue comments:
@ckmah: The HE image is larger than the default maximum. Try this:

```
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
```@kotr98: Thanks for the fast reply. Unfortunately it did not resolve the error:

I can reproduce it with this data sample: https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?acc=GSM7990532 
I downloaded the GSM7990532_warped_THD0008-TMA3_HE.tif.gz HE image and GSM7990532_output-XETG00048__0003392__THD0008__20230313__191400.tar.gz and subsequently unzipped them. 

```
import spatialdata_io as sdio
from PIL import Image
Image.MAX_IMAGE_PIXELS = None


sdata = sdio.xenium("relabel_output-XETG00048__0003392__THD0008__20230313__191400/outs")

he_image_datatree = sdio.xenium_aligned_image(
    image_path="GSM7990532_warped_THD0008-TMA3_HE.tif",
    alignment_file=None,
    dims=None  
)

sdata.images['he_image'] = he_image_datatree

sdata.write("test.zarr")
```
@ckmah: I would try to manually set the chunk sizes, assuming your H&E images is shaped [3, y, x], something like:

```python
...

he_image_datatree = sdio.xenium_aligned_image(
    image_path="GSM7990532_warped_THD0008-TMA3_HE.tif",
    alignment_file=None,
    dims=["c", "y", "x"] ,
    image_models_kwargs={"chunks": (1, 4000, 4000)}} # adjust to be smaller/larger as preferred
)

...
@kotr98: Thanks for the suggestions. Setting the chunks parameters did not solve the problem unfortunately. I also tried different/smaller chunk sizes. Additionally, I tried saving it on my local disc instead of an external ssd because someone else had this issue. It did also not solve the error. 

Was anyone able to reproduce the error with the data I provided?@timtreis: The error is caused by the compression algo not being able to compress anything bigger than 2GB at once. Since your image is loaded as a single chunk, that fails. You can fix this by simply chunking your image, f.e. like this:

```
he_image_datatree = sdio.xenium_aligned_image(
    image_path="hne.tif",
    alignment_file=None,
    dims=None  
)
he_image_datatree = he_image_datatree.chunk({"y": 512, "x": 512})
he_image_datatree
```

I get the same error without re-chunking and it works after adding this line. Could you try and confirm @kotr98 ?

![Image](https://github.com/user-attachments/assets/8850263c-3b5a-4c2d-8391-a8dc10523d8c)


@LucaMarconato do we want to add a safety check here? If Korbinian has been adding the argument to `image_model_kwargs`, there might be a bug that the chunking is not correctly parsed.

FYI: Changed the title to be more discoverable@kotr98: With the new code, it works without errors. Thanks for resolving this issue.@LucaMarconato: Thanks @timtreis!

> @LucaMarconato do we want to add a safety check here? If Korbinian has been adding the argument to image_model_kwargs, there might be a bug that the chunking is not correctly parsed.

Sounds great to me! Could you please make a small PR?@LucaMarconato: I reopen this as we should add a warning indeed. 

Solution patch:
diff --git a/src/spatialdata/config.py b/src/spatialdata/config.py
new file mode 100644
index 0000000..6d8dcf2
--- /dev/null
+++ b/src/spatialdata/config.py
@@ -0,0 +1,4 @@
+# chunk sizes bigger than this value can trigger a compression error
+# https://github.com/scverse/spatialdata/issues/812#issuecomment-2559380276
+# so if we detect this during parsing/validation we raise a warning
+MAX_N_ELEMS_CHUNK_SIZE = 2147483647
diff --git a/src/spatialdata/models/models.py b/src/spatialdata/models/models.py
index b2663e5..257baaa 100644
--- a/src/spatialdata/models/models.py
+++ b/src/spatialdata/models/models.py
@@ -35,6 +35,7 @@ from spatialdata._core.validation import validate_table_attr_keys
 from spatialdata._logging import logger
 from spatialdata._types import ArrayLike
 from spatialdata._utils import _check_match_length_channels_c_dim
+from spatialdata.config import MAX_N_ELEMS_CHUNK_SIZE
 from spatialdata.models import C, X, Y, Z, get_axes_names
 from spatialdata.models._utils import (
     DEFAULT_COORDINATE_SYSTEM,
@@ -127,7 +128,13 @@ class RasterSchema(DataArraySchema):
             are `[2, 2, 2]`, the returned multiscale image will have 4 scales. The original image and then the 2x, 4x
             and 8x downsampled images.
         method
-            Method to use for multiscale downsampling. Please refer to :class:`multiscale_spatial_image.to_multiscale`.
+            Method to use for multiscale downsampling (default is `'nearest'`). Please refer to
+            :class:`multiscale_spatial_image.to_multiscale` for details.\n
+            Note (advanced): the default choice (`'nearest'`) will keep the original scale lazy and compute each
+            downscaled version. On the other hand `'xarray_coarsen'` will compute each scale lazily (this implies that
+            each scale will be recomputed each time it is accessed unless `.persist()` is manually called to cache the
+            intermediate results). Please refer direclty to the source code of `to_multiscale()` in the
+            `multiscale-spatial-image` for precise information on how this is handled.
         chunks
             Chunks to use for dask array.
         kwargs
@@ -225,6 +232,7 @@ class RasterSchema(DataArraySchema):
                 chunks=chunks,
             )
             _parse_transformations(data, parsed_transform)
+        cls()._check_chunk_size_not_too_large(data)
         # recompute coordinates for (multiscale) spatial image
         return compute_coordinates(data)
 
@@ -251,6 +259,7 @@ class RasterSchema(DataArraySchema):
     @validate.register(DataArray)
     def _(self, data: DataArray) -> None:
         super().validate(data)
+        self._check_chunk_size_not_too_large(data)
 
     @validate.register(DataTree)
     def _(self, data: DataTree) -> None:
@@ -263,6 +272,48 @@ class RasterSchema(DataArraySchema):
         name = list(name)[0]
         for d in data:
             super().validate(data[d][name])
+        self._check_chunk_size_not_too_large(data)
+
+    def _check_chunk_size_not_too_large(self, data: DataArray | DataTree) -> None:
+        if isinstance(data, DataArray):
+            try:
+                max_per_dimension: dict[int, int] = {}
+                if isinstance(data.chunks, list | tuple):
+                    for i, sizes in enumerate(data.chunks):
+                        max_per_dimension[i] = max(sizes)
+                else:
+                    assert isinstance(data.chunks, dict)
+                    for i, sizes in enumerate(data.chunks.values()):
+                        max_per_dimension[i] = max(sizes)
+            except ValueError:
+                warnings.warn(
+                    f"Unable to estimate the maximum chunk size for the data: {sizes}. Please report this bug.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                return
+            n_elems = np.array(list(max_per_dimension.values())).prod().item()
+            if n_elems > MAX_N_ELEMS_CHUNK_SIZE:
+                warnings.warn(
+                    f"Detected a large number of elements for some chunks: {n_elems}. This can lead to low "
+                    "performance and memory issues downstream, and sometimes cause compression errors when writing "
+                    "(https://github.com/scverse/spatialdata/issues/812#issuecomment-2575983527). Please consider using"
+                    " 1) smaller chunks and/or 2) using a multiscale representation for the raster data.\n"
+                    "1) Smaller chunks can be achieved by using the `chunks` argument in the `parse()` function or by "
+                    "calling the `chunk()` method on `DataArray`/`DataTree` objects.\n"
+                    "2) Multiscale representations can be achieved by using the `scale_factors` argument in the "
+                    "`parse()` function.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+        else:
+            assert isinstance(data, DataTree)
+            name = {list(data[i].data_vars.keys())[0] for i in data}
+            assert len(name) == 1
+            name = list(name)[0]
+            for d in data:
+                super().validate(data[d][name])
+            self._check_chunk_size_not_too_large(data[d][name])
 
 
 class Labels2DModel(RasterSchema):
@@ -582,7 +633,12 @@ class PointsModel:
         SUGGESTION = " Please use PointsModel.parse() to construct data that is guaranteed to be valid."
         for ax in [X, Y, Z]:
             # TODO: check why this can return int32 on windows.
-            if ax in data.columns and data[ax].dtype not in [np.int32, np.float32, np.float64, np.int64]:
+            if ax in data.columns and data[ax].dtype not in [
+                np.int32,
+                np.float32,
+                np.float64,
+                np.int64,
+            ]:
                 raise ValueError(f"Column `{ax}` must be of type `int` or `float`.")
         if cls.TRANSFORM_KEY not in data.attrs:
             raise ValueError(
@@ -675,7 +731,10 @@ class PointsModel:
 
         table: DaskDataFrame = dd.from_pandas(pd.DataFrame(**df_kwargs), **kwargs)
         return cls._add_metadata_and_validate(
-            table, feature_key=feature_key, instance_key=instance_key, transformations=transformations
+            table,
+            feature_key=feature_key,
+            instance_key=instance_key,
+            transformations=transformations,
         )
 
     @parse.register(pd.DataFrame)
@@ -719,7 +778,11 @@ class PointsModel:
             )
         if isinstance(data, pd.DataFrame):
             table: DaskDataFrame = dd.from_pandas(
-                pd.DataFrame(data[[coordinates[ax] for ax in axes]].to_numpy(), columns=axes, index=data.index),
+                pd.DataFrame(
+                    data[[coordinates[ax] for ax in axes]].to_numpy(),
+                    columns=axes,
+                    index=data.index,
+                ),
                 # we need to pass sort=True also when the index is sorted to ensure that the divisions are computed
                 sort=sort,
                 **kwargs,
@@ -752,11 +815,21 @@ class PointsModel:
                 )
         if Z not in axes and Z in data.columns:
             logger.info(f"Column `{Z}` in `data` will be ignored since the data is 2D.")
-        for c in set(data.columns) - {feature_key, instance_key, *coordinates.values(), X, Y, Z}:
+        for c in set(data.columns) - {
+            feature_key,
+            instance_key,
+            *coordinates.values(),
+            X,
+            Y,
+            Z,
+        }:
             table[c] = data[c]
 
         validated = cls._add_metadata_and_validate(
-            table, feature_key=feature_key, instance_key=instance_key, transformations=transformations
+            table,
+            feature_key=feature_key,
+            instance_key=instance_key,
+            transformations=transformations,
         )
 
         # when `coordinates` is None, and no columns have been added or removed, preserves the original order
@@ -1047,7 +1120,11 @@ class TableModel:
                 f"Instance key column for region(s) `{', '.join(not_unique)}` does not contain only unique values"
             )
 
-        attr = {"region": region, "region_key": region_key, "instance_key": instance_key}
+        attr = {
+            "region": region,
+            "region_key": region_key,
+            "instance_key": instance_key,
+        }
         adata.uns[cls.ATTRS_KEY] = attr
         cls().validate(adata)
         return convert_region_column_to_categorical(adata)
@@ -1124,7 +1201,11 @@ def get_table_keys(table: AnnData) -> tuple[str | list[str], str, str]:
     """
     if table.uns.get(TableModel.ATTRS_KEY):
         attrs = table.uns[TableModel.ATTRS_KEY]
-        return attrs[TableModel.REGION_KEY], attrs[TableModel.REGION_KEY_KEY], attrs[TableModel.INSTANCE_KEY]
+        return (
+            attrs[TableModel.REGION_KEY],
+            attrs[TableModel.REGION_KEY_KEY],
+            attrs[TableModel.INSTANCE_KEY],
+        )
 
     raise ValueError(
         "No spatialdata_attrs key found in table.uns, therefore, no table keys found. Please parse the table."


Test patch:
diff --git a/tests/models/test_models.py b/tests/models/test_models.py
index 012ca03..529ca24 100644
--- a/tests/models/test_models.py
+++ b/tests/models/test_models.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import os
 import re
 import tempfile
+import warnings
 from collections.abc import Callable
 from functools import partial
 from pathlib import Path
@@ -27,6 +28,7 @@ from xarray import DataArray, DataTree
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._core.validation import ValidationError
 from spatialdata._types import ArrayLike
+from spatialdata.config import MAX_N_ELEMS_CHUNK_SIZE
 from spatialdata.models._utils import (
     force_2d,
     points_dask_dataframe_to_geopandas,
@@ -155,7 +157,11 @@ class TestModels:
     @pytest.mark.parametrize("permute", [True, False])
     @pytest.mark.parametrize("kwargs", [None, {"name": "test"}])
     def test_raster_schema(
-        self, converter: Callable[..., Any], model: RasterSchema, permute: bool, kwargs: dict[str, str] | None
+        self,
+        converter: Callable[..., Any],
+        model: RasterSchema,
+        permute: bool,
+        kwargs: dict[str, str] | None,
     ) -> None:
         dims = np.array(model.dims.dims).tolist()
         if permute:
@@ -346,14 +352,18 @@ class TestModels:
     ) -> None:
         region_key = "reg"
         obs = pd.DataFrame(
-            RNG.choice(np.arange(0, 100, dtype=float), size=(10, 3), replace=False), columns=["A", "B", "C"]
+            RNG.choice(np.arange(0, 100, dtype=float), size=(10, 3), replace=False),
+            columns=["A", "B", "C"],
         )
         obs[region_key] = region
         adata = AnnData(RNG.normal(size=(10, 2)), obs=obs)
         with pytest.raises(TypeError, match="Only int"):
             model.parse(adata, region=region, region_key=region_key, instance_key="A")
 
-        obs = pd.DataFrame(RNG.choice(np.arange(0, 100), size=(10, 3), replace=False), columns=["A", "B", "C"])
+        obs = pd.DataFrame(
+            RNG.choice(np.arange(0, 100), size=(10, 3), replace=False),
+            columns=["A", "B", "C"],
+        )
         obs[region_key] = region
         adata = AnnData(RNG.normal(size=(10, 2)), obs=obs)
         table = model.parse(adata, region=region, region_key=region_key, instance_key="A")
@@ -409,7 +419,10 @@ class TestModels:
             model.parse(adata, region=region, region_key=region_key, instance_key="A")
 
         adata.obs["A"] = [1] * 10
-        with pytest.raises(ValueError, match=re.escape("Instance key column for region(s) `sample_1, sample_2`")):
+        with pytest.raises(
+            ValueError,
+            match=re.escape("Instance key column for region(s) `sample_1, sample_2`"),
+        ):
             model.parse(adata, region=region, region_key=region_key, instance_key="A")
 
     @pytest.mark.parametrize(
@@ -431,7 +444,10 @@ class TestModels:
         if attr in ("obs", "var"):
             df = pd.DataFrame([[None]], columns=[key], index=["1"])
             adata = AnnData(np.array([[0]]), **{attr: df})
-            with pytest.raises(ValueError, match=f"Table contains invalid names(.|\n)*\n  {attr}/{re.escape(key)}"):
+            with pytest.raises(
+                ValueError,
+                match=f"Table contains invalid names(.|\n)*\n  {attr}/{re.escape(key)}",
+            ):
                 if parse:
                     TableModel.parse(adata)
                 else:
@@ -440,14 +456,20 @@ class TestModels:
             if attr in ("obsm", "varm", "obsp", "varp", "layers"):
                 array = np.array([[0]])
                 adata = AnnData(np.array([[0]]), **{attr: {key: array}})
-                with pytest.raises(ValueError, match=f"Table contains invalid names(.|\n)*\n  {attr}/{re.escape(key)}"):
+                with pytest.raises(
+                    ValueError,
+                    match=f"Table contains invalid names(.|\n)*\n  {attr}/{re.escape(key)}",
+                ):
                     if parse:
                         TableModel.parse(adata)
                     else:
                         TableModel().validate(adata)
             elif attr == "uns":
                 adata = AnnData(np.array([[0]]), **{attr: {key: {}}})
-                with pytest.raises(ValueError, match=f"Table contains invalid names(.|\n)*\n  {attr}/{re.escape(key)}"):
+                with pytest.raises(
+                    ValueError,
+                    match=f"Table contains invalid names(.|\n)*\n  {attr}/{re.escape(key)}",
+                ):
                     if parse:
                         TableModel.parse(adata)
                     else:
@@ -570,11 +592,27 @@ def test_force2d():
 
     expected_circles_2d = ShapesModel.parse(GeoDataFrame({"geometry": (Point(1, 1), Point(2, 2)), "radius": [2, 2]}))
     expected_polygons_2d = ShapesModel.parse(
-        GeoDataFrame({"geometry": [Polygon([(0, 0), (1, 0), (1, 1)]), Polygon([(0, 0), (1, 0), (1, 1)])]})
+        GeoDataFrame(
+            {
+                "geometry": [
+                    Polygon([(0, 0), (1, 0), (1, 1)]),
+                    Polygon([(0, 0), (1, 0), (1, 1)]),
+                ]
+            }
+        )
     )
     expected_multipolygons_2d = ShapesModel.parse(
         GeoDataFrame(
-            {"geometry": [MultiPolygon([Polygon([(0, 0), (1, 0), (1, 1)]), Polygon([(0, 0), (1, 0), (1, 1)])])]}
+            {
+                "geometry": [
+                    MultiPolygon(
+                        [
+                            Polygon([(0, 0), (1, 0), (1, 1)]),
+                            Polygon([(0, 0), (1, 0), (1, 1)]),
+                        ]
+                    )
+                ]
+            }
         )
     )
 
@@ -611,7 +649,11 @@ def test_dask_points_unsorted_index_with_xfail(points):
 # helper function to create random points data and write to a Parquet file, used in the test below
 def create_parquet_file(temp_dir, num_points=20, sorted_index=True):
     df = pd.DataFrame(
-        {"x": RNG.uniform(size=num_points), "y": RNG.uniform(size=num_points), "z": RNG.uniform(size=num_points)}
+        {
+            "x": RNG.uniform(size=num_points),
+            "y": RNG.uniform(size=num_points),
+            "z": RNG.uniform(size=num_points),
+        }
     )
     if not sorted_index:
         new_order = RNG.permutation(len(df))
@@ -655,10 +697,47 @@ def test_c_coords_2d(scale_factors: list[int] | None):
         )
 
     with pytest.raises(ValueError, match="The number of channel names"):
-        Image2DModel().parse(data, c_coords=["1st", "2nd", "3rd", "too_much"], scale_factors=scale_factors)
+        Image2DModel().parse(
+            data,
+            c_coords=["1st", "2nd", "3rd", "too_much"],
+            scale_factors=scale_factors,
+        )
 
 
 @pytest.mark.parametrize("model", [Labels2DModel, Labels3DModel])
 def test_label_no_c_coords(model: Labels2DModel | Labels3DModel):
     with pytest.raises(ValueError, match="`c_coords` is not supported"):
         model().parse(np.zeros((30, 30)), c_coords=["1st", "2nd", "3rd"])
+
+
+def test_warning_on_large_chunks():
+    data_small = DataArray(dask.array.zeros((100, 100), chunks=(50, 50)), dims=["x", "y"])
+    data_large = DataArray(dask.array.zeros((50000, 50000), chunks=(50000, 50000)), dims=["x", "y"])
+    assert np.array(data_large.shape).prod().item() > MAX_N_ELEMS_CHUNK_SIZE
+
+    # single and multiscale, small chunk size
+    with warnings.catch_warnings(record=True) as w:
+        _ = Labels2DModel.parse(data_small)
+        _ = Labels2DModel.parse(data_small, scale_factors=[2, 2])
+        # method 'xarray_coarsen' is used to downsample the data lazily (otherwise the test would be too slow)
+        _ = Labels2DModel.parse(data_large, scale_factors=[2, 2], method="xarray_coarsen")
+        warnings.simplefilter("always")
+        assert len(w) == 0, "Warning should not be raised for valid chunk size"
+
+    # single scale, large chunk size
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        _ = Labels2DModel.parse(data_large)
+        assert len(w) == 1, "Warning should be raised for large chunk size"
+        assert issubclass(w[-1].category, UserWarning)
+        assert "Detected a large number of elements for some chunks" in str(w[-1].message)
+
+    # multiscale, large chunk size
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        multiscale = Labels2DModel.parse(data_large, scale_factors=[2, 2], method="xarray_coarsen")
+        multiscale = multiscale.chunk({"x": 50000, "y": 50000})
+        Labels2DModel().validate(multiscale)
+        assert len(w) == 1, "Warning should be raised for large chunk size"
+        assert issubclass(w[-1].category, UserWarning)
+        assert "Detected a large number of elements for some chunks" in str(w[-1].message)