quantlab.backend.zarr.XrBackend#

class quantlab.backend.zarr.XrBackend[source]#

Bases: DataBackend

Zarr-backed storage for an xarray.Dataset panel.

read and write move the whole panel between memory and a Zarr directory. append extends an existing store along one dimension (timestamp by default) after checking that the incoming window cannot silently corrupt it. widen_symbol_axis, widen_data_vars and widen_and_append handle the case where the panel written today has more symbols or more variables than the store already holds.

Two class constants govern the on-disk layout: APPEND_DIM_CHUNK is the chunk length pinned along the append dimension when a store is created, and MAX_WIDEN_BYTES is the largest panel a widen will hold in memory at once before switching to a block-by-block rewrite.

Examples

>>> backend = XrBackend().to_internal(panel)
>>> backend.write("prices.zarr")
>>> backend.to_internal(next_month).widen_and_append("prices.zarr")
>>> ds = XrBackend().read("prices.zarr").get_xarray_dataset(
...     ["timestamp", "symbol"]
... )

APPEND_DIM_CHUNK = 512#: Chunk length pinned along the append dimension when a store is created by append, widen_symbol_axis or widen_data_vars. Without an explicit encoding Zarr would adopt the first window’s own length as the chunk size, and later windows of a different length would be misaligned with the on-disk grid. The grid a store ends up with is min(APPEND_DIM_CHUNK, extent), where extent is the panel’s length along the append dimension or, when the caller states it with append_dim_size, the store’s eventual total length. Zarr fixes the grid at creation; an append cannot revise it, and a mode="w" rewrite re-pins it.

MAX_WIDEN_BYTES = 4294967296#: Largest number of bytes a symbol-axis widen may materialise at once. At or under this budget widen_symbol_axis reindexes the whole store in memory and writes it once; above it, the store is rewritten block by block along the append dimension. The value routes between the two strategies and never refuses a widen. It is a constant of this module rather than an import because the storage layer must not depend on the acquisition layer for a scalar.

SYMBOL_DIM_CHUNK = 1024#: Chunk width pinned along the symbol axis when a store is created by append, widen_and_append or a widen rewrite. A store’s symbol axis grows at its end (extend_symbol_axis, ADR 0032), so a width taken from the first window would leave a store that started with a few symbols in thousands of narrow chunks. Stores created before the constant keep their own width; Zarr fixes it at creation.

WIDENING_SUFFIX = '.widening.tmp'#: Sidecar suffixes used by widen_symbol_axis’s atomic directory swap. .widening.tmp holds the rewritten store before it becomes authoritative and is never read by anything. .superseded.tmp briefly holds the original store between the two renames and is the one artefact a crash can leave behind that still contains real data.

append(path, append_dim='timestamp', *, append_dim_size=None, **kwargs)[source]#

Create the store at path, or extend it along append_dim.

When no store exists, data is written as the first window with the chunk grid pinned by _append_encoding. Otherwise the window is appended after _assert_append_compatible has checked that it cannot silently corrupt the store. Raw to_zarr(mode="a") performs none of those checks: a changed coordinate overwrites the stored labels, a float NaN written into an integer variable becomes 0, an overlapping window leaves the axis non-monotonic, and a changed variable set leaves the store unopenable. A gap between the stored end and the incoming start is allowed; overlap is refused and has no opt-out, because append extends a store and does not recompute one. To replace a range the store already holds, rewrite the store with write.

After the append the store’s 1-D index coordinates are rewritten in one chunk by rechunk_index_coordinates: Zarr fixes a coordinate’s chunk grid at creation, so without it a store created from a one-bar window would keep timestamp in one chunk per bar.

A panel that has grown a variable must go through widen_data_vars or widen_and_append first; a panel missing a stored variable is refused outright, since the only fill would be NaN over the incoming dates of a variable that was complete.

Parameters:

pathstr: Directory of the Zarr store.
append_dimstr, default “timestamp”: The dimension the store grows along.
append_dim_sizeint, optional: The store’s eventual total length along append_dim, if the caller knows it. Only the creating write reads it; it pins the chunk grid to the value a single whole-range write would have chosen instead of the first window’s length. Ignored against an existing store, so an incremental writer may pass it on every window.
**kwargs: Passed through to Dataset.to_zarr. An encoding entry is dropped on the append path because xarray rejects it there.

Raises:

ValueError: If the window fails any compatibility check.

Examples

Two windows on the same symbol axis; the first call creates the store, the second extends it:

>>> XrBackend().to_internal(first_window).append("prices.zarr")
XrBackend()
>>> XrBackend().to_internal(next_window).append("prices.zarr")
XrBackend()
>>> xr.open_zarr("prices.zarr").sizes["timestamp"]
6

create(path, fill_values=None, *, dim='symbol', append_dim='timestamp', **kwargs)[source]#

Write data as a new store at path, laid out to grow.

The dim axis is put in sort_symbol_axis order (sort_symbols) and the chunk grid and Zarr fill values are pinned as append and widen_and_append pin them (_append_encoding): APPEND_DIM_CHUNK along append_dim, SYMBOL_DIM_CHUNK along dim, and the fill_values entry as the fill of each variable on dim, so a symbol appended later (extend_symbol_axis) reads as it without being written. Any store already at path is replaced, as by write.

Parameters:

pathstr: Directory of the Zarr store.
fill_valuesMapping[str, object], optional: Per-variable fill of the cells a later symbol extension adds.
dimstr, default “symbol”: The symbol axis.
append_dimstr, default “timestamp”: The dimension the store grows along.
**kwargs: Passed through to write; an encoding given here wins.

Examples

>>> 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 XrBackend
>>> panel = xr.Dataset(
...     {"close": (("timestamp", "symbol"), np.ones((2, 2)))},
...     coords={"timestamp": pd.date_range("2024-01-02", periods=2),
...             "symbol": np.array(["BBB", "AAA"], dtype=object)},
... )
>>> path = str(Path(tempfile.mkdtemp()) / "prices.zarr")
>>> _ = XrBackend().to_internal(panel).create(path)
>>> xr.open_zarr(path)["symbol"].values.tolist()
['AAA', 'BBB']
>>> zarr.open_group(path)["close"].chunks
(2, 1024)

extend_symbol_axis(path, symbols, dim='symbol', fill_values=None)[source]#

Append symbols at the end of the store’s dim axis, in place.

Each Zarr array carrying dim is resized; no stored chunk is rewritten, so the cost is a few metadata writes whatever the store’s size. The new cells read as each array’s fill value: NaN for a floating-point variable, the fill_values entry otherwise, an empty string for a text coordinate on dim. A variable whose Zarr fill value is not that (a store created before _append_encoding pinned it, a datetime variable) has its new cells written, which rewrites at most the chunks the old axis ends in. The coordinate is then swapped for one holding the old labels followed by symbols, in one chunk, and the metadata consolidated; arrays without dim are untouched. data is not touched.

Labels the store already holds are skipped; the others keep the order given (widen_and_append sorts them among themselves). They are re-spelled in the stored axis’s dtype first, so a ticker handed to an integer axis raises. A fixed-width text axis is stored at the width of its longest label.

The swap of the coordinate is the commit: a process killed before it leaves arrays longer than the coordinate, which the next writing call cuts back (_recover_axis_extension), so the store is as it was before. Do not run this while another process reads the store.

Parameters:

pathstr: Directory of the Zarr store.
symbolsSequence: The labels to append.
dimstr, default “symbol”: The axis to extend.
fill_valuesMapping[str, object], optional: Per-variable fill of the new cells. A variable on dim that is not floating-point needs one, as in widen_symbol_axis.

Raises:

FileNotFoundError: If no store exists at path.
ValueError: If a non-floating variable on dim has no fill value.

Examples

>>> import tempfile
>>> from pathlib import Path
>>> import numpy as np
>>> import pandas as pd
>>> import xarray as xr
>>> from quantlab.backend.zarr import XrBackend
>>> panel = xr.Dataset(
...     {"close": (("timestamp", "symbol"), np.ones((2, 2))),
...      "volume": (("timestamp", "symbol"), np.ones((2, 2)))},
...     coords={"timestamp": pd.date_range("2024-01-02", periods=2),
...             "symbol": np.array(["BBB", "DDD"], dtype=object)},
... )
>>> path = str(Path(tempfile.mkdtemp()) / "prices.zarr")
>>> _ = XrBackend().to_internal(panel).append(path)
>>> _ = XrBackend().extend_symbol_axis(
...     path, ["AAA"], fill_values={"volume": 0.0}
... )
>>> stored = xr.open_zarr(path)
>>> stored["symbol"].values.tolist()
['BBB', 'DDD', 'AAA']
>>> stored["close"].sel(symbol="AAA").values.tolist()
[nan, nan]
>>> stored["volume"].sel(symbol="AAA").values.tolist()
[0.0, 0.0]

filter_by_date(col, start_date, end_date)[source]#

Narrow data in place to the label slice start_date..end_date.

Parameters:

colstr: The date dimension to slice.
start_datestr: First label to keep, inclusive.
end_datestr: Last label to keep, inclusive.

Examples

>>> backend.filter_by_date("timestamp", "2024-01-02", "2024-01-03")
XrBackend()
>>> backend.data["timestamp"].values.astype("datetime64[D]")
array(['2024-01-02', '2024-01-03'], dtype='datetime64[D]')

filter_by_symbol(col, symbols)[source]#

Narrow data in place to the given labels on col.

Parameters:

colstr: The dimension to select on.
symbolstuple[str, …]: The labels to keep.

Examples

>>> backend.filter_by_symbol("symbol", ("BBB",))
XrBackend()
>>> backend.data["symbol"].values.tolist()
['BBB']

get_lazyframe()[source]#

Return data as a long-format polars.LazyFrame.

Examples

>>> backend.get_lazyframe().collect().columns
['timestamp', 'symbol', 'close']

get_xarray_dataset(indexes=None)[source]#

Return data indexed by exactly the dimensions in indexes.

indexes names the result’s dimensions, in order: data variables laid out on any other dimension are dropped, dimensions no longer used are dropped with their coordinates, and the survivors are transposed onto the requested order. This is the same meaning PlBackend.get_xarray_dataset gives the argument. None returns data itself, not a copy, since callers rely on receiving the object the backend holds. data is never modified here; the filter_by_* methods are the in-place ones.

Parameters:

indexeslist[str], optional: The dimensions to index by, or None for no shape request.

Raises:

ValueError: If a requested name is not a dimension of data.

Examples

>>> ds = backend.get_xarray_dataset(["timestamp", "symbol"])
>>> tuple(ds.dims)
('timestamp', 'symbol')
>>> backend.get_xarray_dataset(["symbol", "timestamp"])["close"].dims
('symbol', 'timestamp')
>>> backend.get_xarray_dataset() is backend.data
True

head(path, n)[source]#

Return at most n rows of the store at path as a lazy frame.

Parameters:

pathstr: Directory of the Zarr store.
nint: Maximum number of rows (and of labels per dimension) to read.
The store is opened by path with the same opener ``read`` uses, and
every dimension is sliced to ``n`` before conversion, so the whole
store is never materialised. ``data`` is neither read nor written:
a probe must not observe or inherit the in-place narrowing that
``filter_by_date`` leaves behind.

Raises:

FileNotFoundError: If path does not exist.

Examples

>>> XrBackend().head("prices.zarr", 2).collect().shape
(2, 3)

read(path, **kwargs)[source]#

Open the Zarr store at path lazily into data.

Every call opens the store again and replaces whatever data held; nothing is cached across calls. Zarr reads a string dimension coordinate (such as symbol) back as numpy’s variable-width StringDType; it is handed back as object strings, the form a panel built from a frame has, so labels from a store and from memory compare, align and cast to str alike.

Parameters:

pathstr or os.PathLike: Directory of the Zarr store.
**kwargs: Passed through to xarray.open_dataset.

Raises:

FileNotFoundError: If path does not exist.

Examples

>>> import tempfile
>>> from pathlib import Path
>>> import numpy as np
>>> import pandas as pd
>>> import xarray as xr
>>> from quantlab.backend.zarr import XrBackend
>>> panel = xr.Dataset(
...     {"close": (("timestamp", "symbol"), np.ones((4, 2)))},
...     coords={"timestamp": pd.date_range("2024-01-02", periods=4),
...             "symbol": np.array(["AAA", "BBB"], dtype=object)},
... )
>>> path = Path(tempfile.mkdtemp()) / "prices.zarr"
>>> _ = XrBackend().to_internal(panel).write(str(path))
>>> backend = XrBackend().read(path)
>>> dict(backend.data.sizes)
{'timestamp': 4, 'symbol': 2}
>>> backend.data["symbol"].dtype
dtype('O')

resample(labels, how)[source]#

Aggregate data onto the bars labels assigns, in place.

Every variable with a timestamp dimension is grouped by its label and reduced with its method from how; the other dimensions are kept as they are. A variable without a timestamp dimension is carried over unchanged. count yields integers; the other methods keep pandas’ skip-NaN semantics, so an all-NaN bar is NaN under first, last, max, min and mean and 0 under sum.

Parameters:

labelspd.Series: Source timestamp to target timestamp; see DataBackend.
howdict[str, str]: Variable name to aggregation method, one entry per variable.

Raises:

ValueError: If a held timestamp has no label, or a variable has no method.

Examples

>>> labels = pd.Series(minutes.floor("D"), index=minutes)
>>> backend.resample(labels, {"close": "last", "volume": "sum"})
>>> backend.data["timestamp"].values.astype("datetime64[D]").tolist()
[datetime.date(2024, 1, 2), datetime.date(2024, 1, 3)]

shorten_axis(path, last, dim='timestamp')[source]#

Drop the labels of the store’s dim axis after last, in place.

A write killed midway is settled first (_recover_interrupted_writes), and the axis is then read without the consolidated metadata, so a shortening killed midway is finished by running it again; with no label after last nothing is written. The index coordinate of dim is cut first, then every other array carrying dim, and the metadata is consolidated again. No stored chunk is rewritten: the cut cells stay in the axis’s last chunk until the axis grows over them, and an append writes every cell it adds. Cutting the coordinate is the commit: a process killed after it leaves arrays longer than their axis, which the next writing call cuts back. The other axes are untouched. Do not run this while another process reads the store.

Parameters:

pathstr: Directory of the Zarr store.
lastlabel: The last label of dim kept (a pd.Timestamp on the time axis); the axis is ordered.
dimstr, default “timestamp”: The axis to shorten.

Raises:

FileNotFoundError: If no store exists at path.

Examples

>>> import tempfile
>>> from pathlib import Path
>>> import numpy as np
>>> import pandas as pd
>>> import xarray as xr
>>> from quantlab.backend.zarr import XrBackend
>>> panel = xr.Dataset(
...     {"close": (("timestamp", "symbol"), np.arange(6.0).reshape(3, 2))},
...     coords={"timestamp": pd.date_range("2024-01-02", periods=3),
...             "symbol": np.array(["AAA", "BBB"], dtype=object)},
... )
>>> path = str(Path(tempfile.mkdtemp()) / "prices.zarr")
>>> _ = XrBackend().to_internal(panel).append(path)
>>> _ = XrBackend().shorten_axis(path, pd.Timestamp("2024-01-03"))
>>> xr.open_zarr(path)["close"].values.tolist()
[[0.0, 1.0], [2.0, 3.0]]

sort_symbols(dim='symbol')[source]#

Put data’s dim axis in sort_symbol_axis order.

Every store keeps its symbol axis in that order, the one widen_and_append writes, so a store has one layout whether it was written whole or appended to. Labels and values move together; a panel without dim is left as it is.

Parameters:

dimstr, default “symbol”: The axis to order.

Returns:

Self: self, for chaining.

Examples

>>> backend = XrBackend().to_internal(panel).sort_symbols()
>>> backend.data["symbol"].values.tolist()
['AAA', 'BBB', 'CCC']

to_internal(data)[source]#

Adopt an in-memory xarray.Dataset as data.

Parameters:

dataxr.Dataset: The panel to hold.

Examples

>>> backend = XrBackend().to_internal(panel)
>>> backend.data is panel
True

widen_and_append(path, append_dim='timestamp', dim='symbol', fill_values=None, **kwargs)[source]#

Reconcile every axis with the store at path, then append.

The opt-in sibling of append for a panel that has grown. In order: the incoming labels the store lacks are appended at the end of its dim axis, sorted among themselves (extend_symbol_axis, which rewrites no stored chunk), and data is reindexed onto the stored axis followed by them; then any variable the store lacks is added with widen_data_vars; then the unchanged append runs. A window carrying the store’s labels in another order is only reindexed. Variables are widened after the symbol axis so the filler is built once at the final width. Calling this unconditionally is cheap: when both axes already agree it delegates straight to append, and when no store exists yet it puts the dim axis in sorted order (sort_symbols) and creates the store with its chunk grid and the Zarr fill values of fill_values pinned (_append_encoding).

The closing append is what keeps the compatibility checks in force. Because the widens commit before it runs, a window it then refuses (an overlapping append_dim range, say) can leave the store with the grown symbol axis and variable set while its append dimension and stored history are untouched. That is the accepted cost of inheriting the refusal rather than duplicating it.

append_dim_size may ride in kwargs. It is read, not consumed, so the same value reaches the symbol widen, the variable widen and the closing append; those three cannot disagree on the chunk grid for what this call writes. Variables the store held from an earlier write keep whatever grid they were created with.

Parameters:

pathstr: Directory of the Zarr store.
append_dimstr, default “timestamp”: The dimension the store grows along.
dimstr, default “symbol”: The symbol axis to reconcile.
fill_valuesMapping[str, object], optional: Per-variable fill of the cells a new symbol or a new variable adds (NaN when absent), passed to both widens, pinned as the Zarr fill value of a new store and used to reindex data.
**kwargs: Passed through to append.

Examples

A store holding AAA and CCC, a window carrying DDD, BBB and a variable the store has not seen:

>>> XrBackend().to_internal(window).widen_and_append(
...     "prices.zarr", fill_values={"flag": False}
... )
XrBackend()
>>> stored = xr.open_zarr("prices.zarr")
>>> stored["symbol"].values.tolist(), sorted(stored.data_vars)
(['AAA', 'CCC', 'BBB', 'DDD'], ['amount', 'close', 'flag', 'volume'])

widen_data_vars(path, variables, append_dim='timestamp', fill_values=None, *, append_dim_size=None)[source]#

Add data variables to the store at path, backfilled over its extent.

The variable-set counterpart of widen_symbol_axis. Each name in variables that the store does not yet hold is materialised across the store’s whole existing extent with the fill value (NaN for a floating-point dtype, the fill_values entry otherwise). Names the store already holds are left alone, and when nothing is new the method returns without writing. This is the explicit opt-in for a panel that grew a column; a plain append still refuses a changed variable set because Zarr would otherwise write the new variable over the incoming window only and leave the store unopenable.

The filler carries dimensions but no coordinates: the store already holds them, and writing a decoded coordinate back can land on a different dtype than the one Zarr recorded. No directory swap is needed because a failed partial write raises and leaves the store intact.

Parameters:

pathstr: Directory of the Zarr store.
variablesMapping[str, object]: Maps each variable name to its dtype. The dtype is all the filler needs, so the caller’s panel is never held.
append_dimstr, default “timestamp”: The store’s append dimension, used to pin the filler’s chunk grid to the store’s.
fill_valuesMapping[str, object], optional: Per-variable fill for new variables that are not floating-point. Without an entry such a variable is refused, because np.full with NaN yields 0 for integers and True for booleans, fabricating history.
append_dim_sizeint, optional: The store’s eventual length along append_dim, so a filler added to a store that has not reached its final extent joins on the same grid as the variables already there.

Raises:

FileNotFoundError: If no store exists at path.
ValueError: If a new non-floating variable has no fill value.

Examples

>>> XrBackend().widen_data_vars("prices.zarr", {"volume": "float64"})
XrBackend()
>>> XrBackend().widen_data_vars(
...     "prices.zarr", {"flag": "bool"}, fill_values={"flag": False}
... )
XrBackend()
>>> stored = xr.open_zarr("prices.zarr")
>>> sorted(stored.data_vars), stored["flag"].dtype
(['close', 'flag', 'volume'], dtype('bool'))

widen_symbol_axis(path, symbols, dim='symbol', append_dim='timestamp', fill_values=None, *, append_dim_size=None)[source]#

Rewrite the store at path onto a superset dim axis.

Every label already in the store keeps its history unchanged; each new label is filled across the whole stored extent (NaN for a floating-point variable, the value named in fill_values otherwise). data is not touched, so a caller can widen a store without holding a panel. This is the explicit opt-in for a roster that grew between runs; a plain append still refuses a changed coordinate.

The rewrite lands in a .widening.tmp sidecar and is then swapped in with two same-parent renames (store to .superseded.tmp, sidecar to store) followed by removal of the superseded copy. A crash between the renames leaves no store at path and a superseded sidecar holding the only copy; the next call refuses and names the manual move rather than guessing which directory is authoritative. A non-empty superseded sidecar beside a live store is also refused, because it may be a real store left by an interrupted rebuild. Empty or never-authoritative residue is removed and the widen proceeds.

The strategy is chosen by size: at or under MAX_WIDEN_BYTES the whole store is reindexed in memory and written once; above it the store is rewritten block by block along append_dim and a warning names the figures involved. Both leave an identical store.

symbols is re-spelled in the stored axis’s dtype before anything is compared, so a ticker handed to an integer axis raises from normalize_to_axis_dtype instead of matching nothing and writing an all-NaN store.

Widening backfills NaN over a new label’s history; it does not recover data a vendor may have had. Rebuilding from raw data is the route when that history matters.

Parameters:

pathstr: Directory of the Zarr store.
symbolsSequence[str]: The labels the rewritten dim axis must contain. Must be a superset of the stored axis.
dimstr, default “symbol”: The axis being widened.
append_dimstr, default “timestamp”: The store’s append dimension, used to size blocks and to pin the rewritten chunk grid.
fill_valuesMapping[str, object], optional: Per-variable fill for variables carrying dim that are not floating-point. Without an entry such a variable is refused, because an unfilled reindex would silently upcast it to float64.
append_dim_sizeint, optional: The store’s eventual length along append_dim. This rewrite re-pins the chunk grid, so a caller widening a store that has not yet reached its final extent states the extent here; None sizes the grid from the store as it is now.

Raises:

FileNotFoundError: If no store exists at path.
ValueError: If crash residue makes the store’s identity ambiguous, if symbols would drop a stored label, or if a non-floating variable has no fill value.

Examples

>>> XrBackend().widen_symbol_axis("prices.zarr", ["AAA", "BBB", "CCC"])
XrBackend()
>>> stored = xr.open_zarr("prices.zarr")
>>> stored["symbol"].values.tolist()
['AAA', 'BBB', 'CCC']
>>> stored["close"].sel(symbol="CCC").values  # backfilled history
array([nan, nan, nan, nan, nan, nan])

write(path, **kwargs)[source]#

Write data to path, replacing any store already there.

The parent directory is created if needed and mode defaults to "w". No chunk encoding is applied, so the store lands on Zarr’s default chunk grid.

Parameters:

pathstr: Directory of the Zarr store.
**kwargs: Passed through to Dataset.to_zarr.

Examples

>>> XrBackend().to_internal(panel).write("prices.zarr")
XrBackend()
>>> Path("prices.zarr").is_dir()
True