diff --git a/covjsonkit/stream.py b/covjsonkit/stream.py new file mode 100644 index 0000000..ab490a5 --- /dev/null +++ b/covjsonkit/stream.py @@ -0,0 +1,514 @@ +"""Streaming CoverageJSON encoder over a block stream. + +``CovjsonStreamEncoder`` consumes the block IR produced by polytope-mars (``polytope_mars.blocks``) and +returns the CoverageJSON document in pieces: ``begin(header)`` -> collection opening, +``encode_iter(block)`` -> the fragments of one block, ``end()`` -> collection closing. Concatenating the +pieces gives exactly the bytes the legacy path produced with +``json.dumps(encoder.from_polytope*(tree)).encode()``, except that missing values (NaN) are written as +``null``. + +``encode_iter`` yields one fragment per slice of the block's arrays, each at most +``max_fragment_bytes`` (constructor argument, default :data:`DEFAULT_MAX_FRAGMENT_BYTES` = 8 MiB), so +neither the encoder nor its consumer ever holds the text of a whole block: a whole-world O2560 +coordinates block is ~800 MB of JSON but is handed over 8 MiB at a time. The fragments of a block must +be consumed completely and in order before the next block is encoded (the encoder's state advances as +they are produced). ``encode(block)`` is the buffered equivalent, ``b"".join(encode_iter(block))``, kept +for callers that want one ``bytes`` per block. The structural fragments (a coverage's metadata and ``t`` +axis, the collection's ``referencing`` and ``parameters``) are written whole: they are a few hundred +bytes, and only point-feature collections, which are small by construction, are buffered to ``end()``. + +Blocks are read structurally (attributes only), so this module imports neither polytope-feature nor +polytope-mars: + +* header: ``feature_type``, ``domain_type``, ``time_axis``, ``parameters`` (each with ``id``, + ``shortname``, ``name``, ``unit``, ``description``), ``referencing_coordinates``, ``extra``; +* group (``block.group``): ``index``, ``t``, ``params``, ``levels``, ``n_points``, ``mars_metadata``; +* coordinates block: ``lat``, ``lon`` (float arrays); values block: ``param``, ``level``, ``values`` + (float array, NaN = missing); group end: only ``group``. + +A group's points may arrive in one block or in several consecutive ones: blocks are placed by arrival +order, so a producer that splits a field writes the same bytes as one that does not. + +MultiPoint coverages are written as the blocks arrive (memory independent of the coverage size, apart +from the coordinates of the current coverage when it has several levels). PointSeries, +VerticalProfile and Trajectory collections are small and their legacy layout is point-major across field +groups, so their groups are buffered and written by ``end()``. +""" + +from __future__ import annotations + +import json +import math +from typing import Iterator + +import numpy as np +import orjson + +__all__ = [ + "CovjsonStreamEncoder", + "DEFAULT_MAX_FRAGMENT_BYTES", + "pointseries_coverages", +] + +#: Default upper bound on the size of one fragment of :meth:`CovjsonStreamEncoder.encode_iter`. +DEFAULT_MAX_FRAGMENT_BYTES = 8 * 1024 * 1024 + +#: Longest text one float64 can take (``-1.2345678901234567e-308``), used to size array slices. +_FLOAT_CHARS = 24 + + +_CRS = {"type": "GeographicCRS", "id": "http://www.opengis.net/def/crs/OGC/1.3/CRS84"} + + +def _dumps(obj) -> bytes: + """``json.dumps`` bytes (default separators, ASCII escapes) for small structural fragments.""" + return json.dumps(obj).encode("utf-8") + + +def _needs_repr(arr: np.ndarray) -> np.ndarray: + """Mask of the values orjson prints differently from ``json.dumps`` (0 < |x| < 1e-4).""" + with np.errstate(invalid="ignore"): + return (np.abs(arr) < 1e-4) & (arr != 0) + + +def _fragment_limit(value) -> int: + """``value`` as a positive fragment size in bytes.""" + try: + limit = int(value) + except (TypeError, ValueError): + raise ValueError(f"max_fragment_bytes must be an integer number of bytes, got {value!r}") from None + if limit <= 0: + raise ValueError(f"max_fragment_bytes must be positive, got {limit}") + return limit + + +def _fragment_slices(n: int, item_bytes: int, max_fragment_bytes: int) -> Iterator[tuple]: + """Per fragment, the ``(start, stop)`` point range to serialise into it. + + ``item_bytes`` is an upper bound on the text of one point (its floats at their longest plus + separators), so a fragment stays within ``max_fragment_bytes``. A fragment always holds at least one + point: with a limit smaller than one point's text, the fragment is one point long and exceeds it. + """ + per_fragment = max(1, max_fragment_bytes // max(1, item_bytes)) + for start in range(0, n, per_fragment): + yield start, min(start + per_fragment, n) + + +def format_floats(values) -> list: + """Each value as ``json.dumps`` writes a Python float, with NaN/inf as ``null``. + + orjson and CPython both print the shortest round-trip representation; they differ only in the + exponent form of numbers below 1e-4 (orjson ``1e-5``/``0.00002`` vs Python ``1e-05``/``2e-05``), + which are re-printed with ``repr``. + """ + arr = np.ascontiguousarray(values, dtype=np.float64) + if arr.size == 0: + return [] + parts = orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY)[1:-1].split(b",") + for i in np.flatnonzero(_needs_repr(arr)): + parts[i] = repr(arr[i].item()).encode("ascii") + return parts + + +def float_list_bytes(values) -> bytes: + """``a, b, c`` (no brackets), as ``json.dumps`` separates list items.""" + arr = np.ascontiguousarray(values, dtype=np.float64) + if not _needs_repr(arr).any(): + # fast path without one Python object per value + return orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY)[1:-1].replace(b",", b", ") + return b", ".join(format_floats(arr)) + + +def composite_tuples(lat, lon, level) -> bytes: + """``[lat, lon, level], ...`` (no outer brackets) as ``json.dumps`` writes the composite values.""" + lat = np.ascontiguousarray(lat, dtype=np.float64) + lon = np.ascontiguousarray(lon, dtype=np.float64) + if lat.size == 0: + return b"" + suffix = b", " + _dumps(level) + b"]" + if not (_needs_repr(lat).any() or _needs_repr(lon).any()): + pairs = orjson.dumps(np.column_stack((lat, lon)), option=orjson.OPT_SERIALIZE_NUMPY)[1:-1] + return pairs.replace(b",", b", ").replace(b"]", suffix) + return b", ".join([b"[" + a + b", " + o + suffix for a, o in zip(format_floats(lat), format_floats(lon))]) + + +def _py_values(values) -> list: + """Floats for ``json.dumps``; NaN -> None.""" + return [None if (v is None or (isinstance(v, float) and not math.isfinite(v))) else v for v in values] + + +def _parameter(p) -> dict: + return { + "type": "Parameter", + "description": {"en": p.description}, + "unit": {"symbol": p.unit}, + "observedProperty": {"id": p.shortname, "label": {"en": p.name}}, + } + + +def _range(name, values: list) -> dict: + return {"type": "NdArray", "dataType": "float", "shape": [len(values)], "axisNames": [name], "values": values} + + +def _series(groups) -> list: + """Consecutive buffered groups with equal metadata and levels form one coverage series.""" + series = [] + for g in groups: + if series and series[-1][0]["meta"] == g["meta"] and series[-1][0]["levels"] == g["levels"]: + series[-1].append(g) + else: + series.append([g]) + return series + + +def pointseries_coverages(groups, header, shortname) -> list: + """Legacy PointSeries layout (TimeSeries / Position encoders): one coverage per (point, series, level). + + A series is a run of field groups (time instants) sharing their coverage metadata; its ``t`` axis is + the groups' ``t`` values in order. Coverages are point-major (``for point: for series``) unless + ``header.extra["pointseries_order"] == "series_major"``. Position coverages from ``from_polytope`` + write the first level on their levelist axis whatever their level (preserved quirk). + """ + position_quirk = header.feature_type == "position" and header.time_axis == "date" + order = (header.extra or {}).get("pointseries_order", "point_major") + all_series = _series(groups) + + def coverage(series, i, level): + first = series[0] + levels = first["levels"] + if position_quirk and levels: + levelist = [levels[0]] + else: + levelist = [0 if level is None else level] + t = [v for g in series for v in g["t"]] + ranges = {} + for p in header.parameters: + if not any(p.id in g["values_params"] for g in series): + continue + vals = [] + for g in series: + arr = g["values"].get((p.id, level)) + vals.append(arr[i].item() if arr is not None else None) + ranges[shortname[p.id]] = _range(shortname[p.id], _py_values(vals)) + return { + "mars:metadata": first["meta"], + "type": "Coverage", + "domain": { + "type": "Domain", + "axes": { + "latitude": {"values": [first["lat"][i].item()]}, + "longitude": {"values": [first["lon"][i].item()]}, + "levelist": {"values": levelist}, + "t": {"values": t}, + }, + }, + "ranges": ranges, + } + + out = [] + if order == "series_major": + for series in all_series: + for i in range(len(series[0]["lat"])): + for level in series[0]["levels"] or (None,): + out.append(coverage(series, i, level)) + return out + n_points = len(all_series[0][0]["lat"]) if all_series else 0 + for i in range(n_points): + for series in all_series: + for level in series[0]["levels"] or (None,): + out.append(coverage(series, i, level)) + return out + + +def verticalprofile_coverages(groups, header, shortname) -> list: + """Legacy VerticalProfile layout: one coverage per (point, group), all levels on the levelist axis.""" + out = [] + n_points = len(groups[0]["lat"]) if groups else 0 + for i in range(n_points): + for g in groups: + levels = list(g["levels"]) if g["levels"] else [0] + keys = list(g["levels"]) if g["levels"] else [None] + ranges = {} + for p in header.parameters: + if p.id not in g["values_params"]: + continue + vals = [g["values"][(p.id, k)][i].item() for k in keys] + ranges[shortname[p.id]] = { + "type": "NdArray", + "dataType": "float", + "shape": [len(vals)], + "axisNames": ["levelist"], + "values": _py_values(vals), + } + out.append( + { + "mars:metadata": g["meta"], + "type": "Coverage", + "domain": { + "type": "Domain", + "axes": { + "latitude": {"values": [g["lat"][i].item()]}, + "longitude": {"values": [g["lon"][i].item()]}, + "levelist": {"values": levels}, + "t": {"values": list(g["t"])}, + }, + }, + "ranges": ranges, + } + ) + return out + + +def trajectory_coverages(groups, header, shortname, coordinates) -> list: + """Legacy Trajectory layout: one coverage per series; composite ``[t, lat, lon, level]``.""" + out = [] + for series in _series(groups): + composite = [] + values = {p.id: [] for p in header.parameters} + for g in series: + t = g["t"][0] if g["t"] else 0 + keys = list(g["levels"]) if g["levels"] else [None] + for k in keys: + level = 0 if k is None else k + composite.extend([t, la.item(), lo.item(), level] for la, lo in zip(g["lat"], g["lon"])) + for p in header.parameters: + arr = g["values"].get((p.id, k)) + values[p.id].extend(arr.tolist() if arr is not None else [None] * len(g["lat"])) + ranges = { + shortname[p.id]: _range(shortname[p.id], _py_values(values[p.id])) + for p in header.parameters + if any(p.id in g["values_params"] for g in series) + } + out.append( + { + "mars:metadata": series[0]["meta"], + "type": "Coverage", + "domain": { + "type": "Domain", + "axes": {"composite": {"dataType": "tuple", "coordinates": coordinates, "values": composite}}, + }, + "ranges": ranges, + } + ) + return out + + +class _OpenCoverage: + """State of the MultiPoint coverage being written.""" + + __slots__ = ("group", "n_tuples", "coords", "composite_closed", "param", "n_values") + + def __init__(self, group): + self.group = group + self.n_tuples = 0 + #: (lat, lon) arrays of every band, kept only when the coverage has more than one level + self.coords: list = [] + self.composite_closed = False + self.param = None + self.n_values = 0 + + +class CovjsonStreamEncoder: + """CoverageJSON encoder for the polytope-mars block stream (one instance per request).""" + + content_type = "application/prs.coverage+json" + file_extension = "covjson" + + def __init__(self, config=None, max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES): + self.config = dict(config or {}) + #: upper bound on one fragment of :meth:`encode_iter` (``config["max_fragment_bytes"]`` wins) + self.max_fragment_bytes = _fragment_limit(self.config.get("max_fragment_bytes") or max_fragment_bytes) + self._header = None + self.n_coverages = 0 + self._streaming = True + self._buffered: list = [] + self._current: _OpenCoverage | None = None + + @property + def header(self): + if self._header is None: + raise RuntimeError("CovjsonStreamEncoder.begin() must be called first") + return self._header + + def _open(self) -> _OpenCoverage: + if self._current is None: + raise RuntimeError("values block before the coordinates of its group") + return self._current + + # -- protocol ---------------------------------------------------------------------------------------- + + def begin(self, header) -> bytes: + self._header = header + self.n_coverages = 0 + self._shortname = {p.id: p.shortname for p in header.parameters} + self._coordinates = list(header.referencing_coordinates) + #: ids of the parameters some coverage has a range for; the collection's ``parameters`` lists + #: these, so a reader can take it as the list of what the result holds + self._present: set = set() + self._streaming = header.domain_type == "MultiPoint" or ( + header.domain_type == "Trajectory" and header.time_axis == "hdate" + ) + self._with_t = header.domain_type == "MultiPoint" + self._buffered = [] + self._current = None + domain = "MultiPoint" if header.domain_type in ("shapefile",) else header.domain_type + return b'{"type": "CoverageCollection", "domainType": ' + _dumps(domain) + b', "coverages": [' + + def encode(self, block) -> bytes: + """The whole block as one ``bytes`` (``b"".join(self.encode_iter(block))``).""" + return b"".join(self.encode_iter(block)) + + def encode_iter(self, block) -> Iterator[bytes]: + """The block as fragments of at most ``max_fragment_bytes``, each from a slice of its arrays. + + The fragments of a block must be consumed completely, in order, before the next block is + encoded: the encoder's state (separators, open ranges) advances as they are produced. + """ + if hasattr(block, "lat"): + return self._coords(block) + if hasattr(block, "values") and hasattr(block, "param"): + return self._values(block) + return self._group_end(block) + + def end(self) -> bytes: + out = [] + if not self._streaming: + covs = self._buffered_coverages() + for i, cov in enumerate(covs): + out.append((b", " if self.n_coverages or i else b"") + _dumps(cov)) + self.n_coverages += len(covs) + self._buffered = [] + parameters = {p.shortname: _parameter(p) for p in self.header.parameters if p.id in self._present} + referencing = [{"coordinates": self._coordinates, "system": dict(_CRS)}] + out.append(b'], "referencing": ' + _dumps(referencing) + b', "parameters": ' + _dumps(parameters) + b"}") + return b"".join(out) + + # -- MultiPoint: stream ------------------------------------------------------------------------------- + + def _coords(self, block) -> Iterator[bytes]: + g = block.group + if not self._streaming: + cur = self._buffer_group(g) + cur["lat"].append(np.asarray(block.lat, dtype=np.float64)) + cur["lon"].append(np.asarray(block.lon, dtype=np.float64)) + return + if self._current is None or self._current.group is not g: + yield self._open_coverage(g) + cur = self._open() + levels = list(g.levels) if g.levels else [0] + yield from self._tuples(cur, block.lat, block.lon, levels[0]) + if len(levels) > 1: + cur.coords.append((np.array(block.lat, dtype=np.float64), np.array(block.lon, dtype=np.float64))) + + def _tuples(self, cur: _OpenCoverage, lat, lon, level) -> Iterator[bytes]: + """The composite tuples of one block of points at one level, sliced into fragments. + + The ``(n, 3)`` tuples are built per slice, so the text of a whole block never exists at once. + """ + lat = np.asarray(lat) + lon = np.asarray(lon) + item = 2 * _FLOAT_CHARS + 8 + len(_dumps(level)) # "[lat, lon, level], " + for a, b in _fragment_slices(lat.size, item, self.max_fragment_bytes): + body = composite_tuples(lat[a:b], lon[a:b], level) + if not body: + continue + sep = b", " if cur.n_tuples else b"" + cur.n_tuples += b - a + yield sep + body + + def _open_coverage(self, g) -> bytes: + self._current = _OpenCoverage(g) + head = (b", " if self.n_coverages else b"") + b'{"mars:metadata": ' + _dumps(g.mars_metadata) + head += b', "type": "Coverage", "domain": {"type": "Domain", "axes": {' + if self._with_t: + head += b'"t": {"values": ' + _dumps(list(g.t)) + b"}, " + head += b'"composite": {"dataType": "tuple", "coordinates": ' + _dumps(self._coordinates) + b', "values": [' + self.n_coverages += 1 + return head + + def _close_composite(self) -> Iterator[bytes]: + cur = self._open() + if cur.composite_closed: + return + cur.composite_closed = True + levels = list(cur.group.levels) + for level in levels[1:]: + for lat, lon in cur.coords: + yield from self._tuples(cur, lat, lon, level) + cur.coords = [] + yield b']}}}, "ranges": {' + + def _values(self, block) -> Iterator[bytes]: + g = block.group + self._present.add(block.param) + if not self._streaming: + cur = self._buffer_group(g) + cur["values"].setdefault((block.param, block.level), []).append(np.asarray(block.values, dtype=np.float64)) + return + cur = self._open() + yield from self._close_composite() + if cur.param != block.param: + name = self._shortname.get(block.param, block.param) + n = g.n_points * max(1, len(g.levels)) + head = b"]}, " if cur.param is not None else b"" + head += _dumps(name) + b': {"type": "NdArray", "dataType": "float", "shape": [' + str(n).encode() + head += b'], "axisNames": [' + _dumps(str(name)) + b'], "values": [' + cur.param = block.param + cur.n_values = 0 + yield head + values = np.asarray(block.values) + for a, b in _fragment_slices(values.size, _FLOAT_CHARS + 2, self.max_fragment_bytes): + body = float_list_bytes(values[a:b]) + if not body: + continue + sep = b", " if cur.n_values else b"" + cur.n_values += b - a + yield sep + body + + def _group_end(self, block) -> Iterator[bytes]: + if not self._streaming: + cur = self._buffered[-1] if self._buffered else None + if cur is not None and cur["group"] is block.group: + cur["lat"] = np.concatenate(cur["lat"]) if cur["lat"] else np.empty(0) + cur["lon"] = np.concatenate(cur["lon"]) if cur["lon"] else np.empty(0) + cur["values"] = {k: np.concatenate(v) for k, v in cur["values"].items()} + cur["done"] = True + return + cur = self._current + if cur is None or cur.group is not block.group: + return + yield from self._close_composite() + yield b"]}}}" if cur.param is not None else b"}}" + self._current = None + + # -- buffered domains --------------------------------------------------------------------------------- + + def _buffer_group(self, g) -> dict: + if not self._buffered or self._buffered[-1]["group"] is not g: + self._buffered.append( + { + "group": g, + "meta": g.mars_metadata, + "levels": tuple(g.levels), + "t": tuple(g.t), + "values_params": set(g.params), + "lat": [], + "lon": [], + "values": {}, + "done": False, + } + ) + return self._buffered[-1] + + def _buffered_coverages(self) -> list: + groups = [g for g in self._buffered if g["done"]] + if not groups: + return [] + domain = self.header.domain_type + if domain == "PointSeries": + return pointseries_coverages(groups, self.header, self._shortname) + if domain == "VerticalProfile": + return verticalprofile_coverages(groups, self.header, self._shortname) + if domain == "Trajectory": + return trajectory_coverages(groups, self.header, self._shortname, self._coordinates) + raise ValueError(f"Unsupported domain type {domain!r}") diff --git a/tests/stream_memory_probe.py b/tests/stream_memory_probe.py new file mode 100644 index 0000000..dabb899 --- /dev/null +++ b/tests/stream_memory_probe.py @@ -0,0 +1,117 @@ +"""Peak RSS of streaming one large MultiPoint coverage, as a subprocess probe. + +Run it directly to measure; it prints one JSON line:: + + python tests/stream_memory_probe.py [n_points] + +The blocks are duck-typed stand-ins for polytope-mars' ``CoordsBlock`` / ``ValuesBlock`` / ``GroupEnd``. +The process' peak RSS is reset once the arrays are allocated (``/proc/self/clear_refs``, so the peak a +parent process left behind does not mask the measurement), hence the reported growth is what encoding +costs on top of the data itself. ``test_stream_memory.py`` asserts the bounds. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import resource +from types import SimpleNamespace + +import numpy as np + +from covjsonkit.stream import CovjsonStreamEncoder + + +def peak_rss_bytes() -> int: + """Peak resident set size of this process since the last :func:`reset_peak_rss`.""" + try: + with open("/proc/self/status", "r") as status: + for line in status: + if line.startswith("VmHWM:"): + return int(line.split()[1]) * 1024 + except OSError: + pass + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024 + + +def reset_peak_rss() -> bool: + """Forget the peak RSS so far (Linux >= 4.0); ``False`` when the kernel does not allow it.""" + try: + with open("/proc/self/clear_refs", "w") as clear_refs: + clear_refs.write("5") + except OSError: + return False + return True + + +def one_coverage(n_points: int): + """Header and blocks of one coverage of ``n_points`` points with one parameter.""" + param = SimpleNamespace(id="167", shortname="2t", name="2 metre temperature", unit="K", description="") + header = SimpleNamespace( + feature_type="boundingbox", + domain_type="MultiPoint", + time_axis="date", + parameters=(param,), + mars_metadata={}, + referencing_coordinates=("latitude", "longitude", "levelist"), + extra={}, + ) + group = SimpleNamespace( + index=0, + path={}, + t=("2024-01-01T00:00:00Z",), + params=("167",), + levels=(), + n_points=n_points, + mars_metadata={"class": "od", "stream": "oper", "number": 0, "step": 0}, + ) + coords = SimpleNamespace( + group=group, + lat=np.linspace(-90.0, 90.0, n_points), + lon=np.linspace(0.0, 359.9, n_points), + ) + values = SimpleNamespace( + group=group, + param="167", + level=None, + values=np.linspace(200.0, 320.0, n_points), + ) + return header, [coords, values, SimpleNamespace(group=group)] + + +def measure(n_points: int) -> dict: + header, blocks = one_coverage(n_points) + encoder = CovjsonStreamEncoder() + gc.collect() + peak_reset = reset_peak_rss() + baseline = peak_rss_bytes() + + total = len(encoder.begin(header)) + n_fragments = 0 + largest = 0 + for block in blocks: + for fragment in encoder.encode_iter(block): + n_fragments += 1 + largest = max(largest, len(fragment)) + total += len(fragment) + total += len(encoder.end()) + + peak = peak_rss_bytes() + return { + "n_points": n_points, + "document_bytes": total, + "n_fragments": n_fragments, + "largest_fragment_bytes": largest, + "max_fragment_bytes": encoder.max_fragment_bytes, + "peak_was_reset": peak_reset, + "baseline_rss_bytes": baseline, + "peak_rss_bytes": peak, + "peak_rss_growth_bytes": peak - baseline, + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Measure the peak RSS of streaming one large coverage.") + parser.add_argument("n_points", nargs="?", type=int, default=5_000_000) + print(json.dumps(measure(parser.parse_args().n_points))) diff --git a/tests/test_stream_encoder.py b/tests/test_stream_encoder.py new file mode 100644 index 0000000..6f388c8 --- /dev/null +++ b/tests/test_stream_encoder.py @@ -0,0 +1,405 @@ +"""CovjsonStreamEncoder on synthetic blocks (no polytope needed). + +The blocks are local stand-ins with the attributes polytope-mars' block IR exposes; the encoder only +reads attributes. Expected bytes are ``json.dumps`` of the document the legacy encoders built. +""" + +import json +import subprocess +import sys +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import pytest + +from covjsonkit.stream import ( + DEFAULT_MAX_FRAGMENT_BYTES, + CovjsonStreamEncoder, + composite_tuples, + float_list_bytes, + format_floats, +) + + +@dataclass(frozen=True) +class Param: + id: str + shortname: str + name: str + unit: str + description: str = "" + + +@dataclass(frozen=True) +class Header: + feature_type: str + domain_type: str + time_axis: str + parameters: tuple + #: polytope-mars decides these (``legacy_format.referencing_coordinates``); the encoder writes them + referencing_coordinates: tuple = ("latitude", "longitude", "levelist") + mars_metadata: dict = field(default_factory=dict) + extra: dict = field(default_factory=dict) + + +@dataclass(frozen=True) +class Group: + index: int + path: dict + t: tuple + params: tuple + levels: tuple + n_points: int + mars_metadata: dict + + +@dataclass(frozen=True) +class Coords: + group: Any + lat: np.ndarray + lon: np.ndarray + + +@dataclass(frozen=True) +class Values: + group: Any + param: str + level: Any + values: np.ndarray + + +@dataclass(frozen=True) +class End: + group: Any + + +T2 = Param("167", "2t", "2 metre temperature", "K", "Température à 2 m") +U10 = Param("165", "10u", "10 metre U wind component", "m s**-1") +CRS = {"type": "GeographicCRS", "id": "http://www.opengis.net/def/crs/OGC/1.3/CRS84"} + + +def _parameters(params): + return { + p.shortname: { + "type": "Parameter", + "description": {"en": p.description}, + "unit": {"symbol": p.unit}, + "observedProperty": {"id": p.shortname, "label": {"en": p.name}}, + } + for p in params + } + + +def _nan_to_none(values): + return [None if v != v else v for v in np.asarray(values, dtype=np.float64).tolist()] + + +def multipoint_blocks(fields, lat, lon, levels=(), meta=None, t=("2024-01-01T00:00:00Z",), n_blocks=1, index=0): + """Blocks of one MultiPoint group; ``fields`` = {(param, level): values}. + + ``n_blocks`` splits the group's points over that many consecutive blocks, which a producer that + cannot fetch a field whole would do and which must not change the bytes. + """ + params = tuple(dict.fromkeys(p for p, _ in fields)) + bounds = np.array_split(np.arange(len(lat)), n_blocks) + g = Group(index, {}, t, params, tuple(levels), len(lat), meta or {"class": "od", "number": 0}) + blocks: list[Any] = [Coords(g, lat[idx], lon[idx]) for idx in bounds] + for p in params: + for lev in levels or (None,): + for idx in bounds: + blocks.append(Values(g, p, lev, fields[(p, lev)][idx])) + blocks.append(End(g)) + return blocks + + +def encode(header, blocks) -> bytes: + enc = CovjsonStreamEncoder({"param_db": "ecmwf"}) + out = [enc.begin(header)] + out += [enc.encode(b) for b in blocks] + out.append(enc.end()) + return b"".join(out) + + +def fragments(header, blocks, max_fragment_bytes) -> list: + """Every fragment of the document, block by block, with ``encode_iter``.""" + enc = CovjsonStreamEncoder({"param_db": "ecmwf"}, max_fragment_bytes=max_fragment_bytes) + out = [enc.begin(header)] + for block in blocks: + out += list(enc.encode_iter(block)) + out.append(enc.end()) + return out + + +def legacy_multipoint(header, groups): + """What BoundingBox.from_polytope + json.dumps produced for the same data. + + The collection's ``parameters`` lists the parameters some coverage has a range for, in header order. + """ + coords = list(header.referencing_coordinates) + covs = [] + present = set() + for fields, lat, lon, levels, meta, t in groups: + composite = [[float(a), float(o), lev] for lev in (levels or (0,)) for a, o in zip(lat, lon)] + ranges = {} + params = tuple(dict.fromkeys(p for p, _ in fields)) + present.update(params) + short = {p.id: p.shortname for p in header.parameters} + for p in params: + vals = [v for lev in (levels or (None,)) for v in _nan_to_none(fields[(p, lev)])] + ranges[short[p]] = { + "type": "NdArray", + "dataType": "float", + "shape": [len(vals)], + "axisNames": [short[p]], + "values": vals, + } + covs.append( + { + "mars:metadata": meta, + "type": "Coverage", + "domain": { + "type": "Domain", + "axes": { + "t": {"values": list(t)}, + "composite": {"dataType": "tuple", "coordinates": coords, "values": composite}, + }, + }, + "ranges": ranges, + } + ) + doc = { + "type": "CoverageCollection", + "domainType": "MultiPoint", + "coverages": covs, + "referencing": [{"coordinates": coords, "system": CRS}], + "parameters": _parameters([p for p in header.parameters if p.id in present]), + } + return json.dumps(doc).encode() + + +def _data(n, seed=0): + rng = np.random.default_rng(seed) + lat = np.round(rng.uniform(-90, 90, n), 12) + lon = np.round(rng.uniform(0, 360, n), 12) + return lat, lon + + +BBOX = Header("boundingbox", "MultiPoint", "date", (U10, T2)) + + +def test_multipoint_matches_legacy_json_dumps(): + lat, lon = _data(7) + fields = {("165", None): np.linspace(1, 7, 7) * 1.1e10 + 0.123, ("167", None): np.arange(7) * 0.5} + meta = {"class": "od", "Forecast date": "2024-01-01T00:00:00Z", "number": 0, "step": 0} + out = encode(BBOX, multipoint_blocks(fields, lat, lon, meta=meta)) + expected = legacy_multipoint(BBOX, [(fields, lat, lon, (), meta, ("2024-01-01T00:00:00Z",))]) + assert out == expected + json.loads(out) + + +def test_several_blocks_per_group_give_the_same_bytes(): + """Blocks are placed by arrival order, so a group split over several blocks encodes identically.""" + lat, lon = _data(13, 1) + levels = ("500", "850") + fields = {(p, lev): np.random.default_rng(2).normal(size=13) * 1e3 for p in ("165", "167") for lev in levels} + ref = encode(BBOX, multipoint_blocks(fields, lat, lon, levels=levels)) + assert encode(BBOX, multipoint_blocks(fields, lat, lon, levels=levels, n_blocks=3)) == ref + group = (fields, lat, lon, levels, {"class": "od", "number": 0}, ("2024-01-01T00:00:00Z",)) + expected = legacy_multipoint(BBOX, [group]) + assert ref == expected + + +def _mixed_fields(n, levels, seed=11): + """Values covering both float paths: a NaN, a ``1e-5``-style number, ``-0.0`` and plain floats.""" + rng = np.random.default_rng(seed) + fields = {(p, lev): np.round(rng.normal(size=n) * 1e3, 9) for p in ("165", "167") for lev in levels or (None,)} + first = fields[("165", levels[0] if levels else None)] + first[0], first[1], first[2] = np.nan, 1.5e-5, -0.0 + return fields + + +@pytest.mark.parametrize("max_fragment_bytes", [64, 1024, DEFAULT_MAX_FRAGMENT_BYTES], ids=["64B", "1KiB", "8MiB"]) +def test_fragments_join_to_the_single_fragment_bytes(max_fragment_bytes): + levels = ("500", "850") + lat, lon = _data(500, 8) + lon[4] = 3e-5 # forces the per-value path for the composite tuples as well + fields = _mixed_fields(500, levels) + blocks = multipoint_blocks(fields, lat, lon, levels=levels, n_blocks=3) + joined = b"".join(fragments(BBOX, blocks, max_fragment_bytes)) + assert joined == encode(BBOX, blocks) + group = (fields, lat, lon, levels, {"class": "od", "number": 0}, ("2024-01-01T00:00:00Z",)) + assert joined == legacy_multipoint(BBOX, [group]) + + +@pytest.mark.parametrize("max_fragment_bytes", [1024, DEFAULT_MAX_FRAGMENT_BYTES], ids=["1KiB", "8MiB"]) +def test_no_fragment_exceeds_the_limit(max_fragment_bytes): + lat, lon = _data(5000, 9) + fields = {("167", None): np.linspace(200.0, 320.0, 5000)} + frags = fragments(BBOX, multipoint_blocks(fields, lat, lon, n_blocks=2), max_fragment_bytes) + assert max(len(f) for f in frags) <= max_fragment_bytes + assert all(f for f in frags) # no empty fragment is handed out + assert len(frags) > 5 # the arrays really were sliced + + +def test_small_magnitude_values_are_formatted_per_value_in_every_slice(): + vals = np.array([1e-5, 2e-5, 0.00002, 9.99e-5, 3.0, np.nan]) + lat, lon = _data(6, 10) + blocks = multipoint_blocks({("167", None): vals}, lat, lon, n_blocks=2) + out = b"".join(fragments(BBOX, blocks, 64)) + assert out == encode(BBOX, blocks) + assert b"1e-05" in out and b"2e-05" in out and b"1e-5" not in out + assert json.loads(out)["coverages"][0]["ranges"]["2t"]["values"] == [1e-5, 2e-5, 2e-5, 9.99e-5, 3.0, None] + + +def test_fragment_limit_is_configurable_and_validated(): + assert CovjsonStreamEncoder().max_fragment_bytes == DEFAULT_MAX_FRAGMENT_BYTES + assert CovjsonStreamEncoder({"max_fragment_bytes": 4096}).max_fragment_bytes == 4096 + assert CovjsonStreamEncoder(max_fragment_bytes=4096).max_fragment_bytes == 4096 + with pytest.raises(ValueError): + CovjsonStreamEncoder(max_fragment_bytes=-1) + with pytest.raises(ValueError): + CovjsonStreamEncoder({"max_fragment_bytes": "many"}) + + +def test_nan_is_null_and_missing_param_is_omitted(): + lat, lon = _data(4, 3) + vals = np.array([1.0, np.nan, 3.0, np.nan]) + out = encode(BBOX, multipoint_blocks({("167", None): vals}, lat, lon)) + doc = json.loads(out) + assert b"NaN" not in out + ranges = doc["coverages"][0]["ranges"] + assert list(ranges) == ["2t"] # 10u not in group.params -> no range + assert ranges["2t"]["values"] == [1.0, None, 3.0, None] + assert list(doc["parameters"]) == ["2t"] # a parameter with no range in any coverage is not listed + + +def test_several_groups_and_empty_collection(): + lat, lon = _data(3, 4) + blocks = multipoint_blocks({("167", None): np.ones(3)}, lat, lon, index=0) + blocks += multipoint_blocks({("167", None): np.zeros(3)}, lat, lon, index=1, t=("2024-01-02T00:00:00Z",)) + doc = json.loads(encode(BBOX, blocks)) + assert [c["domain"]["axes"]["t"]["values"] for c in doc["coverages"]] == [ + ["2024-01-01T00:00:00Z"], + ["2024-01-02T00:00:00Z"], + ] + empty = json.loads(encode(BBOX, [])) + assert empty["coverages"] == [] and list(empty) == ["type", "domainType", "coverages", "referencing", "parameters"] + assert empty["parameters"] == {} + + +def test_parameters_lists_what_the_coverages_hold_in_header_order(): + lat, lon = _data(3, 4) + blocks = multipoint_blocks({("167", None): np.ones(3)}, lat, lon, index=0) + blocks += multipoint_blocks({("165", None): np.zeros(3)}, lat, lon, index=1, t=("2024-01-02T00:00:00Z",)) + doc = json.loads(encode(BBOX, blocks)) + assert [list(c["ranges"]) for c in doc["coverages"]] == [["2t"], ["10u"]] + assert list(doc["parameters"]) == ["10u", "2t"] + + +def test_float_formatting_matches_json_dumps(): + rng = np.random.default_rng(5) + vals = np.concatenate( + [ + rng.normal(size=5000) * 10.0 ** rng.integers(-12, 20, 5000), + [0.0, -0.0, 1.0, 1e16, 1e-4, 9.99e-5, 2e-5, 1e-7, -3e-9, 123456789.123, 5e-324, 1.7976931348623157e308], + ] + ) + assert b", ".join(format_floats(vals)) == json.dumps(vals.tolist()).encode()[1:-1] + assert format_floats(np.array([np.nan, np.inf])) == [b"null", b"null"] + + +@pytest.mark.parametrize("small", [False, True], ids=["fast-path", "repr-fallback"]) +@pytest.mark.parametrize("level", [0, "500", 850]) +def test_composite_and_value_bytes_match_json_dumps(small, level): + lat, lon = _data(2000, 6) + lon[3] = -0.0 + if small: + lon[7] = 3e-5 # forces the per-value path + pairs = [[a, o, level] for a, o in zip(lat.tolist(), lon.tolist())] + assert composite_tuples(lat, lon, level) == json.dumps(pairs).encode()[1:-1] + values = np.concatenate([lat * 1e9, [np.nan, 2e-7 if small else 2.0]]) + expected = json.dumps([None if v != v else v for v in values.tolist()]).encode()[1:-1] + assert float_list_bytes(values) == expected + assert composite_tuples(np.empty(0), np.empty(0), level) == b"" + + +@pytest.mark.parametrize("coords", [("latitude", "longitude", "levelist"), ("x", "y", "z")]) +def test_the_header_decides_the_referencing_coordinates(coords): + """Which names a collection declares is polytope-mars' rule; the encoder writes what it is given.""" + header = Header("boundingbox", "MultiPoint", "date", (T2,), referencing_coordinates=coords) + lat, lon = _data(3, 7) + doc = json.loads(encode(header, multipoint_blocks({("167", None): np.ones(3)}, lat, lon))) + assert doc["referencing"][0]["coordinates"] == list(coords) + assert doc["coverages"][0]["domain"]["axes"]["composite"]["coordinates"] == list(coords) + + +def _point_group(index, t, meta, values, levels=(), lat=(51.5, -33.9), lon=(0.1, 18.4)): + params = tuple(dict.fromkeys(p for p, _ in values)) + g = Group(index, {}, t, params, tuple(levels), len(lat), meta) + blocks: list[Any] = [Coords(g, np.array(lat), np.array(lon))] + for (p, lev), v in values.items(): + blocks.append(Values(g, p, lev, np.asarray(v, dtype=float))) + return blocks + [End(g)] + + +def test_pointseries_layout_is_point_major_over_consecutive_groups(): + header = Header("timeseries", "PointSeries", "date", (T2,)) + meta = {"class": "od", "number": 0, "levelist": 0} + blocks = [] + for i, t in enumerate(["2024-01-01T00:00:00Z", "2024-01-01T01:00:00Z", "2024-01-01T02:00:00Z"]): + blocks += _point_group(i, (t,), meta, {("167", None): [i, 10 + i]}) + doc = json.loads(encode(header, blocks)) + assert len(doc["coverages"]) == 2 + first, second = doc["coverages"] + assert first["domain"]["axes"]["t"]["values"] == [ + "2024-01-01T00:00:00Z", + "2024-01-01T01:00:00Z", + "2024-01-01T02:00:00Z", + ] + assert first["ranges"]["2t"]["values"] == [0.0, 1.0, 2.0] + assert second["ranges"]["2t"]["values"] == [10.0, 11.0, 12.0] + assert second["domain"]["axes"]["latitude"]["values"] == [-33.9] + assert doc["referencing"][0]["coordinates"] == ["latitude", "longitude", "levelist"] + + +def test_buffered_domains_produce_their_coverages_from_end(): + """Point features are buffered: their blocks yield no fragment, ``end()`` writes the coverages.""" + header = Header("timeseries", "PointSeries", "date", (T2,)) + meta = {"class": "od", "number": 0, "levelist": 0} + blocks = _point_group(0, ("2024-01-01T00:00:00Z",), meta, {("167", None): [1, 2]}) + enc = CovjsonStreamEncoder(max_fragment_bytes=64) + begin = enc.begin(header) + assert [f for block in blocks for f in enc.encode_iter(block)] == [] + assert begin + enc.end() == encode(header, blocks) + + +def test_verticalprofile_layout(): + header = Header("verticalprofile", "VerticalProfile", "date", (T2,)) + blocks = _point_group( + 0, ("2024-01-01T00:00:00Z",), {"number": 0}, {("167", 500): [1, 2], ("167", 850): [3, 4]}, levels=(500, 850) + ) + doc = json.loads(encode(header, blocks)) + assert [c["ranges"]["2t"]["values"] for c in doc["coverages"]] == [[1.0, 3.0], [2.0, 4.0]] + assert doc["coverages"][0]["domain"]["axes"]["levelist"]["values"] == [500, 850] + assert doc["coverages"][0]["ranges"]["2t"]["axisNames"] == ["levelist"] + + +def test_trajectory_layout(): + header = Header("trajectory", "Trajectory", "date", (T2,), referencing_coordinates=("t", "x", "y", "z")) + blocks = _point_group(0, (0,), {"number": 0}, {("167", None): [1, 2]}) + doc = json.loads(encode(header, blocks)) + (cov,) = doc["coverages"] + assert cov["domain"]["axes"]["composite"]["values"] == [[0, 51.5, 0.1, 0], [0, -33.9, 18.4, 0]] + assert cov["domain"]["axes"]["composite"]["coordinates"] == ["t", "x", "y", "z"] + + +def test_no_polytope_imports(): + code = ( + "import sys, covjsonkit.stream; " + "roots = ('polytope_feature', 'polytope_mars', 'covjson_pydantic'); " + "bad = [m for m in sys.modules if m.split('.')[0] in roots]; " + "print(bad); sys.exit(1 if bad else 0)" + ) + proc = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert proc.returncode == 0, proc.stdout + proc.stderr diff --git a/tests/test_stream_memory.py b/tests/test_stream_memory.py new file mode 100644 index 0000000..a1a6f0e --- /dev/null +++ b/tests/test_stream_memory.py @@ -0,0 +1,40 @@ +"""The streaming encoder's memory bound: peak RSS must not follow the size of a block. + +The measurement runs in a subprocess (``stream_memory_probe.py``) so the peak RSS it reports is only +this encoding; run that module directly to see the numbers. +""" + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +from covjsonkit.stream import DEFAULT_MAX_FRAGMENT_BYTES + +PROBE = str(Path(__file__).with_name("stream_memory_probe.py")) +MB = 1024 * 1024 + + +def _probe(n_points: int) -> dict: + run = subprocess.run([sys.executable, PROBE, str(n_points)], capture_output=True, text=True) + assert run.returncode == 0, run.stdout + run.stderr + last = run.stdout.strip().splitlines()[-1] if run.stdout.strip() else "" + try: + return json.loads(last) + except json.JSONDecodeError as exc: + raise AssertionError(f"probe printed no measurement: {last!r}\n{run.stderr}") from exc + + +def test_peak_rss_growth_is_bounded_and_independent_of_the_block_size(): + small = _probe(1_000_000) + big = _probe(5_000_000) + if not (small["peak_was_reset"] and big["peak_was_reset"]): + pytest.skip("this kernel does not support resetting the peak RSS via /proc/self/clear_refs") + assert big["document_bytes"] > 300 * MB # ~320 MB of JSON text out of 5M points + assert big["largest_fragment_bytes"] <= DEFAULT_MAX_FRAGMENT_BYTES + assert small["largest_fragment_bytes"] <= DEFAULT_MAX_FRAGMENT_BYTES + assert big["peak_rss_growth_bytes"] < 100 * MB, big + # five times the points for 5 times the text, but not for more memory + assert big["peak_rss_growth_bytes"] < small["peak_rss_growth_bytes"] + 32 * MB, f"{small} vs {big}"