chore: add pytest benchmarks

.gitignore

@@ -114,3 +114,4 @@ dev-docs/brainstorms/
 dev-docs/plans/
 
 .worktrees/
+.benchmarks/

pyproject.toml

@@ -71,6 +71,7 @@ dev = [
     "netcdf4>=1.7.4",
     "pre-commit>=4.6.0",
     "pytest-asyncio>=0.24.0",
+    "pytest-benchmark>=5.1",
     "pytest-cov>=7.1.0",
     "pytest>=9.0.3",
     "pytest-xdist",

tests/benchmarks/README.md

@@ -0,0 +1,31 @@
+# Reader benchmarks
+
+In-process `XarrayReader` timings (pytest-benchmark) on a generated zarr
+store, with the store optionally wrapped in zarr's `LatencyStore` so each
+`get` costs a fixed round trip, as it would from S3. Every case also
+records the number of store `get` calls one operation makes
+(`extra_info.store_gets`), which is exact and branch-comparable.
+
+```sh
+uv run pytest tests/benchmarks --benchmark-only --no-cov -q \
+    --benchmark-json=bench.json --benchmark-columns=median,min,rounds
+```
+
+## Reporting a PR's performance
+
+Same machine, same interpreter, nothing else running. Build `main` in a
+worktree so the branch stays checked out:
+
+```sh
+git worktree add ../titiler-multidim-main main
+(cd ../titiler-multidim-main && uv run pytest tests/benchmarks --benchmark-only --no-cov -q --benchmark-json=../before.json)
+uv run pytest tests/benchmarks --benchmark-only --no-cov -q --benchmark-json=../after.json
+uv run pytest-benchmark compare ../before.json ../after.json --columns=median,min,rounds --sort=name
+```
+
+Paste that table into the PR's **Performance** section, together with the
+diff of `tests/test_perf_counts.py` if the PR changed any pinned count
+(opens, storage requests or boto3 sessions per request). Medians under
+~20 ms move by a few percent between runs; call a change real only when it
+is well outside that. Do not pass `-n`: pytest-benchmark refuses
+`--benchmark-only` under xdist.

tests/benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Reader benchmarks (run with --benchmark-only)."""

tests/benchmarks/conftest.py

@@ -0,0 +1,19 @@
+"""Collect benchmarks only when pytest-benchmark is explicitly enabled.
+
+`pytest` alone runs the normal suite; `pytest tests/benchmarks
+--benchmark-only` runs these (see README in this directory).
+"""
+
+import pytest
+
+
+def pytest_collection_modifyitems(config, items):
+    """Skip benchmark tests unless --benchmark-only / --benchmark-enable."""
+    if config.getoption("benchmark_only", False) or config.getoption(
+        "benchmark_enable", False
+    ):
+        return
+    skip = pytest.mark.skip(reason="run with --benchmark-only")
+    for item in items:
+        if "benchmarks" in str(item.fspath):
+            item.add_marker(skip)

tests/benchmarks/test_app_bench.py

@@ -0,0 +1,39 @@
+"""App-level benchmarks on a 17,228-step Icechunk store (TEMPO's time axis length).
+
+Each case is one request through `TestClient`, so it includes dependency
+parsing, the mosaic backend, the reader and rendering. Compare runs with
+`pytest-benchmark compare` (see README).
+"""
+
+import pytest
+
+TIME = "2023-08-01T00:00:00"
+
+
+@pytest.mark.parametrize(
+    ("path", "params"),
+    [
+        ("/info", {"sel": f"time={TIME}"}),
+        # rio-tiler info() builds band metadata for every time step: Issue 4
+        ("/info", {}),
+        # plus the times list (cheap since #194; kept so a regression shows)
+        ("/info", {"show_times": True}),
+        ("/tiles/WebMercatorQuad/2/0/1.png", {"sel": f"time={TIME}"}),
+        ("/point/-95,35", {"sel": f"time={TIME}"}),
+    ],
+    ids=["info_sel", "info_no_sel", "info_show_times", "tile_z2", "point"],
+)
+def test_request(benchmark, app, long_icechunk_store, path, params):
+    """One request, cold: every round opens the store again (as on main).
+
+    The data chunks are inlined in the manifest (see the fixture), so the
+    tile and point cases measure the open and request path, not chunk
+    fetches; the chunk counts live in tests/test_perf_counts.py.
+    """
+    query = {"url": long_icechunk_store, "variable": "data", **params}
+
+    def run():
+        response = app.get(path, params=query)
+        assert response.status_code == 200
+
+    benchmark(run)

tests/benchmarks/test_reader_bench.py

@@ -0,0 +1,167 @@
+"""XarrayReader benchmarks: open, tile and point.
+
+Cases are parametrized by store latency (0 and 30 ms per `get`) so the
+difference between "reads the window" and "reads the whole slice", and
+between one batched round trip and several, shows up as time and not
+just as a `get` count. Values compared across branches must come from
+the same machine, interpreter and store layout.
+"""
+
+import numpy as np
+import pytest
+import xarray as xr
+from zarr.abc.store import Store
+from zarr.storage import LocalStore, WrapperStore
+from zarr.testing.store import LatencyStore
+
+from titiler.multidim import mosaic, reader
+
+LATENCIES = (0.0, 0.03)
+# WebMercatorQuad z/x/y inside the fixture's bounds
+TILE = (3, 2, 3)
+POINT = (-95.0, 35.0)
+
+
+class CountingStore(WrapperStore[Store]):
+    """Count `get` calls on the wrapped store.
+
+    Wrap a read-only store: zarr then uses this instance as-is rather than
+    deriving a read-only copy, with its own count, through `_with_store`.
+    """
+
+    gets = 0
+
+    async def get(self, key, prototype, byte_range=None):
+        """Count, then delegate."""
+        self.gets += 1
+        return await self._store.get(key, prototype, byte_range=byte_range)
+
+    def get_sync(self, key, prototype, byte_range=None):
+        """Count, then delegate (zarr >= 3.2 sync read path)."""
+        self.gets += 1
+        return self._store.get_sync(key, prototype=prototype, byte_range=byte_range)
+
+
+@pytest.fixture(scope="session")
+def store_path(tmp_path_factory) -> str:
+    """4 x 720 x 1440 float32 data, 2-D mask and flag, chunked 1 x 180 x 180."""
+    rng = np.random.default_rng(0)
+    ny, nx, nt = 720, 1440, 4
+    flag = rng.integers(0, 3, (ny, nx)).astype("float32")
+    flag[::7, ::11] = np.nan  # fill pixels
+    ds = xr.Dataset(
+        {
+            "data": (("time", "lat", "lon"), rng.random((nt, ny, nx), dtype="float32")),
+            "mask2d": (("lat", "lon"), rng.random((ny, nx), dtype="float32")),
+            "flag": (("lat", "lon"), flag),
+        },
+        coords={
+            "time": np.arange(nt),
+            "lat": np.linspace(-89.875, 89.875, ny),
+            "lon": np.linspace(-179.875, 179.875, nx),
+        },
+    )
+    path = str(tmp_path_factory.mktemp("bench") / "store.zarr")
+    ds.to_zarr(
+        path,
+        consolidated=False,
+        encoding={
+            "data": {"chunks": (1, 180, 180)},
+            "mask2d": {"chunks": (180, 180)},
+            "flag": {"chunks": (180, 180)},
+        },
+    )
+    return path
+
+
+@pytest.fixture(params=LATENCIES, ids=lambda v: f"latency={int(v * 1000)}ms")
+def latency(request) -> float:
+    """Per-`get` store latency in seconds."""
+    return request.param
+
+
+@pytest.fixture
+def counter(store_path, latency, monkeypatch) -> CountingStore:
+    """Route the reader's opener through counting + latency store wrappers."""
+    counting = CountingStore(LocalStore(store_path, read_only=True))
+    store: Store = LatencyStore(counting, get_latency=latency) if latency else counting
+
+    def opener(src_path: str, **kwargs) -> xr.Dataset:
+        return xr.open_zarr(store, chunks=None, consolidated=False, decode_coords="all")
+
+    monkeypatch.setattr(reader, "guess_opener", opener)
+    return counting
+
+
+def _reader():
+    return reader.XarrayReader(
+        src_path="bench://store", variable="data", sel=["time=0"]
+    )
+
+
+def _record(benchmark, counter, fn):
+    """Time `fn`; record the `get` count of one call as extra_info."""
+    counter.gets = 0
+    fn()
+    assert counter.gets, "gets went to a copy of the counting store"
+    benchmark.extra_info["store_gets"] = counter.gets
+    benchmark(fn)
+
+
+def test_open(benchmark, counter):
+    """Reader construction: open and extract the variable (metadata only)."""
+
+    def run():
+        with _reader():
+            pass
+
+    _record(benchmark, counter, run)
+
+
+def test_tile(benchmark, counter):
+    """Open + one z3 tile: the request path a map viewer exercises."""
+
+    def run():
+        with _reader() as src:
+            return src.tile(TILE[1], TILE[2], TILE[0])
+
+    _record(benchmark, counter, run)
+
+
+def test_point(benchmark, counter):
+    """Open + one point read."""
+
+    def run():
+        with _reader() as src:
+            return src.point(*POINT)
+
+    _record(benchmark, counter, run)
+
+
+def _backend():
+    """Single-URL request path: the backend validates, then rio-tiler reads."""
+    return mosaic.XarrayMosaicBackend(
+        ["bench://store"],
+        reader=reader.XarrayReader,
+        reader_options={"variable": "data", "sel": ["time=0"]},
+    )
+
+
+def test_backend_tile(benchmark, counter):
+    """Open through XarrayMosaicBackend + one z3 tile, as /tiles does."""
+
+    def run():
+        with _backend() as src:
+            return src.tile(TILE[1], TILE[2], TILE[0])
+
+    _record(benchmark, counter, run)
+
+
+def test_backend_point(benchmark, counter):
+    """Open through XarrayMosaicBackend + one point read, as /point does."""
+
+    def run():
+        with _backend() as src:
+            return src.point(*POINT)
+
+    _record(benchmark, counter, run)

tests/conftest.py

@@ -30,3 +30,53 @@ def app(monkeypatch):
 
     with TestClient(app) as client:
         yield client
+
+
+@pytest.fixture
+def storage_requests(monkeypatch):
+    """Icechunk storage requests made by the test's local repositories, by kind."""
+    from helpers import serve_counting
+
+    server, counts = serve_counting(monkeypatch)
+    yield counts
+    server.shutdown()
+    server.server_close()
+
+
+@pytest.fixture(scope="session")
+def long_icechunk_store(tmp_path_factory) -> str:
+    """Icechunk store with TEMPO's time axis length (17,228 steps) on an 8x8 grid.
+
+    Data is chunked coarsely: one chunk per step takes ~8 s to write and 17k
+    files, and nothing here depends on it. The `time` coordinate is the only
+    chunk object in the store (~18 KB stored): the all-zero data chunks
+    compress below icechunk's inline threshold and live in the manifest, so
+    tile and point requests on this store never fetch a data chunk.
+    """
+    import icechunk
+    import numpy as np
+    import xarray as xr
+
+    nt, ny, nx = 17228, 8, 8
+    ds = xr.Dataset(
+        {
+            "data": (
+                ("time", "latitude", "longitude"),
+                np.zeros((nt, ny, nx), "float32"),
+            )
+        },
+        coords={
+            "time": np.datetime64("2023-08-01T00:00:00", "ns")
+            + np.arange(nt) * np.timedelta64(40, "m"),
+            "latitude": np.linspace(50, 15, ny),
+            "longitude": np.linspace(-125, -65, nx),
+        },
+    )
+    path = str(tmp_path_factory.mktemp("icechunk") / "long")
+    repo = icechunk.Repository.create(icechunk.local_filesystem_storage(path))
+    session = repo.writable_session("main")
+    ds.to_zarr(
+        session.store, consolidated=False, encoding={"data": {"chunks": (1000, ny, nx)}}
+    )
+    session.commit("17,228 time steps")
+    return path