zarr-python/tests/benchmarks/test_e2e.py at b81844a81f2dc3d03891ebefdd4f5d193fb8fbbf · abishop1990/zarr-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
"""
Benchmarks for end-to-end read/write performance of Zarr
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from tests.benchmarks.common import Layout

if TYPE_CHECKING:
    from pytest_benchmark.fixture import BenchmarkFixture

    from zarr.abc.store import Store
    from zarr.core.common import NamedConfig

from operator import getitem, setitem
from typing import Any, Literal

import pytest

from zarr import create_array

CompressorName = Literal["gzip"] | None
ZarrFormat = Literal[2, 3]

compressors: dict[CompressorName, NamedConfig[Any, Any] | None] = {
    None: None,
    "gzip": {"name": "gzip", "configuration": {"level": 1}},
}

layouts: tuple[Layout, ...] = (
    # No shards, just 1000 chunks
    Layout(shape=(1_000_000,), chunks=(1000,), shards=None),
    # 1:1 chunk:shard shape, should measure overhead of sharding
    Layout(shape=(1_000_000,), chunks=(1000,), shards=(1000,)),
    # One shard with all the chunks, should measure overhead of handling inner shard chunks
    Layout(shape=(1_000_000,), chunks=(100,), shards=(10000 * 100,)),
)


@pytest.mark.parametrize("zarr_format", [2, 3])
@pytest.mark.parametrize("compression_name", [None, "gzip"])
@pytest.mark.parametrize("layout", layouts, ids=str)
@pytest.mark.parametrize("store", ["memory", "local"], indirect=["store"])
def test_write_array(
    store: Store,
    layout: Layout,
    compression_name: CompressorName,
    zarr_format: ZarrFormat,
    benchmark: BenchmarkFixture,
) -> None:
    """
    Test the time required to fill an array with a single value.
    Parametrized to benchmark both Zarr v2 and v3 formats.
    """
    # Skip v2 tests with sharding (not supported in v2)
    if zarr_format == 2 and layout.shards is not None:
        pytest.skip("Sharding not supported in Zarr v2")

    arr = create_array(
        store,
        zarr_format=zarr_format,
        dtype="uint8",
        shape=layout.shape,
        chunks=layout.chunks,
        shards=layout.shards,
        compressors=compressors[compression_name],  # type: ignore[arg-type]
        fill_value=0,
    )

    benchmark(setitem, arr, Ellipsis, 1)


@pytest.mark.parametrize("zarr_format", [2, 3])
@pytest.mark.parametrize("compression_name", [None, "gzip"])
@pytest.mark.parametrize("layout", layouts, ids=str)
@pytest.mark.parametrize("store", ["memory", "local"], indirect=["store"])
def test_read_array(
    store: Store,
    layout: Layout,
    compression_name: CompressorName,
    zarr_format: ZarrFormat,
    benchmark: BenchmarkFixture,
) -> None:
    """
    Test the time required to read an array.
    Parametrized to benchmark both Zarr v2 and v3 formats.
    """
    # Skip v2 tests with sharding (not supported in v2)
    if zarr_format == 2 and layout.shards is not None:
        pytest.skip("Sharding not supported in Zarr v2")

    arr = create_array(
        store,
        zarr_format=zarr_format,
        dtype="uint8",
        shape=layout.shape,
        chunks=layout.chunks,
        shards=layout.shards,
        compressors=compressors[compression_name],  # type: ignore[arg-type]
        fill_value=0,
    )
    arr[:] = 1
    benchmark(getitem, arr, Ellipsis)