|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +import numpy as np |
| 3 | +import pytest |
| 4 | +import zarr |
| 5 | + |
| 6 | +from arrayroute import ArrayReader |
| 7 | +from arrayroute.backends import ZarristaBackend |
| 8 | + |
| 9 | +zarrista = pytest.importorskip('zarrista') |
| 10 | + |
| 11 | + |
| 12 | +@pytest.mark.parametrize('sharded', [False, True]) |
| 13 | +def test_strides_match_bits_across_partial_shards_and_missing_chunks(tmp_path, sharded): |
| 14 | + rng = np.random.default_rng(289) |
| 15 | + values = rng.integers(0, 2**32, size=(9, 13, 7), dtype=np.uint32).view(np.float32) |
| 16 | + path = tmp_path / 'field.zarr' |
| 17 | + options = {'shards': (6, 8, 4)} if sharded else {} |
| 18 | + stored = zarr.create_array(path, shape=values.shape, dtype='f4', chunks=(3, 4, 2), |
| 19 | + fill_value=-7, **options) |
| 20 | + stored[:6] = values[:6] |
| 21 | + values[6:] = -7 |
| 22 | + backend = ZarristaBackend(zarrista.Array.open(zarrista.store.FilesystemStore(str(path)))) |
| 23 | + for budget in [400, 1 << 20]: |
| 24 | + reader = ArrayReader(backend, buffer_bytes=budget) |
| 25 | + for selection in [ |
| 26 | + (slice(1, 9, 2), slice(2, 13, 3), 3), |
| 27 | + (7, slice(None, None, 2), slice(1, 7, 3)), |
| 28 | + (slice(None, None, 4), 0, slice(None, None, 2)), |
| 29 | + (slice(0, 0, 2), slice(None), slice(None)), |
| 30 | + ]: |
| 31 | + assert reader.read(selection).tobytes() == values[selection].tobytes() |
| 32 | + |
| 33 | + |
| 34 | +def test_shard_index_does_not_outlive_read_or_accumulate(tmp_path): |
| 35 | + path = tmp_path / 'field.zarr' |
| 36 | + values = np.arange(16 * 24, dtype='f4').reshape(16, 24) |
| 37 | + stored = zarr.create_array(path, data=values, chunks=(2, 3), shards=(4, 6)) |
| 38 | + native = zarrista.Array.open(zarrista.store.FilesystemStore(str(path))) |
| 39 | + caches, sizes = [], [] |
| 40 | + |
| 41 | + class Observed: |
| 42 | + def __getattr__(self, name): |
| 43 | + return getattr(native, name) |
| 44 | + |
| 45 | + def shard_cache(self): |
| 46 | + cache = native.shard_cache() |
| 47 | + caches.append(cache) |
| 48 | + return cache |
| 49 | + |
| 50 | + def retrieve_subchunk(self, coordinates, *, shard_cache): |
| 51 | + value = native.retrieve_subchunk(coordinates, shard_cache=shard_cache) |
| 52 | + sizes.append(shard_cache.size()) |
| 53 | + return value |
| 54 | + |
| 55 | + reader = ArrayReader(ZarristaBackend(Observed()), buffer_bytes=1 << 20) |
| 56 | + selection = (slice(1, 16, 2), slice(0, 24, 2)) |
| 57 | + assert reader.read(selection).tobytes() == values[selection].tobytes() |
| 58 | + stored[:] = values + 7 |
| 59 | + assert reader.read(selection).tobytes() == (values + 7)[selection].tobytes() |
| 60 | + assert len(caches) == 2 and caches[0] is not caches[1] |
| 61 | + assert sizes and max(sizes) == 1 |
| 62 | + |
| 63 | + |
| 64 | +def test_outer_shard_codec_strides_use_whole_shard_budget(tmp_path): |
| 65 | + from zarr.codecs import BytesCodec, ShardingCodec, ZstdCodec |
| 66 | + |
| 67 | + path = tmp_path / 'field.zarr' |
| 68 | + values = np.arange(32 * 32, dtype='f4').reshape(32, 32) |
| 69 | + with pytest.warns(UserWarning, match='disables partial'): |
| 70 | + stored = zarr.create_array(path, shape=values.shape, chunks=(16, 16), dtype='f4', |
| 71 | + serializer=ShardingCodec(chunk_shape=(8, 8), codecs=[BytesCodec()]), |
| 72 | + compressors=[ZstdCodec()]) |
| 73 | + stored[:] = values |
| 74 | + backend = ZarristaBackend(zarrista.Array.open(zarrista.store.FilesystemStore(str(path)))) |
| 75 | + selection = (slice(1, 32, 3), slice(2, 32, 4)) |
| 76 | + reader = ArrayReader(backend, buffer_bytes=2048) |
| 77 | + assert reader.read(selection).tobytes() == values[selection].tobytes() |
0 commit comments