|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +from typing import TYPE_CHECKING |
| 4 | + |
| 5 | +import mmap |
| 6 | +import pytest |
| 7 | + |
| 8 | +import zarr |
| 9 | +from zarr.core.buffer import Buffer, cpu |
| 10 | +from zarr.storage import LocalStore |
| 11 | +from zarr.testing.store import StoreTests |
| 12 | + |
| 13 | +if TYPE_CHECKING: |
| 14 | + import pathlib |
| 15 | + |
| 16 | + |
| 17 | +class MemoryMappedDirectoryStore(LocalStore): |
| 18 | + def _fromfile(self, fn): |
| 19 | + with open(fn, "rb") as fh: |
| 20 | + return memoryview(mmap.mmap(fh.fileno(), 0, prot=mmap.PROT_READ)) |
| 21 | + |
| 22 | + |
| 23 | +class TestMemoryMappedDirectoryStore(StoreTests[MemoryMappedDirectoryStore, cpu.Buffer]): |
| 24 | + store_cls = MemoryMappedDirectoryStore |
| 25 | + buffer_cls = cpu.Buffer |
| 26 | + |
| 27 | + async def get(self, store: MemoryMappedDirectoryStore, key: str) -> Buffer: |
| 28 | + return self.buffer_cls.from_bytes((store.root / key).read_bytes()) |
| 29 | + |
| 30 | + async def set(self, store: MemoryMappedDirectoryStore, key: str, value: Buffer) -> None: |
| 31 | + parent = (store.root / key).parent |
| 32 | + if not parent.exists(): |
| 33 | + parent.mkdir(parents=True) |
| 34 | + (store.root / key).write_bytes(value.to_bytes()) |
| 35 | + |
| 36 | + @pytest.fixture |
| 37 | + def store_kwargs(self, tmpdir) -> dict[str, str]: |
| 38 | + return {"root": str(tmpdir)} |
| 39 | + |
| 40 | + def test_store_repr(self, store: MemoryMappedDirectoryStore) -> None: |
| 41 | + assert str(store) == f"file://{store.root.as_posix()}" |
| 42 | + |
| 43 | + def test_store_supports_writes(self, store: MemoryMappedDirectoryStore) -> None: |
| 44 | + assert store.supports_writes |
| 45 | + |
| 46 | + def test_store_supports_partial_writes(self, store: MemoryMappedDirectoryStore) -> None: |
| 47 | + assert store.supports_partial_writes |
| 48 | + |
| 49 | + def test_store_supports_listing(self, store: MemoryMappedDirectoryStore) -> None: |
| 50 | + assert store.supports_listing |
| 51 | + |
| 52 | + async def test_empty_with_empty_subdir(self, store: MemoryMappedDirectoryStore) -> None: |
| 53 | + assert await store.is_empty("") |
| 54 | + (store.root / "foo/bar").mkdir(parents=True) |
| 55 | + assert await store.is_empty("") |
| 56 | + |
| 57 | + def test_creates_new_directory(self, tmp_path: pathlib.Path): |
| 58 | + target = tmp_path.joinpath("a", "b", "c") |
| 59 | + assert not target.exists() |
| 60 | + |
| 61 | + store = self.store_cls(root=target) |
| 62 | + zarr.group(store=store) |
| 63 | + |
| 64 | + async def test_mmap_slice_reads(self, store: MemoryMappedDirectoryStore) -> None: |
| 65 | + """Test reading slices with memory mapping""" |
| 66 | + # Create array with large chunks |
| 67 | + z = zarr.create_array(store=store, shape=(2000, 2000), chunks=(1000, 1000), |
| 68 | + dtype='float64') |
| 69 | + # Write test data |
| 70 | + data = zarr.full(shape=(2000, 2000), chunks=(1000, 1000), fill_value=42.0, |
| 71 | + dtype='float64') |
| 72 | + z[:] = data[:] |
| 73 | + |
| 74 | + # Test reading various slices |
| 75 | + slices = [ |
| 76 | + # Within single chunk |
| 77 | + (slice(100, 200), slice(100, 200)), |
| 78 | + # Across chunk boundaries |
| 79 | + (slice(900, 1100), slice(900, 1100)), |
| 80 | + # Full chunk |
| 81 | + (slice(0, 1000), slice(0, 1000)) |
| 82 | + ] |
| 83 | + |
| 84 | + for test_slice in slices: |
| 85 | + assert (z[test_slice] == data[test_slice]).all() |
0 commit comments