writing to a sharded array is 10x slower

[nbren12]

Zarr version

3.1.3

Numcodecs version

n/a

Python Version

3.12

Operating System

Linux

Installation

uv

Description

The attached script shows that writing multiple chunks to an uncompressed sharded array is about 10x slower than writing to a non-sharded array. You can reproduce this bechmark with this command:

uv run benchmark_zarr_sharding.py --num-iterations 2 --shape 1000 100000 --chunks 1 100000 --shard-chunks 100 1000000 ^C

Some profiling with cProfile reveals that expenses is dominated by repeated calls to Buffer._add__ from _ShardWriter. Buffer.__add__ concatenates the arrays manually, which will results in many repeated allocations and copies.

            np.concatenate((np.asanyarray(self._data), np.asanyarray(other_array)))

It would probably be much faster to coalesce these concats into a single operation. A similar pattern is used for gpu buffers, so it seems a new api is required to support this use case.

[nbren12]

cc @akshaysubr

[nbren12]

It might be worth looking into how pyarrow handles this concern: https://arrow.apache.org/docs/3.0/python/generated/pyarrow.cuda.BufferWriter.html. I think that project is quite focused on these kinds of performance issues and avoiding unnecessary copies.

[d-v-b]

@nbren12 I don't see your benchmark script, could you share it? And maybe we add something that hits the same notes in #3562

[nbren12]

Here is the script. benchmark_zarr_sharding.py

[d-v-b]

ok looks like the main thing to add to #3562 would be a partial read / write workload