quantlab.backend.zarr.rechunk_index_coordinates

quantlab.backend.zarr.rechunk_index_coordinates(path)

Rewrite each 1-D index coordinate of a store in one chunk, in place.

Zarr fixes an array’s chunk grid when the array is created, and xarray writes a dimension coordinate as one chunk of the creating write’s length. A store created from a short window and grown by appends therefore holds timestamp in many small chunks (one per bar when the first window was one bar), and every open of the store reads each of them. This function rewrites every such coordinate (a 1-D array whose only dimension is its own name, such as timestamp or symbol) as one chunk spanning the whole axis. XrBackend.append calls it after every append, and it is the maintenance step for a store written before that.

Only the chunk grid changes. The stored values are copied as they are encoded on disk (the int64 ticks of a datetime coordinate, the text of a string one), so no value, dtype, codec, attribute or data variable is touched and a data fingerprint of the store is unchanged. A store with no fragmented coordinate is not written at all.

Each rewritten array is built and read back in a .rechunk.tmp sidecar beside the store, then swapped in by two renames: the original array to a .replaced.tmp sidecar, the rewritten one into its place. Consolidated metadata, when the store has it, is rewritten after every swap, and only then are the originals removed. A process killed at any point leaves sidecars that tell how far it got; the next call of this function, or of any XrBackend method that writes to the store, resolves them first: before the first rename the store is intact and the sidecar is dropped, between the renames the original array is renamed back, after them the swap is finished. Nothing beside the store is ever needed to read it, except in two windows a reader can fall into: between the two renames the store has no array for that coordinate, and between the swap and the consolidation its consolidated metadata records the old chunk grid. Both last microseconds to milliseconds when the process runs to the end; after a kill inside one, the store stays unreadable until the next writing call recovers it. Do not run this while another process reads or writes the store.

Parameters:

pathstr or os.PathLike

Directory of the Zarr store (format 2 or 3).

Returns:

list[str]

Names of the coordinates rewritten, sorted; empty when every index coordinate was already one chunk.

Raises:

FileNotFoundError

If path does not exist.

RuntimeError

If a rewritten coordinate does not read back equal to the original; the store is then left as it was.

Examples

A store created from a one-bar window and grown by an append holds timestamp in one-element chunks:

>>> import tempfile
>>> from pathlib import Path
>>> import numpy as np
>>> import pandas as pd
>>> import xarray as xr
>>> import zarr
>>> from quantlab.backend.zarr import rechunk_index_coordinates
>>> panel = xr.Dataset(
...     {"close": (("timestamp", "symbol"), np.ones((3, 2)))},
...     coords={"timestamp": pd.date_range("2024-01-02", periods=3),
...             "symbol": ["AAA", "BBB"]},
... )
>>> path = Path(tempfile.mkdtemp()) / "prices.zarr"
>>> _ = panel.isel(timestamp=[0]).to_zarr(path, mode="w")
>>> _ = panel.isel(timestamp=[1, 2]).to_zarr(
...     path, mode="a", append_dim="timestamp"
... )
>>> zarr.open_group(path, mode="r")["timestamp"].chunks
(1,)
>>> rechunk_index_coordinates(path)
['timestamp']
>>> zarr.open_group(path, mode="r")["timestamp"].chunks
(3,)
>>> xr.open_zarr(path)["timestamp"].dt.day.values.tolist()
[2, 3, 4]
>>> rechunk_index_coordinates(path)
[]