Support sharding through config and raster_write_kwargs by melonora · Pull Request #1106 · scverse/spatialdata

melonora · 2026-04-14T07:52:13Z

This PR adds the following:

passing kwargs for zarr.create_array directly as raster_write_kwargs for io functions like .write and .write_element. This also adds the ability to write sharded arrays. Support for anndata sharding is to be added in a follow up PR.
proper docstrings for the new raster_write_kwargs argument.
Extension of the current config to include raster_chunks and raster_shards. The config can now be stored in a default location or a custom location. Additionally, environment variables can be set to temporarily override the values.
Adding zarrs as a dependency and enabled the codec by default to allow for faster io when writing shards. This is a discussion point of whether we should do this or provide more of an opt-in for advanced users.

Additional changes

Minimal supported version of dask is 2026.3.0. The reason here is that only this provides the api in such a way that you don't risk zarr format 2 being written in a zarr v3 group and vice versa + it includes the setting that prevents collaps of partitions of dask dataframes after reading from parquet.

@LucaMarconato

- Added additional settings - Allow environment variables that overwrite config

melonora · 2026-04-14T07:55:09Z

Failing atm due to ome-zarr not yet being released. You can test locally with ome-zarr-py from main.

Also, need to add support for zarrs to improve speed of shard io

codecov · 2026-04-14T14:10:37Z

Codecov Report

❌ Patch coverage is 90.09009% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.44%. Comparing base (68dade6) to head (1cd0eb3).

Files with missing lines	Patch %	Lines
src/spatialdata/config.py	87.50%	9 Missing ⚠️
src/spatialdata/_core/_utils.py	90.00%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main    #1106    +/-   ##
========================================
  Coverage   92.44%   92.44%            
========================================
  Files          51       51            
  Lines        7811     7918   +107     
========================================
+ Hits         7221     7320    +99     
- Misses        590      598     +8

Files with missing lines	Coverage Δ
src/spatialdata/_core/spatialdata.py	`93.93% <100.00%> (+0.07%)`	⬆️
src/spatialdata/_io/io_raster.py	`90.90% <100.00%> (+1.65%)`	⬆️
src/spatialdata/_core/_utils.py	`97.05% <90.00%> (-2.95%)`	⬇️
src/spatialdata/config.py	`88.60% <87.50%> (-11.40%)`	⬇️

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

The reason for only supporting these versions is that they provide the proper use of the zarr api inside dask and also the possibility for setting the tune optimization. The latter is required to prevent errors due to collapsing dask partitions when reading data back in from parquet.

Mr-Milk · 2026-04-15T10:30:26Z

Should we also allow the control of sharding for anndata?

melonora · 2026-04-16T11:41:53Z

Yes, but not as part of this PR. I will adjust the config though to accommodate.

LucaMarconato · 2026-05-12T13:53:54Z

    "xarray>=2024.10.0",
    "xarray-spatial>=0.3.5",
    "zarr>=3.0.0",
+    "zarrs",


Let's follow up on this.

LucaMarconato · 2026-05-12T14:02:49Z

+

 @dataclass
 class Settings:


General comment. Sounds like this file is a duplicate of the settings mechanism from scverse-misc: https://github.com/scverse/scverse-misc/blob/main/src/scverse_misc/_settings.py.

I think scverse-misc doesn't support reading settings from .json (they support .env files), but since it's based on Pydantic, it would be a quick extension.

I'll go trough the PR and review. In a follow-up PR we should probably align to use scverse-misc.

@ilia-kats @flying-sheep, devs behind scverse-misc, may have extra comments.

would this be something for this PR or follow up PR before release? I believe before settings were not stored at all in any case and I think in general we need a PR to move towards using scverse-misc, also with adding the context manager there.

Since config.py works, we can do that on a follow up release to avoid delaying this PR even more. On the other hand, if there is time/energy/will, I'd go directly with scverse-misc so that we don't merge code that will be removed soon after. Up to you!

LucaMarconato · 2026-05-12T14:07:00Z

+    raster_chunks
+        The chunksize to use for chunking an array. Length of the tuple must match
+        the number of dimensions.
+    raster_shards
+        The default shard size (zarr v3) to use when storing arrays. Length of the tuple
+        must match the number of dimensions.


Feels a bit weird to have these 2 variables in a class that is for global settings, since, if I am getting this right, we for cyx data we would need len(raster_chunks) == 3, while for yx data we would need len(raster_chunks) == 2.

want to split in image_shards/chunks and label or just remove all together? I would maybe vote for the former.

there is still the problem of 2D vs 3D. I'd remove altogether. We could use instead an helper function that computes reasonable defaults based on the input shape.

LucaMarconato · 2026-05-12T14:08:18Z

+        must match the number of dimensions.
    """

+    custom_config_path: Path | None = None


I don't think the settings should contain this path. Feels out of place.

depends on whether certain settings are project specific I believe. This allows for that. Otherwise global settings on hpc could be shared by more than 1 user even though there are project specific needs.

LucaMarconato · 2026-05-12T14:28:19Z

+def _config_path() -> Path:
+    """Return the platform-appropriate path to the user config file."""
+    return Path(user_config_dir(appname="spatialdata")) / "settings.json"


why not moving this inside config_path()?

LucaMarconato · 2026-05-12T14:32:56Z

+        assert s.large_chunk_threshold_bytes == 1_000_000_000
+        assert s.raster_chunks == [512, 512]
+        assert s.raster_shards == [1024, 1024]
+        assert s.custom_config_path is None


Fine. But since we called save(), shouldn't custom_config_path be equal to the path we are reading from?

This is a consequence of https://github.com/melonora/spatialdata/blob/f44221ab23d6b136084157203e1d0b466979318c/src/spatialdata/config.py#L65

path is None, so it gets assigned the value from _config_path(), without saving it in the in-memory value of custom_config_path.

LucaMarconato · 2026-05-12T14:33:43Z

+
+        s.reset()
+        s.save()
+        assert s.custom_config_path is None  # This returns False


what's the comment about?

LucaMarconato · 2026-05-12T14:37:31Z

I haven't made it a default dependency for anndata but it will be enabled by default if discovered.

@melonora I would do the same that @ilan-gold is doing in anndata, and therefore drop the dependency (we can add this as an optional dependency).

LucaMarconato · 2026-05-12T14:51:07Z

+    assert arr.shards == write_shards
+
+    other_arr = zarr.open_group(path / zarr_subpath / other_name, mode="r")["s0"]
+    assert other_arr.chunks == base_chunks


we could add an explicit check that shards are None (?) here.

LucaMarconato · 2026-05-12T14:53:16Z

+    assert other_arr.chunks == base_chunks
+
+
+def test_write_raster_elements_sharding_chunking(tmp_path: Path) -> None:


What are we testing that it is not tested before? That write_element() supports sharding? In that case I'd rename the test.

yeah did that independently because of issues downstream with non matching shards / chunks. These errors are specific and not connected to our implementation per se so I kept it separate. For me it indicates that by _sharding_chunking which is seen not to be tested in the other functions.

LucaMarconato · 2026-05-12T14:54:50Z

Looks good in general. There are a few decisions to be made, but once done most of the requested changes can be one-shot by an agent.

chore: fix typo in docstrings

…to support_sharding

bump ome_zarr remove distributed add platformdirs

LucaMarconato

Changes look good so far.

melonora added 6 commits March 30, 2026 08:44

build: change supported dask versions

7039c00

feat[config]: allow for persisting config

63fbe9a

- Added additional settings - Allow environment variables that overwrite config

test: add config fixture

b7e98b9

test: add tests for config

558fe97

add raster write kwargs to api

8eebdb0

add tests for raster API

790be0c

melonora added 4 commits April 14, 2026 12:07

build: add zarrs-python for improved shard io

1a1c673

CI: change lowerbound dask version

471f72c

build: correct zarrs

cd48574

build: change distributed version constraint

0187fe9

melonora added 5 commits April 14, 2026 17:14

CI: change lowerbound test version of dask

c2b1375

fix: pre-commit error due to incorrect typehint

bf5b910

build: include zarrs as dependency

49441fe

make zarrs codec default

6334fa8

melonora added 11 commits April 16, 2026 13:48

config: change chunks and shards to accomodate for raster and table

73ca72a

chore: adjust raster_write to new config fields

c6041bb

docs: add docstring

278606a

change: add support for providing storage options as list

f6c0a37

docs: make docstring for raster write kwargs more clear

36e2271

tests: complete and refactor sharding tests

50cb6bb

docs: adjust docstring config

6257096

fix: correct parsing chunks, shards argument

5525fbf

test: add testing for adjusting chunks with env variable

44414c1

test: write using settings raster_chunks

e974647

feat: add raster_write_kwargs to write_element

bd6249e