diff --git a/AGENTS.md b/AGENTS.md index 7b432db..36f9de6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -199,9 +199,13 @@ at-least-once path is still `SubscriptionsService::listen`. Two things differ fr `recommend_value_type` sits beside it: a bare `ValueTypeRecommendation`, and an unknown unit is the generic default with `recognized: false`, not an error. -### Binary datapoint ingest (`src/timeseries/binary.rs`) +### Binary datapoints (`src/timeseries/binary.rs`, `src/timeseries/binary_read.rs`) -`TimeSeriesService::insert_datapoints_binary` is the second ingest path, to `POST /timeseries/data/binary`: the same `DatapointsCollection` input, resolved through `/timeseries/byids` (cached per service instance; needs read access on the dataset), checked against each series' value type locally, cut into Arrow IPC frames at the contract's caps, zstd-compressed per frame (mandatory: level 1, 3 or 9, default 9) and posted through `execute_post_bytes_request`. Every refusal on this path is a problem document typed for its status — `invalid-frame` (400), `unknown-timeseries` (404), `request-too-large` (413), `unsupported-media-type` (415), `value-type-mismatch` and `external-id-mismatch` (422), `too-many-in-flight` (429) — with the kebab-case sub-case in a `reason` extension beside it. It was one type, `datapoint-block-rejected`, answering with all six statuses until the api split it (platform #120); the SDK matches the new slugs only, and nothing here recognises the old one — the binary path has never been in a release, so no caller can be on the other side of that change. `is_stale_series_rejection` reads the slug to decide the one re-resolve-and-retry, and reads it from the problem rather than the body text because the SDK's own pre-flight 404 names the series it could not resolve — an external id spelling `unknown-timeseries` used to trigger the retry. `FrameWriter` builds one frame and is public; `cut_into_writers` and `pack_requests` hold the caps. The byte layout is the platform's `binary_datapoints_format.md`. The arrow-rs crates (`arrow-array`, `arrow-schema`, `arrow-ipc`) exist for this path and are the seed of the Arrow read path. In Python it is `insert_datapoints_binary` and `insert_from_lists_binary` on the **sync** client only, each taking an optional `zstd_level`; the async client has no binary methods. The ignored `timeseries::tests::test_datapoints_binary` is the live twin of `test_datapoints` and needs a backend that serves the endpoint; the writer's own tests in `binary.rs` run offline. `python_tests/test_timeseries_binary.py` round-trips the Python binding through the JSON read path. +`TimeSeriesService::insert_datapoints_binary` is the second ingest path, to `POST /timeseries/data/binary`: the same `DatapointsCollection` input, resolved through `/timeseries/byids` (cached per service instance; needs read access on the dataset), checked against each series' value type locally, cut into Arrow IPC frames at the contract's caps, zstd-compressed per frame (mandatory: level 1, 3 or 9, default 3; 9 compressed noisy floats no smaller than 1 and took several times as long) and posted through `execute_post_bytes_request`. `BinaryIngestOptions::compression` picks where the zstd goes: `FrameCompression::Buffers` (compression 2, the default: arrow-rs compresses each buffer in-band, the form a read answers with) or `FrameCompression::Payload` (compression 1: one stream around the Arrow payload). Buffers is the default because the consumer relays those batches into ClickHouse without decompressing them (platform `c9e9351dd`), while a compression 1 frame is decoded first: on 100M float points from polars, time until the last point was readable fell from 15.5 s to 13.2 s, and the client's cost was the same. At compression 2 the envelope's raw length is what the server recounts — every record batch buffer uncompressed, padded to 8 — so `in_band_raw_len` reads it back from the written stream rather than predicting which buffers arrow-rs writes; a mismatch is a 400. The api only takes compression 2 on ingest from platform `357794cce` on. Every refusal on this path is a problem document typed for its status — `invalid-frame` (400), `unknown-timeseries` (404), `request-too-large` (413), `unsupported-media-type` (415), `value-type-mismatch` and `external-id-mismatch` (422), `too-many-in-flight` (429) — with the kebab-case sub-case in a `reason` extension beside it. It was one type, `datapoint-block-rejected`, answering with all six statuses until the api split it (platform #120); the SDK matches the new slugs only, and nothing here recognises the old one — the binary path has never been in a release, so no caller can be on the other side of that change. `is_stale_series_rejection` reads the slug to decide the one re-resolve-and-retry, and reads it from the problem rather than the body text because the SDK's own pre-flight 404 names the series it could not resolve — an external id spelling `unknown-timeseries` used to trigger the retry. `FrameWriter` builds one frame and is public; `cut_into_writers` and `pack_requests` hold the caps. The byte layout is the platform's `binary_datapoints_format.md`. The arrow-rs crates (`arrow-array`, `arrow-schema`, `arrow-ipc`) serve both directions. In Python it is `insert_datapoints_binary` and `insert_from_lists_binary` on the **sync** client only, each taking an optional `zstd_level` and `compression` (`"buffers"`, the default, or `"payload"`); the async client has no binary methods. The ignored `timeseries::tests::test_datapoints_binary` is the live twin of `test_datapoints` and needs a backend that serves the endpoint; the writer's own tests in `binary.rs` run offline. `python_tests/test_timeseries_binary.py` round-trips the Python binding through the JSON read path. + +The read is `TimeSeriesService::retrieve_datapoints_binary`, to `POST /timeseries/data/list/binary`: the JSON retrieve body in, the same frames out, one series per frame, then the 4-byte marker `DHDE`. Read frames are **compression 2**: the payload is an Arrow IPC stream whose record batches zstd-compress each buffer in-band (Arrow's body compression), because the api relays ClickHouse's compressed batches without decoding them; `arrow-ipc`'s `zstd` feature is what decodes that, and a frame holds whole batches up to the row cap, so most hold fewer rows than the cap. Compression 1 (one zstd stream around the payload, what ingest sends) is still read. At compression 2 the envelope's raw length is the batches' buffers uncompressed, which arrow-rs does not report, so only its cap is checked. A body without the marker was cut short (the server streams, so a failure after the first bytes can only end the body early) and is a **502** saying so, never fewer points; bytes after it are refused too. Codec 2 frames carry bucketed aggregates in their own schema (`avg`/`min`/`max`/`sum`, Float64, 50 000 buckets per frame) while the envelope's type byte stays the series' own. `parse_read_response` returns one columnar `BinaryDatapoints` per series, merging consecutive frames of one series; decimals stay exact as unscaled `i128` plus scale. `TimeSeriesService::stream_datapoints_binary` is the primitive: a `FrameAssembler` cuts the body into frames as chunks arrive (a frame may be split anywhere), each frame is an owned `Bytes` decoded on `spawn_blocking`, at most `min(cores, 8)` in flight through a `FuturesOrdered` so items stay in frame order, and a producer task feeds a bounded channel behind the `BinaryDatapointStream` (dropping it aborts the task and closes the connection). One item per frame, so a long series is several consecutive items. A refusal is the `Err` before there is a stream; a response that breaks off is an error item after the items that arrived, and that is the semantic difference stream consumers must handle. `retrieve_datapoints_binary` collects the stream and merges consecutive parts (each series allocated once at its final size); `parse_read_response` does the same for a body in memory, decoding on scoped threads. Measured on 100M float rows, 10 series: the server takes ~6.5 s; the buffered call adds ~0.24 s (the merge copy) and the stream into polars ~0.02 s. Single-threaded zstd used to add 1.3 s on its own. Differences from the JSON read worth knowing: no `cursor` (a 400; advance `start` past the last timestamp), a series with no points is absent rather than empty, and an unknown or unreadable series is a 400 for the whole request — `DataWrapper` serialises only `items`, so the request-level `ignoreUnknownIds` and `limit` are not reachable. `src/timeseries/fixtures/read_*.bin` are response bodies written by the platform's own `DatapointFrameWriter.relayStream` (`read_*_zstd.bin` are the compression 1 bodies its former `wrap` wrote), so the offline tests in `binary_read.rs` prove arrow-rs reads what the Java writer emits; regenerate them when the platform's writer changes. The ignored `timeseries::tests::test_datapoints_binary_read` is the live check against the JSON read. In Python it is `retrieve_datapoints_binary` on the sync client, plus `retrieve_datapoints_binary_stream`, whose single-use `BinaryDatapointReader` iterates per frame or is an Arrow stream of one long table (`timeseries_id`, `timestamp`, values; one schema, so a read mixing value types must be iterated). The long table carries no external id column on purpose: as a dictionary column, polars spent ~0.8 s per 100M rows making it categorical, more than everything else; `reader.external_ids` maps the ids instead. The buffered `retrieve_datapoints_binary` returns a `BinaryReadResult`, a sequence of per-series `BinaryDatapoints` with the same `external_ids`, long-table Arrow stream and `to_polars(wide=False)` as the reader: both call `python/intellistream_datahub_sdk/_tables.py`, so the two tables cannot drift. `wide=True` puts series side by side when every one has the same timestamps (free), pivots on a numeric key when they share one value column type (~1.7 s per 100M rows; joining names first doubled it, `concat(how="align")` tripled it), and aligns otherwise, which is how mixed value types sit side by side. Each per-series `BinaryDatapoints` implements the Arrow PyCapsule stream protocol (`__arrow_c_stream__`, via `arrow-array`'s `ffi` feature enabled in the bindings crate only), so polars, pyarrow and pandas 3 read them zero-copy; the package still declares no runtime dependency, and `to_polars`/`to_pandas` import theirs lazily. `python_tests/test_timeseries_binary_read.py` covers it offline against a stand-in server and live when the backend serves the endpoint. + +`TimeSeriesService::insert_table_binary` (`src/timeseries/binary_table.rs`) is the third ingest entry point, Arrow in: a wide `RecordBatch` (a timestamp column plus one column per series, named by external id). It shares series resolution and its cache, request packing and the stale-series retry with the string path (`resolve_series`, `post_frame_requests`, `assemble_frame` and `ipc_stream` in `binary.rs`), and stores what the string path stores: its offline tests decode both paths' frames with the read parser and compare. Columns already in the canon's type are sliced, not copied; nulls are dropped as "no point"; unsorted rows are sorted and a repeated timestamp keeps its last value; timestamps must carry a time zone and whole milliseconds. Frames hold one value type, series in ascending id order, short series sharing a frame, and are compressed on scoped threads. A `mixed` series takes a column as typed (numeric → `value_numeric`, string → `value_text`), unlike the string path's per-value classification. In Python it is `insert_dataframe_binary` (sync client), importing anything with `__arrow_c_stream__` zero-copy; with no `timestamp=` it takes the pandas index if that is a timestamp (pandas records it in the schema metadata, and other index columns are dropped as not series), then a `timestamp` column, then the only timestamp-typed column. 100M float points from polars took ~3 s with the client using ~3.6 cores, so the api, not the client, sets that pace; `python_tests/test_timeseries_binary_dataframe.py` round-trips offline through a stand-in that replays the posted frames, and live through `to_polars(wide=True)`. ### The `ApiServiceProvider` trait (`src/generic.rs`) diff --git a/Cargo.toml b/Cargo.toml index 885d051..3e27035 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ reqwest = { version = "0.12", default-features = false, features = ["blocking", chrono = { version ="0.4", features = ["serde"] } rand = "0.9" futures = "0.3" +bytes = "1" dotenv = "0.15" uuid = { version = "1.16", features = ["v4", "v7", "serde"] } regex = "1.11.1" @@ -44,7 +45,8 @@ geojson = "1" base64 = "0.22" arrow-array = "59.3" arrow-schema = "59.3" -arrow-ipc = "59.3" +arrow-ipc = { version = "59.3", features = ["zstd"] } +arrow-select = "59.3" #[lib] diff --git a/README.md b/README.md index 0653c0b..14478e1 100644 --- a/README.md +++ b/README.md @@ -124,7 +124,7 @@ the same collections as `insert_datapoints` and sends them through `POST /timeseries/data/binary`: each series is resolved once to its id and value type and cached, values are checked against that type locally, sorted and de-duplicated, cut into Arrow IPC frames at the contract's caps (100 000 points per numeric frame, 10 000 per text or mixed -frame, 32 frames per request), compressed with zstd (level 9 by default, 1 and 3 are the other +frame, 32 frames per request), compressed with zstd (level 3 by default, 1 and 9 are the other choices) and posted. A 204 means every frame was accepted. A series that does not exist is a 404 before anything is sent, a value that does not fit its @@ -139,6 +139,81 @@ and sends again, once. Resolving by external id goes through The durable spool does not cover this path. `binary::FrameWriter` is public for producers that build frames themselves. +### From a dataframe + +`time_series.insert_table_binary(&batches, "timestamp", &options)` takes Arrow `RecordBatch`es in +the wide shape: a timestamp column and one column per series, named by its external id. In Python +that is `insert_dataframe_binary`, which takes a polars or pandas DataFrame or a pyarrow table as +it is, through the Arrow C stream: + +```python +df = pl.DataFrame({"timestamp": stamps, "pump.temp": temps, "pump.count": counts}) +client.timeseries.insert_dataframe_binary(df) # or a pandas frame with a timestamp index +``` + +Columns are used without copying where they already have their series' type; nulls are "no +point here", so a ragged table works; values are stored as the string path stores them. 100M float +points from a polars frame took about 3 s, against about 75 s through Python lists. What +`to_polars(wide=True)` returns from a read goes back in unchanged. + +## Binary datapoint read + +`time_series.retrieve_datapoints_binary(&filters)` takes the request `retrieve_datapoints` +takes and reads the same frames back from `POST /timeseries/data/list/binary`, so a large read +costs no JSON and no object per point. It returns one `BinaryDatapoints` per series that has +points, in the order of the request's items: `timestamps` in epoch milliseconds and `values` as +one typed column (`BinaryValues::Float(Vec)`, exact `Decimal { unscaled, scale }` for +`numeric` and `decimal32`, two optional columns for `mixed`, or `Aggregates { avg, min, max, +sum }` for an item with `aggregates` and a `granularity`). + +```rust +let mut filter = RetrieveFilter::default(); +filter.external_id = Some("engine.temp".to_string()); +filter.start = Some(Utc.with_ymd_and_hms(2026, 10, 1, 0, 0, 0).unwrap()); +filter.end = Some(Utc.with_ymd_and_hms(2026, 10, 2, 0, 0, 0).unwrap()); +for series in api.time_series.retrieve_datapoints_binary(&DataWrapper::from_vec(vec![filter])).await? { + if let BinaryValues::Float(values) = &series.values { + println!("{}: {} points", series.external_id, values.len()); + } +} +``` + +An item with a `limit` (at most 100 000) returns that many of the oldest points from `start`; +an item without one returns its whole window, so bound a long series with `end`, as the whole +answer is held in memory. There is no cursor: to read on in pieces, send the request again with +`start` one millisecond past the last timestamp you got. A series that does not exist is a 400 for the whole request, where +the JSON read answers empty. A response the server could not finish is a 502 rather than fewer +points, and is safe to send again. + +`time_series.stream_datapoints_binary(&filters)` is the same read handed over as it arrives: a +`Stream` of `BinaryDatapoints`, one per frame, decoded on blocking threads while the rest of the +response downloads, so nothing waits for the whole body and memory holds what the consumer keeps. +A series longer than one frame comes as consecutive items. A refusal is the `Err` before there is +a stream; a response that breaks off is an error item **after** the items that arrived, so a +consumer acting on items as they come has to treat what it took as incomplete. The blocking +client has it as an iterator. `retrieve_datapoints_binary` collects this stream, which is where +all-or-nothing lives. + +In Python, `client.timeseries.retrieve_datapoints_binary(filter_or_list)` on the sync client +returns objects that implement the Arrow PyCapsule stream protocol, so +`polars.DataFrame(series)`, `pyarrow.table(series)` and `pandas.DataFrame.from_arrow(series)` +take them without copying; `series.to_polars()` and `series.to_pandas()` do the same. The package +adds no dependency for this — bring polars, or pyarrow for pandas. +`retrieve_datapoints_binary_stream(filter_or_list)` returns a single-use reader of the same read +as it arrives: iterate it for one `BinaryDatapoints` per frame. Both the buffered result and the +reader turn into one table the same way: + +```python +read = client.timeseries.retrieve_datapoints_binary(filters) # or ..._binary_stream(filters) +long = read.to_polars() # timeseries_id, timestamp, value: one row per point +wide = read.to_polars(wide=True) # timestamp, then one column per series named by external id +read.external_ids # {timeseries_id: external_id} +``` + +The wide table matches series on `timestamp`, leaving nulls where one has no point, and holds +series of different value types side by side. Reading 100M float rows into either costs within +half a second of the download itself. + ## Python bindings `datahub_python_bindings/` wraps this SDK as the Python package `intellistream-datahub-sdk` (import name diff --git a/datahub_python_bindings/Cargo.toml b/datahub_python_bindings/Cargo.toml index 96212f4..b6cbb3a 100644 --- a/datahub_python_bindings/Cargo.toml +++ b/datahub_python_bindings/Cargo.toml @@ -16,6 +16,7 @@ crate-type = ["cdylib"] pyo3 = { version = "0.29", features = ["extension-module", "chrono","chrono-tz","uuid","experimental-inspect", "abi3-py39","multiple-pymethods"] } intellistream-datahub-sdk = { path = ".." } tokio = { version = "1", features = ["full"] } +futures = "0.3" chrono = { version = "0.4", features = ["serde"] } serde = { version = "1", features = ["derive"] } serde_json = "1" @@ -25,3 +26,6 @@ pyo3-async-runtimes = { version = "0.29", features = ["tokio-runtime"] } numpy = "0.29" uuid = { version = "1.22.0", features = ["v4"] } strum = { version = "0.28", features = ["derive"] } +# `ffi` for the Arrow PyCapsule stream; the SDK already depends on both crates. +arrow-array = { version = "59.3", features = ["ffi"] } +arrow-schema = "59.3" diff --git a/datahub_python_bindings/python/intellistream_datahub_sdk/__init__.pyi b/datahub_python_bindings/python/intellistream_datahub_sdk/__init__.pyi index 7a233fa..bfa9780 100644 --- a/datahub_python_bindings/python/intellistream_datahub_sdk/__init__.pyi +++ b/datahub_python_bindings/python/intellistream_datahub_sdk/__init__.pyi @@ -8,7 +8,7 @@ registration in src/lib.rs::intellistream_datahub_sdk() also adds them. from __future__ import annotations import datetime -from typing import Any, Iterable, Iterator, Mapping, Optional, Sequence, Union +from typing import Any, Iterable, Iterator, Literal, Mapping, Optional, Sequence, Union from uuid import UUID @@ -590,6 +590,145 @@ class DatapointsCollectionDatapoints: def id(self) -> int | None: ... +class BinaryDatapointReader: + """A binary read as it arrives, from ``TimeSeriesServiceSync.retrieve_datapoints_binary_stream``. + + Single use, in one of two ways; a second use raises ``RuntimeError``. + + **Iterating** yields one :class:`BinaryDatapoints` per frame, as the frames arrive. A series + longer than one frame comes as several consecutive items, oldest first, so an item is a chunk + of a series rather than all of it; series of different value types mix freely. + + **As an Arrow stream** (``polars.DataFrame(reader)``, :meth:`to_polars`, + ``pyarrow.RecordBatchReader.from_stream(reader)``) it is one long table: ``timeseries_id``, + ``timestamp`` and the value columns, a batch per frame. Every series in the read has to share + those columns, so a read mixing value types (or values and aggregates) raises at the first + batch that differs; read such a read with ``to_polars(wide=True)`` or by iterating. A read + with no rows has no schema and raises + ``ValueError``. The network is waited on with the GIL held, as the C stream protocol gives no + way to release it. + + A response that breaks off raises ``DataHubException`` (502, "cut short") from ``next()`` or + from the Arrow stream, **after** the rows that did arrive were handed out: treat what was + already taken as incomplete. ``retrieve_datapoints_binary`` is the all-or-nothing form. + """ + + def __iter__(self) -> Iterator[BinaryDatapoints]: ... + def __next__(self) -> BinaryDatapoints: ... + @property + def external_ids(self) -> dict[int, str]: + """``{timeseries_id: external_id}`` for the series read so far, in the order they came: + the key to the long table's ``timeseries_id``. Complete once the reader is exhausted.""" + def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: + """The Arrow PyCapsule stream protocol, as the long table. ``requested_schema`` is ignored.""" + def to_polars(self, wide: bool = False) -> Any: + """A ``polars.DataFrame`` of the read; needs polars, and raises ``ImportError`` without it. + + Parameters + ---------- + wide + ``False`` (default): the long table, ``timeseries_id``, ``timestamp`` and the value + columns, one row per point. Every series has to share those columns. + ``True``: one row per timestamp and one column per series, named after its external + id (``{external_id}.{column}`` for a series with several value columns: ``mixed``, or + aggregates), matched on ``timestamp`` with nulls where a series has no point. Series + of different value types sit side by side, each keeping its type. + + Raises + ------ + ValueError + The read returned no rows; or, for the long table, its series do not share their + columns; or, for the wide one, a series has repeated timestamps. + + Notes + ----- + The same table whether the read was buffered or streamed. The wide table costs nothing + beyond the long one when every series has the same timestamps, and a pivot (about 1.7 s + per 100M rows) when they do not. + """ + + +class BinaryReadResult(Sequence[BinaryDatapoints]): + """A buffered binary read, from ``TimeSeriesServiceSync.retrieve_datapoints_binary``. + + A sequence of :class:`BinaryDatapoints`, one per series that had points, in the order of the + request; and the same tables a :class:`BinaryDatapointReader` gives, through + :meth:`to_polars` and the Arrow stream of the long table. + """ + + def __len__(self) -> int: ... + def __getitem__(self, index: int) -> BinaryDatapoints: ... # type: ignore[override] + def __iter__(self) -> Iterator[BinaryDatapoints]: ... + @property + def external_ids(self) -> dict[int, str]: + """``{timeseries_id: external_id}`` for the series in the read: the key to the long + table's ``timeseries_id``.""" + def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: + """The Arrow PyCapsule stream protocol, as the long table. ``requested_schema`` is ignored.""" + def to_polars(self, wide: bool = False) -> Any: + """A ``polars.DataFrame`` of the read; needs polars, and raises ``ImportError`` without it. + + Parameters + ---------- + wide + ``False`` (default): the long table, ``timeseries_id``, ``timestamp`` and the value + columns, one row per point. Every series has to share those columns. + ``True``: one row per timestamp and one column per series, named after its external + id (``{external_id}.{column}`` for a series with several value columns: ``mixed``, or + aggregates), matched on ``timestamp`` with nulls where a series has no point. Series + of different value types sit side by side, each keeping its type. + + Raises + ------ + ValueError + The read returned no rows; or, for the long table, its series do not share their + columns; or, for the wide one, a series has repeated timestamps. + + Notes + ----- + The same table whether the read was buffered or streamed. The wide table costs nothing + beyond the long one when every series has the same timestamps, and a pivot (about 1.7 s + per 100M rows) when they do not. + """ + + +class BinaryDatapoints: + """One series' rows from ``TimeSeriesServiceSync.retrieve_datapoints_binary``, as an Arrow table. + + Implements the Arrow PyCapsule stream protocol, so ``polars.DataFrame(series)``, + ``pyarrow.table(series)`` and ``pandas.DataFrame.from_arrow(series)`` (pandas 3) take it + without copying. Every call makes a new stream, so one object converts more than once. + + Columns are ``timestamp`` (``timestamp[ms, UTC]``; the bucket start for aggregates) and then, + by value type: ``value`` as int64 (``bigint``), float64 (``float``), float32 (``float32``), + ``decimal128(18, 6)`` (``numeric``), ``decimal128(9, 4)`` (``decimal32``) or utf8 (``text``); + nullable ``value_numeric`` and ``value_text`` for ``mixed``, each row setting one; or + ``avg``, ``min``, ``max`` and ``sum`` as float64 when :attr:`is_aggregates`. + """ + + @property + def id(self) -> int: ... + @property + def external_id(self) -> str: ... + @property + def value_type(self) -> str: + """The series' value type: ``bigint``, ``float``, ``float32``, ``numeric``, ``decimal32``, + ``text`` or ``mixed``. An aggregate read reports the series' own type.""" + @property + def is_aggregates(self) -> bool: + """Whether the rows are buckets with ``avg``, ``min``, ``max`` and ``sum``.""" + @property + def columns(self) -> list[str]: ... + def __len__(self) -> int: ... + def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: + """The Arrow PyCapsule stream protocol. ``requested_schema`` is ignored.""" + def to_polars(self) -> Any: + """A ``polars.DataFrame``. Needs polars; raises ``ImportError`` without it.""" + def to_pandas(self) -> Any: + """A ``pandas.DataFrame``, converted through pyarrow. Needs pandas and pyarrow; raises + ``ImportError`` without pyarrow.""" + + class RetrieveFilter: def __init__( self, @@ -733,13 +872,72 @@ class TimeSeriesServiceSync: self, input: list[DatapointsCollectionString], zstd_level: int | None = None, + compression: Literal["payload", "buffers"] | None = None, ) -> list[str]: ... + def insert_dataframe_binary( + self, + table: Any, + timestamp: str | None = None, + zstd_level: int | None = None, + compression: Literal["payload", "buffers"] | None = None, + ) -> int: + """Insert a wide table through ``POST /timeseries/data/binary``, straight from Arrow. + + Parameters + ---------- + table + A polars or pandas DataFrame (pandas 2.2 or later, which exports through pyarrow), a + pyarrow Table or RecordBatchReader: anything with ``__arrow_c_stream__``. Every column + but the timestamp is a series, named by its external id. + timestamp + The timestamp column. If not given: the pandas index when it is a timestamp, else a + column named ``timestamp``, else the only timestamp-typed column. It has to carry a time + zone and whole milliseconds, or be ``Int64`` epoch milliseconds. + zstd_level + 1, 3 or 9; 3 by default. + compression + ``"buffers"`` (the default) compresses each Arrow buffer inside the frame, the form a + binary read answers with, which the server stores without decompressing; + ``"payload"`` zstd-compresses each frame's Arrow stream whole. The server stores the + same points either way. + + Returns + ------- + int + The points sent. + + Raises + ------ + DataHubException + 404 naming a column that is not a series you can read, before anything is sent; 422 + for a naive or sub-millisecond timestamp, or a column its series cannot hold (a float + column for a ``bigint`` series, say). + ValueError + No timestamp column, or several and none named. + + Notes + ----- + The table is read without copying it; timestamps get one pass to milliseconds and the + frames are compressed on several threads, so no object or string is made per value. + 100M float points from a polars frame took about 3 s, against about 75 s through + :meth:`insert_from_lists_binary`. + + A null is "no point at this timestamp", so a ragged table works (pandas exports NaN as + null; polars keeps NaN, which is then stored as a value). Rows need not be sorted, and a + timestamp repeated within a series keeps its last value. Values are stored as the string + path stores them, decimals rounded half-up and ``decimal32`` clamped, with one deliberate + difference: a ``mixed`` series takes a numeric column as numbers and a string column as + text, where the string path classifies each value. ``to_polars(wide=True)`` on a read + gives back exactly this shape. A table larger than one request is not atomic: a refusal + can leave the earlier requests stored. + """ def insert_from_lists_binary( self, timestamps: list[datetime.datetime], values: list[float], ts: Identifiable, zstd_level: int | None = None, + compression: Literal["payload", "buffers"] | None = None, ) -> list[str]: ... def insert_from_lists( self, @@ -748,6 +946,70 @@ class TimeSeriesServiceSync: ts: Identifiable, ) -> list[str]: ... def retrieve_datapoints(self, input: RetrieveFilter) -> list[DatapointsCollectionDatapoints]: ... + def retrieve_datapoints_binary( + self, input: RetrieveFilter | list[RetrieveFilter] + ) -> BinaryReadResult: + """Read datapoints through ``POST /timeseries/data/list/binary``: the request + ``retrieve_datapoints`` takes, answered as compressed Arrow frames instead of JSON. + + Parameters + ---------- + input + One filter or a list. Each needs a ``start`` or an ``end``. With a ``limit`` (at most + 100 000) it returns that many of the oldest points from there; without one it returns + every point in its window, so bound a long series with ``end``, as the whole answer is + held in memory. With ``aggregates`` and a ``granularity`` + the series comes back as buckets carrying all four aggregates, whichever were named. + + Returns + ------- + BinaryReadResult + A sequence of :class:`BinaryDatapoints`, one per series that has points in its window, + in the order of ``input`` (a series with none is left out rather than returned empty), + with ``to_polars(wide=False)`` for one table of them all. + + Raises + ------ + DataHubException + 400 for a ``cursor`` (there is no paging: send the request again with ``start`` one + millisecond past the last timestamp read), for aggregates on a text series, and for + a series that does not exist or that you cannot read, where the JSON read answers + empty. 502 when the response was cut short; sending it again is safe. + + Notes + ----- + Sync client only, like the binary insert. The frames are decoded while the rest of the + response downloads; for rows as they arrive, or a window too large to hold, use + :meth:`retrieve_datapoints_binary_stream`. + """ + def retrieve_datapoints_binary_stream( + self, input: RetrieveFilter | list[RetrieveFilter] + ) -> BinaryDatapointReader: + """The read :meth:`retrieve_datapoints_binary` makes, handed over as it arrives. + + Parameters + ---------- + input + As for :meth:`retrieve_datapoints_binary`. + + Returns + ------- + BinaryDatapointReader + Single use: iterate it for one :class:`BinaryDatapoints` per frame, or read it as an + Arrow stream of one long table. + + Raises + ------ + DataHubException + A refusal (400 for an unknown series, a ``cursor``, aggregates on a text series) + raises here, before there is a reader. A response that breaks off raises from the + reader instead, after the rows that did arrive: see :class:`BinaryDatapointReader`. + + Notes + ----- + Sync client only. Holding the reader holds the connection and the server's query open, + so read it through rather than leaving it half-consumed. + """ def delete_datapoints(self, input: list[DeleteFilter]) -> None: ... def retrieve_latest_datapoints( self, input: list[Identifiable] diff --git a/datahub_python_bindings/python/intellistream_datahub_sdk/_tables.py b/datahub_python_bindings/python/intellistream_datahub_sdk/_tables.py new file mode 100644 index 0000000..81fba15 --- /dev/null +++ b/datahub_python_bindings/python/intellistream_datahub_sdk/_tables.py @@ -0,0 +1,60 @@ +"""Polars tables from a binary read, shared by the buffered result and the streamed reader so the +two cannot drift: the long table, and its wide form with one column per series.""" + + +def _polars(method): + try: + import polars + except ImportError: + raise ImportError(f"{method}() needs polars; install it with `pip install polars`") from None + return polars + + +def to_polars(read, wide): + """`read` is a `BinaryReadResult` or a `BinaryDatapointReader`: an Arrow stream of the long + table, and an iterable of `BinaryDatapoints` (whole series, or a reader's per-frame chunks).""" + pl = _polars("to_polars") + if not wide: + return pl.DataFrame(read) + return _wide(pl, read) + + +def _wide(pl, parts): + # Chunks of one series arrive consecutively; a dict keyed by id keeps the read's order. + series = {} + for part in parts: + external_id, frames = series.setdefault(part.id, (part.external_id, [])) + frames.append(part.to_polars()) + if not series: + raise ValueError("the read returned no rows, so there are no columns for a wide table") + + columns = [] + for external_id, frames in series.values(): + frame = pl.concat(frames, rechunk=False) if len(frames) > 1 else frames[0] + if not frame["timestamp"].is_sorted() or frame["timestamp"].n_unique() != frame.height: + raise ValueError( + f"series {external_id} has repeated timestamps, as when one series is asked for " + f"twice with overlapping windows; a wide table needs one value per timestamp" + ) + values = [c for c in frame.columns if c != "timestamp"] + names = {values[0]: external_id} if len(values) == 1 else {c: f"{external_id}.{c}" for c in values} + columns.append(frame.rename(names)) + + first = columns[0]["timestamp"] + if all(c.height == first.len() and c["timestamp"].equals(first) for c in columns[1:]): + # Every series on the same timestamps: the columns side by side, nothing to match up. + return columns[0].hstack([s for c in columns[1:] for s in c.drop("timestamp").get_columns()]) + + single_values = [c.columns[1:] for c in columns] + if all(len(v) == 1 for v in single_values) and len({c.dtypes[1] for c in columns}) == 1: + # One value column of one type: a pivot on a numeric key is the fastest way to match rows. + names = [v[0] for v in single_values] + long = pl.concat( + [c.rename({n: "value"}).with_columns(pl.lit(i, pl.Int32).alias("series")) for i, (c, n) in enumerate(zip(columns, names))], + rechunk=False, + ) + wide = long.pivot(on="series", index="timestamp", values="value").sort("timestamp") + return wide.rename({str(i): n for i, n in enumerate(names)}).select(["timestamp", *names]) + + # Different value types or several value columns: matched on timestamp, each keeping its type. + return pl.concat(columns, how="align") diff --git a/datahub_python_bindings/src/timeseries/binary_read.rs b/datahub_python_bindings/src/timeseries/binary_read.rs new file mode 100644 index 0000000..a30b78c --- /dev/null +++ b/datahub_python_bindings/src/timeseries/binary_read.rs @@ -0,0 +1,566 @@ +use std::sync::Arc; + +use arrow_array::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream}; +use std::sync::Mutex; + +use arrow_array::{ + ArrayRef, Decimal128Array, Float32Array, Float64Array, Int64Array, + RecordBatch, RecordBatchIterator, RecordBatchReader, StringArray, TimestampMillisecondArray, +}; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef, TimeUnit}; +use futures::StreamExt; +use intellistream_datahub_sdk::timeseries::{ + BinaryDatapointStream, BinaryDatapoints, BinaryValues, DatapointValueType, +}; +use tokio::runtime::Runtime; +use pyo3::exceptions::{PyImportError, PyRuntimeError, PyStopIteration, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyDict, PyList}; + +/// One series' rows from `retrieve_datapoints_binary`, as an Arrow table: `timestamp` and the +/// value columns of its type. Read it with anything that takes the Arrow PyCapsule protocol. +#[pyclass(module = "intellistream_datahub_sdk", name = "BinaryDatapoints", frozen)] +pub struct PyBinaryDatapoints { + id: u64, + external_id: String, + value_type: DatapointValueType, + aggregates: bool, + batch: RecordBatch, +} + +/// A series' `timestamp` and value columns, as Arrow arrays that take over its vectors. +fn columns( + value_type: DatapointValueType, + timestamps: Vec, + values: BinaryValues, +) -> Result<(Vec, Vec), ArrowError> { + let timestamp = Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + false, + ); + let mut fields = vec![timestamp]; + let mut columns: Vec = + vec![Arc::new(TimestampMillisecondArray::from(timestamps).with_timezone("UTC"))]; + let mut push = |name: &str, nullable: bool, array: ArrayRef| { + fields.push(Field::new(name, array.data_type().clone(), nullable)); + columns.push(array); + }; + match values { + BinaryValues::Bigint(v) => push("value", false, Arc::new(Int64Array::from(v))), + BinaryValues::Float(v) => push("value", false, Arc::new(Float64Array::from(v))), + BinaryValues::Float32(v) => push("value", false, Arc::new(Float32Array::from(v))), + BinaryValues::Decimal { unscaled, scale } => { + let precision = if value_type == DatapointValueType::Decimal32 { 9 } else { 18 }; + let array = Decimal128Array::from(unscaled).with_precision_and_scale(precision, scale)?; + push("value", false, Arc::new(array)) + } + BinaryValues::Text(v) => push("value", false, Arc::new(StringArray::from(v))), + BinaryValues::Mixed { numeric, text } => { + push("value_numeric", true, Arc::new(Float64Array::from(numeric))); + push("value_text", true, Arc::new(StringArray::from(text))); + } + BinaryValues::Aggregates { avg, min, max, sum } => { + push("avg", false, Arc::new(Float64Array::from(avg))); + push("min", false, Arc::new(Float64Array::from(min))); + push("max", false, Arc::new(Float64Array::from(max))); + push("sum", false, Arc::new(Float64Array::from(sum))); + } + other => { + return Err(ArrowError::NotYetImplemented(format!( + "binary values {other:?} have no Arrow mapping" + ))) + } + } + Ok((fields, columns)) +} + +impl TryFrom for PyBinaryDatapoints { + type Error = ArrowError; + + fn try_from(series: BinaryDatapoints) -> Result { + let aggregates = series.is_aggregates(); + let (fields, columns) = columns(series.value_type, series.timestamps, series.values)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Self { + id: series.id, + external_id: series.external_id, + value_type: series.value_type, + aggregates, + batch, + }) + } +} + +/// One frame's rows as a batch of the long table: `timeseries_id` ahead of the series' own +/// columns. No external id column: as a dictionary column polars spent 0.8 s per 100M rows making +/// it categorical, more than the rest of the conversion; `external_ids` maps the ids instead. +fn long_batch(series: BinaryDatapoints) -> Result { + let id = Arc::new(Int64Array::from(vec![series.id as i64; series.len()])) as ArrayRef; + let (fields, columns) = columns(series.value_type, series.timestamps, series.values)?; + let mut all_fields = vec![Field::new("timeseries_id", DataType::Int64, false)]; + all_fields.extend(fields); + let mut all_columns = vec![id]; + all_columns.extend(columns); + RecordBatch::try_new(Arc::new(Schema::new(all_fields)), all_columns) +} + +#[pymethods] +impl PyBinaryDatapoints { + #[getter] + fn id(&self) -> u64 { + self.id + } + + #[getter] + fn external_id(&self) -> &str { + &self.external_id + } + + /// The series' value type, as a timeseries read reports it: `float`, `bigint`, `text`, ... + #[getter] + fn value_type(&self) -> String { + format!("{:?}", self.value_type).to_lowercase() + } + + /// Whether the rows are buckets with `avg`, `min`, `max` and `sum` rather than values. + #[getter] + fn is_aggregates(&self) -> bool { + self.aggregates + } + + /// The column names, in order. + #[getter] + fn columns(&self) -> Vec { + self.batch.schema().fields().iter().map(|f| f.name().clone()).collect() + } + + fn __len__(&self) -> usize { + self.batch.num_rows() + } + + fn __repr__(&self) -> String { + format!( + "BinaryDatapoints(external_id={:?}, value_type={:?}, rows={}{})", + self.external_id, + self.value_type(), + self.batch.num_rows(), + if self.aggregates { ", aggregates" } else { "" } + ) + } + + /// The Arrow PyCapsule stream protocol. Each call is a new stream over the same buffers, so + /// the object can be converted more than once. `requested_schema` is ignored, as the protocol + /// allows: the columns come in their own types. + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_stream__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult> { + let _ = requested_schema; + let reader = RecordBatchIterator::new(vec![Ok(self.batch.clone())], self.batch.schema()); + let stream = FFI_ArrowArrayStream::new(Box::new(reader)); + PyCapsule::new_with_value(py, stream, c"arrow_array_stream") + } + + /// A `polars.DataFrame`, without copying. Needs polars. + fn to_polars<'py>(slf: &Bound<'py, Self>) -> PyResult> { + let polars = import(slf.py(), "polars", "to_polars")?; + polars.call_method1("DataFrame", (slf,)) + } + + /// A `pandas.DataFrame`. Needs pandas and pyarrow, which pandas converts Arrow through. + fn to_pandas<'py>(slf: &Bound<'py, Self>) -> PyResult> { + let py = slf.py(); + let pyarrow = import(py, "pyarrow", "to_pandas")?; + pyarrow.call_method1("table", (slf,))?.call_method0("to_pandas") + } +} + +/// A binary read as it arrives, from `retrieve_datapoints_binary_stream`. Single use: read it +/// either by iterating, one `BinaryDatapoints` per frame, or as an Arrow stream of the long table. +#[pyclass(module = "intellistream_datahub_sdk", name = "BinaryDatapointReader", frozen)] +pub struct PyBinaryDatapointReader { + runtime: Arc, + state: Mutex, + seen: Seen, +} + +/// The series read so far, by id, in the order they first came. +type Seen = Arc>>; + +fn note(seen: &Seen, series: &BinaryDatapoints) { + let mut seen = seen.lock().unwrap(); + if !seen.iter().any(|(id, _)| *id == series.id) { + seen.push((series.id, series.external_id.clone())); + } +} + +enum ReaderState { + Unread(BinaryDatapointStream), + Iterating(BinaryDatapointStream), + Done, +} + +impl PyBinaryDatapointReader { + pub(crate) fn new(runtime: Arc, stream: BinaryDatapointStream) -> Self { + Self { + runtime, + state: Mutex::new(ReaderState::Unread(stream)), + seen: Arc::default(), + } + } +} + +#[pymethods] +impl PyBinaryDatapointReader { + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + /// The next frame's rows: a chunk of one series, which may continue in the next item. + fn __next__(&self, py: Python<'_>) -> PyResult { + let mut state = self.state.lock().unwrap(); + let mut stream = match std::mem::replace(&mut *state, ReaderState::Done) { + ReaderState::Unread(stream) | ReaderState::Iterating(stream) => stream, + ReaderState::Done => return Err(PyStopIteration::new_err(())), + }; + let item = py.detach(|| self.runtime.block_on(stream.next())); + match item { + Some(Ok(part)) => { + *state = ReaderState::Iterating(stream); + note(&self.seen, &part); + PyBinaryDatapoints::try_from(part).map_err(|e| PyValueError::new_err(e.to_string())) + } + Some(Err(error)) => Err(crate::datahub_err(error)), + None => Err(PyStopIteration::new_err(())), + } + } + + /// `{timeseries_id: external_id}` for the series read so far, in the order they came: the key + /// for the long table's `timeseries_id`. Complete once the reader is exhausted. + #[getter] + fn external_ids<'py>(&self, py: Python<'py>) -> PyResult> { + let ids = PyDict::new(py); + for (id, external_id) in self.seen.lock().unwrap().iter() { + ids.set_item(id, external_id)?; + } + Ok(ids) + } + + /// The Arrow PyCapsule stream protocol, as one long table: `timeseries_id`, `timestamp` and + /// the value columns, a batch per frame. Every series in the read has to + /// share those columns; one that does not is an error at its batch. A read with no rows has + /// no schema to declare and raises. + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_stream__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult> { + let _ = requested_schema; + let mut stream = { + let mut state = self.state.lock().unwrap(); + match std::mem::replace(&mut *state, ReaderState::Done) { + ReaderState::Unread(stream) => stream, + other => { + *state = other; + return Err(PyRuntimeError::new_err( + "this reader has already been read; a binary read streams once", + )); + } + } + }; + let first = match py.detach(|| self.runtime.block_on(stream.next())) { + Some(Ok(part)) => { + note(&self.seen, &part); + long_batch(part).map_err(|e| PyValueError::new_err(e.to_string()))? + } + Some(Err(error)) => return Err(crate::datahub_err(error)), + None => { + return Err(PyValueError::new_err( + "the read returned no rows, so there is no schema for an Arrow stream", + )) + } + }; + let reader = LongTable { + schema: first.schema(), + first: Some(first), + stream, + runtime: self.runtime.clone(), + seen: self.seen.clone(), + }; + let ffi = FFI_ArrowArrayStream::new(Box::new(reader)); + PyCapsule::new_with_value(py, ffi, c"arrow_array_stream") + } + + /// A `polars.DataFrame`: the long table, or with `wide` one column per series, matched on + /// timestamp. Needs polars. Reads the reader, like any other use of it. + #[pyo3(signature = (wide=false))] + fn to_polars<'py>(slf: &Bound<'py, Self>, wide: bool) -> PyResult> { + tables_to_polars(slf.as_any(), wide) + } +} + +fn tables_to_polars<'py>(read: &Bound<'py, PyAny>, wide: bool) -> PyResult> { + read.py() + .import("intellistream_datahub_sdk._tables")? + .call_method1("to_polars", (read, wide)) +} + +/// A buffered binary read, from `retrieve_datapoints_binary`: the series it returned, one +/// `BinaryDatapoints` each, which it behaves as a list of, and the same tables the streamed +/// reader gives. +#[pyclass(module = "intellistream_datahub_sdk", name = "BinaryReadResult", frozen, sequence)] +pub struct PyBinaryReadResult { + series: Vec>, +} + +impl PyBinaryReadResult { + pub(crate) fn new(py: Python<'_>, series: Vec) -> PyResult { + let series = series + .into_iter() + .map(|s| { + let s = PyBinaryDatapoints::try_from(s).map_err(|e| PyValueError::new_err(e.to_string()))?; + Py::new(py, s) + }) + .collect::>()?; + Ok(Self { series }) + } +} + +#[pymethods] +impl PyBinaryReadResult { + fn __len__(&self) -> usize { + self.series.len() + } + + fn __getitem__(&self, py: Python<'_>, index: isize) -> PyResult> { + let len = self.series.len() as isize; + let at = if index < 0 { index + len } else { index }; + if at < 0 || at >= len { + return Err(pyo3::exceptions::PyIndexError::new_err("series index out of range")); + } + Ok(self.series[at as usize].clone_ref(py)) + } + + fn __iter__<'py>(&self, py: Python<'py>) -> PyResult> { + Ok(PyList::new(py, &self.series)?.into_any().try_iter()?.into_any()) + } + + fn __repr__(&self) -> String { + let rows: usize = self.series.iter().map(|s| s.get().batch.num_rows()).sum(); + format!("BinaryReadResult(series={}, rows={rows})", self.series.len()) + } + + /// `{timeseries_id: external_id}` for the series in the read, in order: the key to the long + /// table's `timeseries_id`. + #[getter] + fn external_ids<'py>(&self, py: Python<'py>) -> PyResult> { + let ids = PyDict::new(py); + for s in &self.series { + ids.set_item(s.get().id, &s.get().external_id)?; + } + Ok(ids) + } + + /// The Arrow PyCapsule stream protocol, as the long table: `timeseries_id`, `timestamp` and + /// the value columns, a batch per series, which all have to share those columns. A read with + /// no rows has no schema and raises. + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_stream__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult> { + let _ = requested_schema; + let mut batches = Vec::with_capacity(self.series.len()); + for s in &self.series { + let s = s.get(); + let id = Arc::new(Int64Array::from(vec![s.id as i64; s.batch.num_rows()])) as ArrayRef; + let mut fields = vec![Arc::new(Field::new("timeseries_id", DataType::Int64, false))]; + fields.extend(s.batch.schema().fields().iter().cloned()); + let mut columns = vec![id]; + columns.extend(s.batch.columns().iter().cloned()); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + if let Some(first) = batches.first() { + let first: &RecordBatch = first; + if first.schema().fields() != batch.schema().fields() { + return Err(PyValueError::new_err(format!( + "series {} has columns {:?}, the table's are {:?}; series of different value \ + types do not share one long table, but to_polars(wide=True) holds them", + s.external_id, + batch.schema().fields().iter().map(|f| f.name()).collect::>(), + first.schema().fields().iter().map(|f| f.name()).collect::>(), + ))); + } + } + batches.push(batch); + } + let Some(first) = batches.first() else { + return Err(PyValueError::new_err( + "the read returned no rows, so there is no schema for an Arrow stream", + )); + }; + let schema = first.schema(); + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); + PyCapsule::new_with_value(py, FFI_ArrowArrayStream::new(Box::new(reader)), c"arrow_array_stream") + } + + /// A `polars.DataFrame`: the long table, or with `wide` one column per series, matched on + /// timestamp. Needs polars. + #[pyo3(signature = (wide=false))] + fn to_polars<'py>(slf: &Bound<'py, Self>, wide: bool) -> PyResult> { + tables_to_polars(slf.as_any(), wide) + } +} + +/// The long table behind `__arrow_c_stream__`: the frame peeked for its schema, then the rest as +/// the consumer asks. Waits for the network with the GIL held, as the C stream gives no way to +/// release it; the reader belongs to the sync client. +struct LongTable { + schema: SchemaRef, + first: Option, + stream: BinaryDatapointStream, + runtime: Arc, + seen: Seen, +} + +impl Iterator for LongTable { + type Item = Result; + + fn next(&mut self) -> Option { + if let Some(first) = self.first.take() { + return Some(Ok(first)); + } + let part = match self.runtime.block_on(self.stream.next())? { + Ok(part) => part, + Err(error) => { + let message = format!("{}: {}", error.get_status(), error.get_message()); + return Some(Err(ArrowError::ExternalError(message.into()))); + } + }; + note(&self.seen, &part); + let external_id = part.external_id.clone(); + let batch = match long_batch(part) { + Ok(batch) => batch, + Err(error) => return Some(Err(error)), + }; + if batch.schema().fields() != self.schema.fields() { + return Some(Err(ArrowError::SchemaError(format!( + "series {external_id} has columns {:?}, the stream's are {:?}; read series of \ + different value types with to_polars(wide=True), or by iterating", + batch.schema().fields().iter().map(|f| f.name()).collect::>(), + self.schema.fields().iter().map(|f| f.name()).collect::>(), + )))); + } + Some(Ok(batch)) + } +} + +impl RecordBatchReader for LongTable { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +fn import<'py>(py: Python<'py>, module: &str, method: &str) -> PyResult> { + py.import(module).map_err(|_| { + PyImportError::new_err(format!( + "BinaryDatapoints.{method}() needs {module}; install it with `pip install {module}`" + )) + }) +} + + + +/// The table behind any object with `__arrow_c_stream__`, imported without copying: its batches, +/// and the names pandas records as its index columns. +pub(crate) fn import_table(table: &Bound<'_, PyAny>) -> PyResult<(Vec, Vec)> { + if !table.hasattr("__arrow_c_stream__")? { + return Err(pyo3::exceptions::PyTypeError::new_err( + "expected a table with __arrow_c_stream__: a polars or pandas DataFrame (pandas 2.2 or \ + later, with pyarrow), or a pyarrow Table or RecordBatchReader", + )); + } + let capsule = table.call_method0("__arrow_c_stream__")?; + let capsule = capsule.cast::()?; + // Checks the capsule's name is arrow_array_stream before handing out the pointer. + let pointer = capsule.pointer_checked(Some(c"arrow_array_stream"))?.cast::(); + // Moves the stream out of the capsule and leaves it released, as the protocol requires. + let reader = unsafe { ArrowArrayStreamReader::from_raw(pointer.as_ptr()) } + .map_err(|e| PyValueError::new_err(e.to_string()))?; + let schema = reader.schema(); + let batches = reader + .collect::, _>>() + .map_err(|e| PyValueError::new_err(e.to_string()))?; + let index = schema + .metadata() + .get("pandas") + .and_then(|m| serde_json::from_str::(m).ok()) + .and_then(|m| m.get("index_columns").cloned()) + .and_then(|c| c.as_array().cloned()) + .map(|c| c.iter().filter_map(|v| v.as_str().map(str::to_string)).collect()) + .unwrap_or_default(); + Ok((batches, index)) +} + +/// The timestamp column: as named, else the pandas index if it is a timestamp, else `timestamp`, +/// else the one timestamp column there is. +pub(crate) fn timestamp_column(schema: &Schema, named: Option, index: &[String]) -> PyResult { + if let Some(name) = named { + return Ok(name); + } + let is_time = |name: &str| { + schema + .field_with_name(name) + .is_ok_and(|f| matches!(f.data_type(), DataType::Timestamp(_, _))) + }; + if let [only] = index { + if is_time(only) { + return Ok(only.clone()); + } + } + if schema.field_with_name("timestamp").is_ok() { + return Ok("timestamp".to_string()); + } + let times: Vec<&str> = schema + .fields() + .iter() + .filter(|f| matches!(f.data_type(), DataType::Timestamp(_, _))) + .map(|f| f.name().as_str()) + .collect(); + match times[..] { + [only] => Ok(only.to_string()), + [] => Err(PyValueError::new_err( + "the table has no timestamp column; name one with timestamp=, or give it a \ + timestamp index", + )), + _ => Err(PyValueError::new_err(format!( + "the table has several timestamp columns ({}); name one with timestamp=", + times.join(", ") + ))), + } +} + +/// The batches without the pandas index columns that are not the timestamp: an index is not a +/// series. +pub(crate) fn without_index(batches: Vec, index: &[String], timestamp: &str) -> PyResult> { + let Some(first) = batches.first() else { return Ok(batches) }; + let schema = first.schema(); + let keep: Vec = (0..schema.fields().len()) + .filter(|&i| { + let name = schema.field(i).name(); + name == timestamp || !index.iter().any(|x| x == name) + }) + .collect(); + if keep.len() == schema.fields().len() { + return Ok(batches); + } + batches + .iter() + .map(|b| b.project(&keep).map_err(|e| PyValueError::new_err(e.to_string()))) + .collect() +} diff --git a/datahub_python_bindings/src/timeseries/mod.rs b/datahub_python_bindings/src/timeseries/mod.rs index 5cd6d94..e6c14b9 100644 --- a/datahub_python_bindings/src/timeseries/mod.rs +++ b/datahub_python_bindings/src/timeseries/mod.rs @@ -33,6 +33,7 @@ use std::sync::Arc; use strum::{Display, EnumString}; pub mod async_service; +pub mod binary_read; mod construction; pub mod datapoints; pub mod general; @@ -469,6 +470,9 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -477,16 +481,28 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { Ok(()) } -/// Options for the binary ingest bindings. zstd level 1, 3 or 9; 9 is the default, as in the -/// Rust and Java SDKs, because the client is what pays for it. +/// Options for the binary ingest bindings. zstd level 1, 3 or 9; 3 is the default, as in the +/// Rust SDK. `compression` is `"buffers"` +/// (the default) or `"payload"`. pub(crate) fn binary_options( zstd_level: Option, -) -> intellistream_datahub_sdk::timeseries::BinaryIngestOptions { - use intellistream_datahub_sdk::timeseries::BinaryIngestOptions; - match zstd_level { - Some(level) => BinaryIngestOptions::new().zstd_level(level), - None => BinaryIngestOptions::default(), - } + compression: Option<&str>, +) -> PyResult { + use intellistream_datahub_sdk::timeseries::{BinaryIngestOptions, FrameCompression}; + let mut options = BinaryIngestOptions::default(); + if let Some(level) = zstd_level { + options = options.zstd_level(level); + } + options.compression = match compression { + None | Some("buffers") => FrameCompression::Buffers, + Some("payload") => FrameCompression::Payload, + Some(other) => { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "compression {other:?} is not \"buffers\" or \"payload\"" + ))) + } + }; + Ok(options) } /// Parallel timestamp and value sequences for one series, which is the shape a DataFrame's diff --git a/datahub_python_bindings/src/timeseries/sync_service.rs b/datahub_python_bindings/src/timeseries/sync_service.rs index 0d84537..af8a83e 100644 --- a/datahub_python_bindings/src/timeseries/sync_service.rs +++ b/datahub_python_bindings/src/timeseries/sync_service.rs @@ -4,11 +4,18 @@ use crate::timeseries::datapoints::{ }; use crate::{DatahubIdentity, Identifiable}; use crate::{PyIdCollection, PyRetrieveFilter}; +use crate::timeseries::binary_read::{self, PyBinaryDatapointReader, PyBinaryReadResult}; use intellistream_datahub_sdk::generic::{DataWrapper, IdAndExtId}; use intellistream_datahub_sdk::{ApiService, TimeSeriesUpdateCollection}; use pyo3_async_runtimes::tokio::future_into_py; use std::sync::Arc; +#[derive(FromPyObject)] +enum RetrieveFilters { + One(PyRetrieveFilter), + Many(Vec), +} + #[pyclass(module = "intellistream_datahub_sdk", name = "TimeSeriesServiceSync")] pub struct PyTimeSeriesServiceSync { pub api_service: Arc, @@ -269,19 +276,21 @@ impl PyTimeSeriesServiceSync { }) } /// `POST /timeseries/data/binary`: the same collections as `insert_datapoints`, sent as - /// zstd-compressed Arrow frames. `zstd_level` is 1, 3 or 9 and defaults to 9. - #[pyo3(signature = (input, zstd_level=None))] + /// zstd-compressed Arrow frames. `zstd_level` is 1, 3 or 9 and defaults to 3; `compression` is + /// `"buffers"` (default) or `"payload"`. + #[pyo3(signature = (input, zstd_level=None, compression=None))] fn insert_datapoints_binary<'py>( &self, py: Python<'py>, input: Vec, zstd_level: Option, + compression: Option<&str>, ) -> PyResult> { let service = self.api_service.clone(); let vec: Vec> = input.into_iter().map(|item| item.into()).collect(); let wrapper = DataWrapper::>::from_vec(vec); - let options = binary_options(zstd_level); + let options = binary_options(zstd_level, compression)?; py.detach(|| { let result = self .runtime @@ -291,9 +300,35 @@ impl PyTimeSeriesServiceSync { }) } + /// `POST /timeseries/data/binary` from a table: a polars or pandas DataFrame, or anything else + /// with `__arrow_c_stream__`, one column per series named by its external id beside a + /// timestamp column (or index). Returns the number of points sent. + #[pyo3(signature = (table, timestamp=None, zstd_level=None, compression=None))] + fn insert_dataframe_binary<'py>( + &self, + py: Python<'py>, + table: Bound<'py, PyAny>, + timestamp: Option, + zstd_level: Option, + compression: Option<&str>, + ) -> PyResult { + let (batches, index) = binary_read::import_table(&table)?; + let Some(first) = batches.first() else { return Ok(0) }; + let timestamp = binary_read::timestamp_column(&first.schema(), timestamp, &index)?; + let batches = binary_read::without_index(batches, &index, ×tamp)?; + let service = self.api_service.clone(); + let options = binary_options(zstd_level, compression)?; + let sent = py.detach(|| { + self.runtime + .block_on(service.time_series.insert_table_binary(&batches, ×tamp, &options)) + .map_err(crate::datahub_err) + })?; + Ok(sent.points) + } + /// The binary twin of `insert_from_lists`: parallel timestamp and value sequences for one /// series, which is the shape a DataFrame column pair arrives in. - #[pyo3(signature = (timestamps, values, ts, zstd_level=None))] + #[pyo3(signature = (timestamps, values, ts, zstd_level=None, compression=None))] fn insert_from_lists_binary<'py>( &self, py: Python<'py>, @@ -301,11 +336,12 @@ impl PyTimeSeriesServiceSync { values: Vec, ts: Identifiable, zstd_level: Option, + compression: Option<&str>, ) -> PyResult> { let service = self.api_service.clone(); let collection = lists_to_collection(timestamps, values, ts)?; let wrapper = DataWrapper::>::from_vec(vec![collection]); - let options = binary_options(zstd_level); + let options = binary_options(zstd_level, compression)?; py.detach(|| { let result = self .runtime @@ -355,6 +391,46 @@ impl PyTimeSeriesServiceSync { Ok(result) }) } + /// `POST /timeseries/data/list/binary`: the request `retrieve_datapoints` takes, one filter or + /// a list, answered as Arrow frames. One `BinaryDatapoints` per series that has points. + fn retrieve_datapoints_binary<'py>( + &self, + py: Python<'py>, + input: RetrieveFilters, + ) -> PyResult { + let service = self.api_service.clone(); + let filters: Vec = match input { + RetrieveFilters::One(f) => vec![f.into()], + RetrieveFilters::Many(fs) => fs.into_iter().map(Into::into).collect(), + }; + let wrapper = DataWrapper::::from_vec(filters); + let series = py.detach(|| { + self.runtime + .block_on(service.time_series.retrieve_datapoints_binary(&wrapper)) + .map_err(crate::datahub_err) + })?; + PyBinaryReadResult::new(py, series) + } + /// The read `retrieve_datapoints_binary` makes, handed over as it arrives: a single-use + /// `BinaryDatapointReader`, iterated frame by frame or read as an Arrow stream. + fn retrieve_datapoints_binary_stream<'py>( + &self, + py: Python<'py>, + input: RetrieveFilters, + ) -> PyResult { + let service = self.api_service.clone(); + let filters: Vec = match input { + RetrieveFilters::One(f) => vec![f.into()], + RetrieveFilters::Many(fs) => fs.into_iter().map(Into::into).collect(), + }; + let wrapper = DataWrapper::::from_vec(filters); + let stream = py.detach(|| { + self.runtime + .block_on(service.time_series.stream_datapoints_binary(&wrapper)) + .map_err(crate::datahub_err) + })?; + Ok(PyBinaryDatapointReader::new(self.runtime.clone(), stream)) + } fn delete_datapoints<'py>(&self, py: Python<'py>, input: Vec) -> PyResult<()> { let service = self.api_service.clone(); let wrapper = diff --git a/python_tests/requirements.txt b/python_tests/requirements.txt index c73f18a..5d91ece 100644 --- a/python_tests/requirements.txt +++ b/python_tests/requirements.txt @@ -14,3 +14,5 @@ pytest-asyncio numpy pandas polars +# pandas converts Arrow through pyarrow (BinaryDatapoints.to_pandas). +pyarrow diff --git a/python_tests/test_timeseries_binary.py b/python_tests/test_timeseries_binary.py index d07f5b4..38b3d29 100644 --- a/python_tests/test_timeseries_binary.py +++ b/python_tests/test_timeseries_binary.py @@ -31,15 +31,15 @@ def fetch(): return poll_until(fetch, lambda dps: len(dps) >= expected) -@pytest.mark.parametrize("zstd_level", [None, 1, 3, 9]) -def test_insert_from_lists_binary_round_trips_floats(sync_client, make_ts, zstd_level): +@pytest.mark.parametrize("zstd_level, compression", [(None, None), (1, None), (3, None), (9, None), (3, "buffers")]) +def test_insert_from_lists_binary_round_trips_floats(sync_client, make_ts, zstd_level, compression): ts = make_ts(value_type="float") # Out of order on purpose: the frame writer sorts, the read must still be in time order. values = [i * 1.25 - 20.0 for i in range(len(TIMESTAMPS))] shuffled = list(zip(TIMESTAMPS, values))[::-1] result = sync_client.timeseries.insert_from_lists_binary( - [t for t, _ in shuffled], [v for _, v in shuffled], ts, zstd_level=zstd_level + [t for t, _ in shuffled], [v for _, v in shuffled], ts, zstd_level=zstd_level, compression=compression ) assert result == [] diff --git a/python_tests/test_timeseries_binary_dataframe.py b/python_tests/test_timeseries_binary_dataframe.py new file mode 100644 index 0000000..4c69445 --- /dev/null +++ b/python_tests/test_timeseries_binary_dataframe.py @@ -0,0 +1,242 @@ +"""Binary ingest from a dataframe (`insert_dataframe_binary`): a polars, pandas or pyarrow table, +one column per series named by its external id, handed over through the Arrow C stream. + +Offline, a stand-in server answers the series lookup and keeps what is posted; the posted frames +are then served back through the binary read, so a round trip needs no backend. The live test +round-trips a wide table through a real one. +""" +import datetime +import http.server +import json +import threading + +import pytest + +import intellistream_datahub_sdk +from intellistream_datahub_sdk import DataHubException +from python_tests.fixtures import * # noqa: F401,F403 + +pl = pytest.importorskip("polars") +UTC = datetime.timezone.utc +T0 = datetime.datetime(2025, 1, 1, tzinfo=UTC) +SERIES = {"pump.temp": (11, "float"), "pump.count": (12, "bigint"), "pump.state": (13, "text")} + + +class _Api(http.server.BaseHTTPRequestHandler): + """Answers /timeseries/byids from SERIES, keeps posted frames, and serves them back on a read.""" + protocol_version = "HTTP/1.1" + posted = [] + + def _reply(self, status, content_type, body): + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_POST(self): + body = self.rfile.read(int(self.headers.get("Content-Length", 0))) + if self.path == "/timeseries/byids": + wanted = [i.get("externalId") for i in json.loads(body)["items"]] + items = [ + {"id": str(SERIES[e][0]), "externalId": e, "name": e, "valueType": SERIES[e][1]} + for e in wanted if e in SERIES + ] + self._reply(200, "application/json", json.dumps({"items": items}).encode()) + elif self.path == "/timeseries/data/binary": + type(self).posted.append(body) + self._reply(204, "application/json", b"") + elif self.path == "/timeseries/data/list/binary": + self._reply(200, "application/vnd.intellistream.datapoint-block", b"".join(type(self).posted) + b"DHDE") + else: + self._reply(404, "text/plain", b"") + + def log_message(self, *args): + pass + + +@pytest.fixture +def api(): + _Api.posted = [] + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _Api) + threading.Thread(target=server.serve_forever, daemon=True).start() + client = intellistream_datahub_sdk.DataHubClient( + base_url=f"http://127.0.0.1:{server.server_address[1]}", token="test-token" + ) + yield client, _Api + server.shutdown() + + +def _read_back(client): + """What the posted frames read back as: the wide table, by external id.""" + return client.timeseries.retrieve_datapoints_binary( + intellistream_datahub_sdk.RetrieveFilter(ts="pump.temp", start=T0) + ).to_polars(wide=True) + + +def test_a_polars_frame_round_trips(api): + client, server = api + df = pl.DataFrame({ + "timestamp": [T0 + datetime.timedelta(seconds=s) for s in (2, 0, 1)], # unsorted on purpose + "pump.temp": [2.5, 0.5, None], # a gap + "pump.count": [3, 1, 2], + "pump.state": ["off", "on", "on"], + }) + + assert client.timeseries.insert_dataframe_binary(df) == 8 + assert len(server.posted) == 1 + + # The stand-in replays the posted frames, which are grouped by value type; a real read follows + # the request's items. So compare by name. + back = _read_back(client) + assert sorted(back.columns) == ["pump.count", "pump.state", "pump.temp", "timestamp"] + assert back["timestamp"].to_list() == [T0 + datetime.timedelta(seconds=s) for s in range(3)] + assert back["pump.temp"].to_list() == [0.5, None, 2.5] + assert back["pump.count"].to_list() == [1, 2, 3] + assert back["pump.state"].to_list() == ["on", "on", "off"] + + +def test_buffers_compression_sends_compression_2_frames_that_read_back_the_same(api): + client, server = api + df = pl.DataFrame({ + "timestamp": [T0 + datetime.timedelta(seconds=s) for s in range(3)], + "pump.temp": [0.5, None, 2.5], + "pump.state": ["on", "on", "off"], + }) + assert client.timeseries.insert_dataframe_binary(df, compression="buffers") == 5 + assert [server.posted[0][7]] == [2] + back = _read_back(client) + assert back["pump.temp"].to_list() == [0.5, None, 2.5] + assert back["pump.state"].to_list() == ["on", "on", "off"] + + with pytest.raises(ValueError, match="compression"): + client.timeseries.insert_dataframe_binary(df, compression="lz4") + + +def test_a_pandas_frame_uses_its_timestamp_index(api): + pd = pytest.importorskip("pandas") + pytest.importorskip("pyarrow") + client, _ = api + index = pd.date_range("2025-01-01", periods=3, freq="s", tz="UTC", name="time") + df = pd.DataFrame({"pump.temp": [1.0, float("nan"), 3.0]}, index=index) + + # pandas exports NaN as null, so the middle row is no point at all. + assert client.timeseries.insert_dataframe_binary(df) == 2 + assert _read_back(client)["pump.temp"].to_list() == [1.0, 3.0] + + +def test_a_pyarrow_table_with_a_named_timestamp_column(api): + pa = pytest.importorskip("pyarrow") + client, _ = api + table = pa.table({ + "at": pa.array([T0, T0 + datetime.timedelta(milliseconds=5)], pa.timestamp("ns", "UTC")), + "pump.temp": [1.0, 2.0], + }) + assert client.timeseries.insert_dataframe_binary(table, timestamp="at") == 2 + assert _read_back(client)["timestamp"].to_list() == [T0, T0 + datetime.timedelta(milliseconds=5)] + + +def test_an_unknown_column_is_not_found_before_anything_is_sent(api): + client, server = api + df = pl.DataFrame({"timestamp": [T0], "pump.temp": [1.0], "no.such.series": [2.0]}) + with pytest.raises(DataHubException, match="no.such.series") as excinfo: + client.timeseries.insert_dataframe_binary(df) + assert excinfo.value.status_code == 404 + assert server.posted == [] + + +@pytest.mark.parametrize("timestamps, message", [ + ([datetime.datetime(2025, 1, 1)], "no time zone"), + ([T0 + datetime.timedelta(microseconds=1500)], "finer than a millisecond"), +]) +def test_timestamps_that_are_not_an_instant_in_milliseconds_are_refused(api, timestamps, message): + client, server = api + df = pl.DataFrame({"timestamp": timestamps, "pump.temp": [1.0]}) + with pytest.raises(DataHubException, match=message) as excinfo: + client.timeseries.insert_dataframe_binary(df) + assert excinfo.value.status_code == 422 + assert server.posted == [] + + +def test_a_column_its_series_cannot_hold_is_refused(api): + client, server = api + df = pl.DataFrame({"timestamp": [T0], "pump.count": [1.5]}) + with pytest.raises(DataHubException, match="cannot hold"): + client.timeseries.insert_dataframe_binary(df) + assert server.posted == [] + + +def test_the_timestamp_column_has_to_be_found(api): + client, _ = api + with pytest.raises(ValueError, match="no timestamp column"): + client.timeseries.insert_dataframe_binary(pl.DataFrame({"pump.temp": [1.0]})) + two = pl.DataFrame({"a": [T0], "b": [T0], "pump.temp": [1.0]}) + with pytest.raises(ValueError, match="several timestamp columns"): + client.timeseries.insert_dataframe_binary(two) + with pytest.raises(TypeError, match="__arrow_c_stream__"): + client.timeseries.insert_dataframe_binary([1, 2, 3]) + + +# --- live --------------------------------------------------------------------------------------- + +START = datetime.datetime(2024, 3, 1, tzinfo=UTC) + + +@pytest.mark.parametrize("compression", ["payload", "buffers"]) +def test_a_wide_table_round_trips_through_the_api(sync_client, make_ts, compression): + temp, count = make_ts(value_type="float"), make_ts(value_type="bigint") + n = 150_000 # past one frame's 100 000 rows + stamps = pl.datetime_range(START, START + datetime.timedelta(seconds=n - 1), "1s", time_zone="UTC", eager=True) + df = pl.DataFrame({ + "timestamp": stamps, + temp.external_id: pl.Series([i * 0.5 for i in range(n)]), + # Every third point missing: a ragged column, its nulls left out. + count.external_id: pl.Series([None if i % 3 == 0 else i for i in range(n)], dtype=pl.Int64), + }) + + sent = sync_client.timeseries.insert_dataframe_binary(df, compression=compression) + assert sent == n + n - (n + 2) // 3 + + end = START + datetime.timedelta(seconds=n) + filters = [intellistream_datahub_sdk.RetrieveFilter(ts=s, start=START, end=end) for s in (temp, count)] + + from polling import poll_until + back = poll_until( + lambda: sync_client.timeseries.retrieve_datapoints_binary(filters), + lambda read: sum(len(s) for s in read) == sent, + ).to_polars(wide=True) + assert back.equals(df) + + # What came back goes back in as it is. + assert sync_client.timeseries.insert_dataframe_binary(back) == sent + + +def test_a_pandas_frame_round_trips_through_the_api(sync_client, make_ts): + pd = pytest.importorskip("pandas") + pytest.importorskip("pyarrow") + temp, count = make_ts(value_type="float"), make_ts(value_type="bigint") + n = 120_000 # past one frame's 100 000 rows + index = pd.date_range(START, periods=n, freq="s", tz="UTC", name="time") + df = pd.DataFrame({ + # NaN every fifth row: pandas exports it as null, so no point there. + temp.external_id: [float("nan") if i % 5 == 0 else i * 0.25 for i in range(n)], + count.external_id: pd.array([None if i % 3 == 0 else i for i in range(n)], dtype="Int64"), + }, index=index) + + sent = sync_client.timeseries.insert_dataframe_binary(df) + temps, counts = df[temp.external_id].dropna(), df[count.external_id].dropna() + assert sent == len(temps) + len(counts) + + end = START + datetime.timedelta(seconds=n) + filters = [intellistream_datahub_sdk.RetrieveFilter(ts=s, start=START, end=end) for s in (temp, count)] + + from polling import poll_until + read = poll_until( + lambda: sync_client.timeseries.retrieve_datapoints_binary(filters), + lambda read: sum(len(s) for s in read) == sent, + ) + by_id = {s.external_id: s.to_pandas() for s in read} + for column in (temps, counts): + back = by_id[column.name] + assert back["timestamp"].tolist() == column.index.tolist() + assert back["value"].tolist() == column.tolist() diff --git a/python_tests/test_timeseries_binary_read.py b/python_tests/test_timeseries_binary_read.py new file mode 100644 index 0000000..60f9de9 --- /dev/null +++ b/python_tests/test_timeseries_binary_read.py @@ -0,0 +1,468 @@ +"""Binary datapoint read (`POST /timeseries/data/list/binary`) through the Python bindings. + +Two halves. The offline one answers the client from a stand-in server with response bodies the +platform's own Java frame writer produced (`src/timeseries/fixtures/`), so it checks the binding +and the Arrow PyCapsule stream without a backend. The live one round-trips through a backend and +skips when that backend does not serve the endpoint yet. +""" +import datetime +import decimal +import http.server +import pathlib +import struct +import threading + +import pytest + +import intellistream_datahub_sdk +from intellistream_datahub_sdk import DataHubException +from polling import poll_until +from python_tests.fixtures import * # noqa: F401,F403 + +FIXTURES = pathlib.Path(__file__).resolve().parent.parent / "src" / "timeseries" / "fixtures" +T0 = datetime.datetime.fromtimestamp(1_700_000_000, tz=datetime.timezone.utc) +MEDIA_TYPE = "application/vnd.intellistream.datapoint-block" + + +class _StandIn(http.server.BaseHTTPRequestHandler): + status = 200 + body = b"" + content_type = MEDIA_TYPE + seen = [] + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + type(self).seen.append((self.path, self.headers.get("Accept"), self.rfile.read(length))) + self.send_response(type(self).status) + self.send_header("Content-Type", type(self).content_type) + self.send_header("Content-Length", str(len(type(self).body))) + self.end_headers() + self.wfile.write(type(self).body) + + def log_message(self, *args): + pass + + +@pytest.fixture +def stand_in(): + _StandIn.status, _StandIn.content_type, _StandIn.seen = 200, MEDIA_TYPE, [] + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _StandIn) + threading.Thread(target=server.serve_forever, daemon=True).start() + client = intellistream_datahub_sdk.DataHubClient( + base_url=f"http://127.0.0.1:{server.server_address[1]}", token="test-token" + ) + yield client, _StandIn + server.shutdown() + + +def _filter(ts="rd_float"): + return intellistream_datahub_sdk.RetrieveFilter(ts=ts, start=T0, limit=1000) + + +# read_values.bin is compression 2, record batches compressing their own buffers, as a read answers +# now; read_values_zstd.bin is the same series as compression 1, one zstd stream per frame. +@pytest.mark.parametrize("fixture", ["read_values.bin", "read_values_zstd.bin"]) +def test_every_value_type_converts_through_the_arrow_stream(stand_in, fixture): + pl = pytest.importorskip("polars") + client, server = stand_in + server.body = (FIXTURES / fixture).read_bytes() + + series = client.timeseries.retrieve_datapoints_binary([_filter("rd_bigint"), _filter("rd_float")]) + + path, accept, _ = server.seen[0] + assert path == "/timeseries/data/list/binary" + assert accept.startswith(MEDIA_TYPE) and "application/problem+json" in accept + assert [s.external_id for s in series] == [ + "rd_bigint", "rd_float", "rd_float32", "rd_numeric", "rd_decimal32", "rd_text", "rd_mixed", "rd_mixed_numbers", + ] + by_name = {s.external_id: s for s in series} + + bigint = by_name["rd_bigint"] + assert (bigint.id, bigint.value_type, len(bigint), bigint.is_aggregates) == (11, "bigint", 3, False) + assert bigint.columns == ["timestamp", "value"] + frame = pl.DataFrame(bigint) + assert frame["value"].to_list() == [1, 42, -9_000_000_000] + assert frame["timestamp"].dtype == pl.Datetime("ms", "UTC") + assert frame["timestamp"][0] == T0 + # A stream is made per call, so one object converts more than once. + assert bigint.to_polars().equals(frame) + + assert pl.DataFrame(by_name["rd_float"])["value"].to_list() == [1.5, -2.25] + assert pl.DataFrame(by_name["rd_float32"])["value"].dtype == pl.Float32 + numeric = pl.DataFrame(by_name["rd_numeric"])["value"] + assert numeric.dtype == pl.Decimal(18, 6) + assert numeric.to_list() == [decimal.Decimal("123456789012.345678"), decimal.Decimal("-0.000001")] + assert pl.DataFrame(by_name["rd_decimal32"])["value"].dtype == pl.Decimal(9, 4) + assert pl.DataFrame(by_name["rd_text"])["value"].to_list() == ["open", "blåbær"] + + mixed = by_name["rd_mixed"] + assert mixed.columns == ["timestamp", "value_numeric", "value_text"] + mixed_frame = pl.DataFrame(mixed) + assert mixed_frame["value_numeric"].to_list() == [3.5, None, -7.0] + assert mixed_frame["value_text"].to_list() == [None, "fault", None] + + +def test_aggregates_convert_to_pandas(stand_in): + pytest.importorskip("pyarrow") + pytest.importorskip("pandas") + client, server = stand_in + server.body = (FIXTURES / "read_aggregates.bin").read_bytes() + + (series,) = client.timeseries.retrieve_datapoints_binary(_filter("rd_aggregates")) + + assert series.is_aggregates and series.value_type == "float" + assert series.columns == ["timestamp", "avg", "min", "max", "sum"] + frame = series.to_pandas() + assert frame["avg"].tolist() == [1.5, 2.5, 3.5] + assert frame["sum"].tolist() == [6.0, 10.0, 14.0] + assert frame["timestamp"].iloc[1] == T0 + datetime.timedelta(hours=1) + + +def test_no_points_is_no_series(stand_in): + client, server = stand_in + server.body = b"DHDE" + assert len(client.timeseries.retrieve_datapoints_binary(_filter())) == 0 + + +def test_a_response_cut_short_raises_rather_than_returning_fewer_points(stand_in): + client, server = stand_in + full = (FIXTURES / "read_values.bin").read_bytes() + server.body = full[:-4] # every frame intact, only the end marker missing + with pytest.raises(DataHubException, match="cut short") as excinfo: + client.timeseries.retrieve_datapoints_binary(_filter()) + assert excinfo.value.status_code == 502 + + +def test_a_refusal_is_the_problem_the_server_sent(stand_in): + client, server = stand_in + server.status, server.content_type = 400, "application/problem+json" + server.body = ( + b'{"type":"https://intellistream.ai/errors/bad-request","title":"Bad Request","status":400,' + b'"detail":"The binary read has no cursors; page by moving `start` past the last timestamp read."}' + ) + with pytest.raises(DataHubException, match="no cursors") as excinfo: + client.timeseries.retrieve_datapoints_binary(_filter()) + assert excinfo.value.status_code == 400 + + +# --- streaming, offline ----------------------------------------------------------------------- + + +def _frames(body): + """The frames of a response body, each as its own bytes, without the end marker.""" + frames, pos = [], 0 + while body[pos:pos + 4] == b"DHDP": + directory, payload = struct.unpack_from("= expected, + ) + + +def test_binary_read_matches_the_json_read(sync_client, make_ts, binary_read): + pl = pytest.importorskip("polars") + ts = make_ts(value_type="float") + values = [i * 0.5 - 30.0 for i in range(len(TIMESTAMPS))] + sync_client.timeseries.insert_from_lists_binary(TIMESTAMPS, values, ts) + _wait_for(sync_client, ts, len(TIMESTAMPS)) + + (series,) = binary_read(intellistream_datahub_sdk.RetrieveFilter(ts=ts, start=START, limit=10_000)) + + assert series.external_id == ts.external_id and series.value_type == "float" + frame = series.to_polars() + assert frame["timestamp"].to_list() == TIMESTAMPS + assert frame["value"].to_list() == values + + +def test_binary_read_of_hourly_aggregates(sync_client, make_ts, binary_read): + pytest.importorskip("pyarrow") + ts = make_ts(value_type="float") + sync_client.timeseries.insert_from_lists_binary(TIMESTAMPS, [float(i) for i in range(len(TIMESTAMPS))], ts) + _wait_for(sync_client, ts, len(TIMESTAMPS)) + + rf = intellistream_datahub_sdk.RetrieveFilter( + ts=ts, start=START, end=TIMESTAMPS[-1] + datetime.timedelta(minutes=1), + aggregates=["avg", "min", "max", "sum"], granularity="1h", + ) + (series,) = binary_read(rf) + json_buckets = sync_client.timeseries.retrieve_datapoints(rf)[0].get_datapoints() + + frame = series.to_pandas() + assert series.is_aggregates and len(frame) == len(json_buckets) == 2 + # Zero-order hold with carry-in, so the second hour's min is the 59 carried into it. + assert frame["min"].tolist() == [dp.min for dp in json_buckets] + assert frame["max"].tolist() == [dp.max for dp in json_buckets] + assert frame["sum"].tolist() == pytest.approx([dp.sum for dp in json_buckets]) + assert frame["avg"].tolist() == pytest.approx([dp.average for dp in json_buckets]) + + +def test_binary_read_of_an_unknown_series_is_a_bad_request(make_ts, binary_read): + missing = unique_id("ts_missing") + with pytest.raises(DataHubException, match=missing) as excinfo: + binary_read(intellistream_datahub_sdk.RetrieveFilter(ts=missing, start=START)) + assert excinfo.value.status_code == 400 + + +def test_streaming_read_matches_the_buffered_read(sync_client, make_ts, binary_read): + pl = pytest.importorskip("polars") + ts = make_ts(value_type="float") + values = [i * 0.5 - 30.0 for i in range(len(TIMESTAMPS))] + sync_client.timeseries.insert_from_lists_binary(TIMESTAMPS, values, ts) + _wait_for(sync_client, ts, len(TIMESTAMPS)) + rf = intellistream_datahub_sdk.RetrieveFilter(ts=ts, start=START, limit=10_000) + + (buffered,) = binary_read(rf) + chunks = list(sync_client.timeseries.retrieve_datapoints_binary_stream(rf)) + reader = sync_client.timeseries.retrieve_datapoints_binary_stream(rf) + long = reader.to_polars() + + assert pl.concat([c.to_polars() for c in chunks]).equals(buffered.to_polars()) + assert long.select("timestamp", "value").equals(buffered.to_polars()) + assert reader.external_ids == {buffered.id: ts.external_id} + + +def test_wide_from_series_on_the_same_and_on_different_timestamps(sync_client, make_ts, binary_read): + pl = pytest.importorskip("polars") + a, b, c = (make_ts(value_type="float") for _ in range(3)) + sync_client.timeseries.insert_from_lists_binary(TIMESTAMPS, [float(i) for i in range(len(TIMESTAMPS))], a) + sync_client.timeseries.insert_from_lists_binary(TIMESTAMPS, [-float(i) for i in range(len(TIMESTAMPS))], b) + every_other = TIMESTAMPS[1::2] + sync_client.timeseries.insert_from_lists_binary(every_other, [0.5] * len(every_other), c) + for ts, n in ((a, len(TIMESTAMPS)), (b, len(TIMESTAMPS)), (c, len(every_other))): + _wait_for(sync_client, ts, n) + rf = lambda *series: [intellistream_datahub_sdk.RetrieveFilter(ts=s, start=START, limit=10_000) for s in series] + + # Same timestamps: the columns side by side. + same = sync_client.timeseries.retrieve_datapoints_binary(rf(a, b)).to_polars(wide=True) + assert same.columns == ["timestamp", a.external_id, b.external_id] + assert same["timestamp"].to_list() == TIMESTAMPS + assert same[b.external_id].to_list() == [-float(i) for i in range(len(TIMESTAMPS))] + + # Different timestamps: matched up, the gaps null; streamed gives the same table. + mixed = sync_client.timeseries.retrieve_datapoints_binary(rf(a, c)).to_polars(wide=True) + assert mixed["timestamp"].to_list() == TIMESTAMPS + assert mixed[c.external_id].to_list() == [None, 0.5] * (len(TIMESTAMPS) // 2) + assert mixed.equals(sync_client.timeseries.retrieve_datapoints_binary_stream(rf(a, c)).to_polars(wide=True)) diff --git a/src/blocking.rs b/src/blocking.rs index 1dd827a..4719394 100644 --- a/src/blocking.rs +++ b/src/blocking.rs @@ -49,7 +49,7 @@ use crate::resources::{ ResourceNetwork, ResourceUpdate, }; use crate::timeseries::{ - BinaryIngestOptions, TimeSeries, TimeSeriesFilter, TimeSeriesUpdateCollection, + BinaryDatapointStream, BinaryDatapoints, BinaryIngestOptions, TimeSeries, TimeSeriesFilter, TimeSeriesUpdateCollection, ValueTypeRecommendation, }; use crate::unit::Unit; @@ -182,6 +182,7 @@ impl TimeSeriesService { fn insert_datapoints(json: &mut DataWrapper>) -> Result, ResponseError>; fn insert_datapoints_binary(json: &DataWrapper>, options: &BinaryIngestOptions) -> Result, ResponseError>; fn retrieve_datapoints(json: &DataWrapper) -> Result>, ResponseError>; + fn retrieve_datapoints_binary(json: &DataWrapper) -> Result, ResponseError>; fn delete_datapoints(json: &DataWrapper) -> Result, ResponseError>; fn retrieve_latest_datapoint(json: &DataWrapper) -> Result>, ResponseError>; } @@ -195,6 +196,37 @@ impl TimeSeriesService { pub fn evict_binary_series_cache(&self) { self.api.time_series.evict_binary_series_cache() } + + /// [`stream_datapoints_binary`](crate::timeseries::TimeSeriesService::stream_datapoints_binary) + /// as an iterator: each `next` waits for the next frame's rows, decoded while the rest download. + pub fn stream_datapoints_binary( + &self, + json: &DataWrapper, + ) -> Result { + let stream = self + .rt + .block_on(self.api.time_series.stream_datapoints_binary(json))?; + Ok(BinaryDatapointIter { + stream, + rt: self.rt.clone(), + }) + } +} + +/// A binary read frame by frame; see +/// [`BinaryDatapointStream`](crate::timeseries::BinaryDatapointStream) for what an item is and how +/// a response that breaks off shows up. +pub struct BinaryDatapointIter { + stream: BinaryDatapointStream, + rt: Arc, +} + +impl Iterator for BinaryDatapointIter { + type Item = Result; + + fn next(&mut self) -> Option { + self.rt.block_on(futures::StreamExt::next(&mut self.stream)) + } } /// Blocking counterpart of [`crate::ResourceService`]. diff --git a/src/generic.rs b/src/generic.rs index 7faad32..697eb6f 100644 --- a/src/generic.rs +++ b/src/generic.rs @@ -882,6 +882,51 @@ pub trait ApiServiceProvider { } } + /// `POST` a JSON body to an endpoint that answers with bytes of type `accept` rather than JSON + /// (currently only `/timeseries/data/list/binary`), returning the response to read as it + /// arrives. + /// + /// `application/problem+json` is offered beside `accept` so a refusal reaches the caller as the + /// problem document it is, rather than as Spring's 406 for a type the handler does not produce. + /// The status is checked here, before any of the body is read, so a refusal is always an + /// `Err` and never a stream. + async fn execute_post_request_for_stream( + &self, + path: &str, + json: &J, + accept: &str, + ) -> Result { + let token = self.get_token().await?; + let response = self + .get_api_service() + .http_client + .post(path) + .json(json) + .header(http::header::ACCEPT, format!("{accept}, application/problem+json")) + .bearer_auth(token.clone()) + .send() + .await + .map_err(ResponseError::from_err)?; + let status = response.status(); + if status.is_success() { + return Ok(response); + } + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + let error = ResponseError { + status, + message: response + .text() + .await + .unwrap_or_else(|_| "Failed to read response body".to_string()), + content_type, + }; + Err(self.on_request_error(error, &token).await) + } + /// `GET` an endpoint that answers with bytes rather than JSON (currently only /// `/files/download/{id}`). /// diff --git a/src/timeseries/binary.rs b/src/timeseries/binary.rs index 45cbfac..aab6f5b 100644 --- a/src/timeseries/binary.rs +++ b/src/timeseries/binary.rs @@ -16,7 +16,8 @@ use arrow_array::{ ArrayRef, Decimal128Array, Float32Array, Float64Array, Int64Array, RecordBatch, StringArray, TimestampMillisecondArray, }; -use arrow_ipc::writer::StreamWriter; +use arrow_ipc::writer::{IpcWriteOptions, StreamWriter}; +use arrow_ipc::CompressionType; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use futures::future::join_all; use oauth2::http::StatusCode; @@ -32,7 +33,11 @@ pub const MEDIA_TYPE: &str = "application/vnd.intellistream.datapoint-block"; pub const MAGIC: &[u8; 4] = b"DHDP"; pub const VERSION: u8 = 1; pub const CODEC_ARROW_IPC: u8 = 1; +/// The payload is one zstd stream around an Arrow IPC stream with uncompressed buffers. pub const COMPRESSION_ZSTD: u8 = 1; +/// The payload is an Arrow IPC stream whose record batches compress each buffer with zstd +/// (Arrow's body compression): what a binary read answers with, ClickHouse's batches relayed. +pub const COMPRESSION_ARROW_ZSTD: u8 = 2; /// Bytes before the series directory. pub const HEADER_BYTES: usize = 28; @@ -50,8 +55,8 @@ pub const NUMERIC_ROWS_PER_FRAME: usize = 100_000; /// Rows per frame for `text` and `mixed`. pub const TEXT_ROWS_PER_FRAME: usize = 10_000; -const NUMERIC_UNSCALED_MAX: i128 = 999_999_999_999_999_999; -const DECIMAL32_UNSCALED_MAX: i128 = 999_999_999; +pub(crate) const NUMERIC_UNSCALED_MAX: i128 = 999_999_999_999_999_999; +pub(crate) const DECIMAL32_UNSCALED_MAX: i128 = 999_999_999; /// Headroom under the raw cap: arrow-rs pads every buffer to 64 bytes, which the estimate ignores. const RAW_BYTES_MARGIN: usize = 64 * 1024; @@ -86,6 +91,20 @@ impl DatapointValueType { self as u8 } + /// The value type a frame envelope numbers `id`. + pub fn from_id(id: u8) -> Option { + Some(match id { + 1 => Self::Bigint, + 2 => Self::Float, + 3 => Self::Numeric, + 4 => Self::Text, + 5 => Self::Decimal32, + 6 => Self::Mixed, + 7 => Self::Float32, + _ => return None, + }) + } + pub fn carries_text(self) -> bool { matches!(self, Self::Text | Self::Mixed) } @@ -135,11 +154,36 @@ impl DatapointValueType { } } +/// Where a frame's zstd goes. Either way the frame is zstd at the chosen level and the server +/// stores the same points. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum FrameCompression { + /// Compression 1: one zstd stream around the whole Arrow IPC payload. + Payload, + /// Compression 2: every record batch compresses each buffer in-band (Arrow's body + /// compression), the form a binary read answers with. The consumer relays these batches to + /// ClickHouse without decompressing them, where compression 1 has to be decoded first. + #[default] + Buffers, +} + +impl FrameCompression { + /// The envelope's compression byte. + pub fn id(self) -> u8 { + match self { + FrameCompression::Payload => COMPRESSION_ZSTD, + FrameCompression::Buffers => COMPRESSION_ARROW_ZSTD, + } + } +} + /// How [`TimeSeriesService::insert_datapoints_binary`] compresses and retries. #[derive(Debug, Clone)] pub struct BinaryIngestOptions { /// zstd level per frame: 1, 3 or 9. Compression is mandatory, so there is no way to turn it off. pub zstd_level: i32, + /// Whether zstd wraps the whole payload or each Arrow buffer. + pub compression: FrameCompression, /// Retries of one request after a 429 or a 5xx, one second apart per attempt. pub max_retries: u32, /// Requests in flight at once when one call packs into more than one, which happens above @@ -150,7 +194,8 @@ pub struct BinaryIngestOptions { impl Default for BinaryIngestOptions { fn default() -> Self { BinaryIngestOptions { - zstd_level: 9, + zstd_level: 3, + compression: FrameCompression::Buffers, max_retries: 3, request_concurrency: 4, } @@ -167,6 +212,11 @@ impl BinaryIngestOptions { self } + pub fn compression(mut self, compression: FrameCompression) -> Self { + self.compression = compression; + self + } + pub fn max_retries(mut self, retries: u32) -> Self { self.max_retries = retries; self @@ -177,7 +227,7 @@ impl BinaryIngestOptions { self } - fn validate(&self) -> Result<(), ResponseError> { + pub(crate) fn validate(&self) -> Result<(), ResponseError> { if ![1, 3, 9].contains(&self.zstd_level) { return Err(ResponseError::bad_request(format!( "zstd level {} is not one of 1, 3 or 9", @@ -403,6 +453,11 @@ impl FrameWriter { /// Sorts by (id, timestamp), keeps the last value of a repeated pair, checks the caps, writes /// the Arrow stream, compresses it and returns the complete frame. pub fn build(&self, zstd_level: i32) -> Result { + self.build_with(FrameCompression::Payload, zstd_level) + } + + /// [`FrameWriter::build`] with the zstd where `compression` puts it. + pub fn build_with(&self, compression: FrameCompression, zstd_level: i32) -> Result { if self.rows.is_empty() { return Err("frame has no rows".to_string()); } @@ -441,45 +496,21 @@ impl FrameWriter { )); } - let raw = self.write_ipc(&kept)?; - if raw.len() > MAX_FRAME_RAW_BYTES { - return Err(format!( - "frame payload is {} bytes, the cap is {MAX_FRAME_RAW_BYTES}; split it", - raw.len() - )); - } - let compressed = - zstd::bulk::compress(&raw, zstd_level).map_err(|e| format!("zstd failed: {e}"))?; - - let directory_len: usize = directory - .iter() - .map(|(_, external_id)| 8 + varint_size(external_id.len()) + external_id.len()) - .sum(); - let mut frame = Vec::with_capacity(HEADER_BYTES + directory_len + compressed.len()); - frame.extend_from_slice(MAGIC); - frame.push(VERSION); - frame.push(self.value_type.id()); - frame.push(CODEC_ARROW_IPC); - frame.push(COMPRESSION_ZSTD); - for n in [rows, directory.len(), directory_len, compressed.len(), raw.len()] { - frame.extend_from_slice(&(n as u32).to_le_bytes()); - } - for (id, external_id) in &directory { - frame.extend_from_slice(&id.to_le_bytes()); - write_varint(&mut frame, external_id.len()); - frame.extend_from_slice(external_id.as_bytes()); - } - frame.extend_from_slice(&compressed); - Ok(Frame { - bytes: frame, + let batch = self.record_batch(&kept)?; + let mut frame = assemble_frame( + self.value_type, rows, - series: directory.len(), - raw_len: raw.len(), - clamped: self.clamped, - }) + &directory, + &batch, + self.estimated_raw_bytes(), + compression, + zstd_level, + )?; + frame.clamped = self.clamped; + Ok(frame) } - fn write_ipc(&self, kept: &[usize]) -> Result, String> { + fn record_batch(&self, kept: &[usize]) -> Result { let ids: Vec = kept.iter().map(|&i| self.rows[i].id).collect(); let timestamps: Vec = kept.iter().map(|&i| self.rows[i].timestamp).collect(); let mut columns: Vec = vec![ @@ -561,20 +592,120 @@ impl FrameWriter { columns.push(Arc::new(StringArray::from(text))); } } - let schema = self.value_type.schema(); - let batch = RecordBatch::try_new(schema.clone(), columns).map_err(|e| e.to_string())?; - let mut raw = Vec::with_capacity(self.estimated_raw_bytes()); - let mut writer = - StreamWriter::try_new(&mut raw, schema.as_ref()).map_err(|e| e.to_string())?; - writer.write(&batch).map_err(|e| e.to_string())?; - writer.finish().map_err(|e| e.to_string())?; - drop(writer); - Ok(raw) + RecordBatch::try_new(self.value_type.schema(), columns).map_err(|e| e.to_string()) + } +} + +/// A batch in the canon as an Arrow IPC stream, its buffers compressed in-band when `in_band` +/// names a zstd level. +fn ipc_stream(batch: &RecordBatch, capacity: usize, in_band: Option) -> Result, String> { + let options = match in_band { + None => IpcWriteOptions::default(), + Some(level) => IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::ZSTD)) + .and_then(|o| o.try_with_compression_level(Some(level))) + .map_err(|e| e.to_string())?, + }; + let mut raw = Vec::with_capacity(capacity); + let mut writer = StreamWriter::try_new_with_options(&mut raw, batch.schema().as_ref(), options) + .map_err(|e| e.to_string())?; + writer.write(batch).map_err(|e| e.to_string())?; + writer.finish().map_err(|e| e.to_string())?; + drop(writer); + Ok(raw) +} + +/// A compression 2 frame's raw length as the server counts it: every record batch buffer +/// uncompressed, each padded to 8 bytes. Read back from the stream, so it holds whichever buffers +/// arrow-rs chose to write. +fn in_band_raw_len(stream: &[u8]) -> Result { + let fail = |at: usize| format!("unreadable IPC message at {at}"); + let mut total = 0; + let mut pos = 0; + loop { + let prefix = stream.get(pos..pos + 8).ok_or_else(|| fail(pos))?; + let meta_len = u32::from_le_bytes(prefix[4..8].try_into().unwrap()) as usize; + if meta_len == 0 { + return Ok(total); + } + let meta = stream.get(pos + 8..pos + 8 + meta_len).ok_or_else(|| fail(pos))?; + let message = arrow_ipc::root_as_message(meta).map_err(|_| fail(pos))?; + let body = pos + 8 + meta_len; + if let Some(batch) = message.header_as_record_batch() { + for buffer in batch.buffers().into_iter().flatten() { + let (offset, length) = (body + buffer.offset() as usize, buffer.length() as usize); + if length == 0 { + continue; + } + let declared = stream.get(offset..offset + 8).ok_or_else(|| fail(pos))?; + let declared = i64::from_le_bytes(declared.try_into().unwrap()); + let raw = if declared == -1 { length - 8 } else { declared as usize }; + total += raw.div_ceil(8) * 8; + } + } + pos = body + message.bodyLength() as usize; } } +/// The envelope around a batch: header, series directory (ascending ids, as the rows are +/// ordered) and the batch as an Arrow IPC stream compressed as `compression` says. +pub(crate) fn assemble_frame( + value_type: DatapointValueType, + rows: usize, + directory: &[(i64, &str)], + batch: &RecordBatch, + capacity: usize, + compression: FrameCompression, + zstd_level: i32, +) -> Result { + let (compressed, raw_len) = match compression { + FrameCompression::Payload => { + let raw = ipc_stream(batch, capacity, None)?; + let compressed = + zstd::bulk::compress(&raw, zstd_level).map_err(|e| format!("zstd failed: {e}"))?; + (compressed, raw.len()) + } + FrameCompression::Buffers => { + let stream = ipc_stream(batch, capacity / 2, Some(zstd_level))?; + let raw_len = in_band_raw_len(&stream)?; + (stream, raw_len) + } + }; + if raw_len > MAX_FRAME_RAW_BYTES { + return Err(format!( + "frame payload is {raw_len} bytes, the cap is {MAX_FRAME_RAW_BYTES}; split it" + )); + } + let directory_len: usize = directory + .iter() + .map(|(_, external_id)| 8 + varint_size(external_id.len()) + external_id.len()) + .sum(); + let mut frame = Vec::with_capacity(HEADER_BYTES + directory_len + compressed.len()); + frame.extend_from_slice(MAGIC); + frame.push(VERSION); + frame.push(value_type.id()); + frame.push(CODEC_ARROW_IPC); + frame.push(compression.id()); + for n in [rows, directory.len(), directory_len, compressed.len(), raw_len] { + frame.extend_from_slice(&(n as u32).to_le_bytes()); + } + for (id, external_id) in directory { + frame.extend_from_slice(&id.to_le_bytes()); + write_varint(&mut frame, external_id.len()); + frame.extend_from_slice(external_id.as_bytes()); + } + frame.extend_from_slice(&compressed); + Ok(Frame { + bytes: frame, + rows, + series: directory.len(), + raw_len, + clamped: 0, + }) +} + /// The server counts the character limit in UTF-16 units, so this does too. -fn checked_text(value: &str) -> Result { +pub(crate) fn checked_text(value: &str) -> Result { if value.is_empty() { return Err("text value is empty".to_string()); } @@ -620,7 +751,7 @@ fn decimal_parts(value: &str) -> Option<(bool, String, i32)> { } /// Parses a decimal literal to its unscaled value at `scale` decimals, rounded half-up. -fn parse_scaled(value: &str, scale: u32) -> Option { +pub(crate) fn parse_scaled(value: &str, scale: u32) -> Option { let (negative, digits, point) = decimal_parts(value)?; let total = i32::try_from(digits.len()).ok()?; // Digits to keep: everything before the point plus `scale` after it. @@ -733,7 +864,7 @@ pub(crate) fn pack_requests(frames: Vec) -> Vec> { requests } -fn unprocessable(message: String) -> ResponseError { +pub(crate) fn unprocessable(message: String) -> ResponseError { ResponseError { status: StatusCode::UNPROCESSABLE_ENTITY, message, @@ -747,7 +878,7 @@ fn unprocessable(message: String) -> ResponseError { /// Read from the problem document, not the body text: the SDK's own pre-flight 404 names the /// series it could not resolve, so an external id spelling one of these slugs used to earn a /// pointless second attempt. -fn is_stale_series_rejection(error: &ResponseError) -> bool { +pub(crate) fn is_stale_series_rejection(error: &ResponseError) -> bool { let Some(problem) = error.problem() else { return false; }; @@ -800,7 +931,11 @@ impl TimeSeriesService { options: &BinaryIngestOptions, ) -> Result, ResponseError> { let collections = json.get_items(); - let resolved = self.resolve_series(collections).await?; + let wanted: Vec<(Option, Option<&str>)> = collections + .iter() + .map(|c| (c.id, c.external_id.as_deref())) + .collect(); + let resolved = self.resolve_series(&wanted).await?; let items: Vec<(ResolvedSeries, &[DatapointString])> = resolved .into_iter() .zip(collections.iter()) @@ -808,10 +943,10 @@ impl TimeSeriesService { .collect(); let writers = cut_into_writers(&items)?; - let level = options.zstd_level; + let (compression, level) = (options.compression, options.zstd_level); let builds = writers .into_iter() - .map(|writer| tokio::task::spawn_blocking(move || writer.build(level))); + .map(|writer| tokio::task::spawn_blocking(move || writer.build_with(compression, level))); let mut frames = Vec::new(); for built in join_all(builds).await { let frame = built @@ -819,12 +954,23 @@ impl TimeSeriesService { .map_err(unprocessable)?; frames.push(frame); } + self.post_frame_requests(frames, options).await?; + let mut result = DataWrapper::new(); + result.set_http_status_code(204); + Ok(result) + } - // Concurrently, not one after another. A call of up to 3.2M points packs into a single - // request and this changes nothing, but a larger one becomes several and the api - // validates them independently, so there is no reason to serialise them. Bounded by - // `request_concurrency` so a very large call cannot open an unbounded number of - // connections. + /// Packs frames into requests under the caps and posts them. + /// + /// Concurrently, not one after another. A call of up to 3.2M points packs into a single + /// request and this changes nothing, but a larger one becomes several and the api validates + /// them independently, so there is no reason to serialise them. Bounded by + /// `request_concurrency` so a very large call cannot open an unbounded number of connections. + pub(crate) async fn post_frame_requests( + &self, + frames: Vec, + options: &BinaryIngestOptions, + ) -> Result<(), ResponseError> { let path = format!("{}/data/binary", self.base_url); let bodies = pack_requests(frames); let limit = options.request_concurrency.max(1); @@ -836,9 +982,7 @@ impl TimeSeriesService { outcome?; } } - let mut result = DataWrapper::new(); - result.set_http_status_code(204); - Ok(result) + Ok(()) } async fn post_frames( @@ -867,17 +1011,18 @@ impl TimeSeriesService { } } - /// One [`ResolvedSeries`] per collection, in order, from the cache or `/timeseries/byids`. - async fn resolve_series( + /// One [`ResolvedSeries`] per `(id, external id)` asked for, in order, from the cache or + /// `/timeseries/byids`; the id wins when both are given. + pub(crate) async fn resolve_series( &self, - collections: &[DatapointsCollection], + wanted: &[(Option, Option<&str>)], ) -> Result, ResponseError> { let mut lookups: Vec = Vec::new(); let mut seen: HashSet = HashSet::new(); { let cache = self.binary_series.lock().unwrap(); - for collection in collections { - match (collection.id, &collection.external_id) { + for &(id, external_id) in wanted { + match (id, external_id) { (Some(id), _) => { if !cache.by_id.contains_key(&id) && seen.insert(format!("id:{id}")) { lookups.push(IdAndExtId::from_id(id)); @@ -916,10 +1061,10 @@ impl TimeSeriesService { } let cache = self.binary_series.lock().unwrap(); - let mut resolved = Vec::with_capacity(collections.len()); + let mut resolved = Vec::with_capacity(wanted.len()); let mut missing = Vec::new(); - for collection in collections { - let hit = match (collection.id, &collection.external_id) { + for &(id, external_id) in wanted { + let hit = match (id, external_id) { (Some(id), _) => cache.by_id.get(&id), (None, Some(external_id)) => cache.by_external_id.get(external_id), (None, None) => None, @@ -927,10 +1072,8 @@ impl TimeSeriesService { match hit { Some(series) => resolved.push(series.clone()), None => missing.push( - collection - .id - .map(|id| id.to_string()) - .or_else(|| collection.external_id.clone()) + id.map(|id| id.to_string()) + .or_else(|| external_id.map(str::to_string)) .unwrap_or_default(), ), } @@ -1060,6 +1203,43 @@ mod tests { assert_eq!(values.values(), &[1.5, 2.0, 2.5]); } + /// Compression 2 stores what compression 1 does, and its raw length is the buffers' + /// uncompressed, padded to 8: compressible ones carry their length, tiny ones -1 and the bytes. + #[test] + fn buffers_compression_reads_back_as_payload_compression() { + use crate::timeseries::binary_read::{parse_read_response, END_OF_RESPONSE}; + let cases: [(DatapointValueType, &[&str]); 7] = [ + (DatapointValueType::Bigint, &["1", "-9000000000", "42"]), + (DatapointValueType::Float, &["1.5", "-2.25", "NaN"]), + (DatapointValueType::Float32, &["0.5", "3"]), + (DatapointValueType::Numeric, &["123456789012.345678", "-0.000001"]), + (DatapointValueType::Decimal32, &["12345.6789", "99999999"]), + (DatapointValueType::Text, &["open", "blåbær"]), + (DatapointValueType::Mixed, &["3.5", "fault", "-7"]), + ]; + for (value_type, values) in cases { + for rows in [values.len(), 5000] { + let mut writer = FrameWriter::new(value_type); + writer.series(1, "s").unwrap(); + for t in 0..rows { + writer.add(1, t as i64 * 1000, values[t % values.len()]).unwrap(); + } + let read = |frame: Frame| { + let mut body = frame.bytes; + body.extend_from_slice(END_OF_RESPONSE); + parse_read_response(&body).unwrap() + }; + let payload = writer.build_with(FrameCompression::Payload, 3).unwrap(); + let buffers = writer.build_with(FrameCompression::Buffers, 3).unwrap(); + assert_eq!(buffers.bytes[7], COMPRESSION_ARROW_ZSTD); + assert_eq!(u32_at(&buffers.bytes, 24) as usize, buffers.raw_len); + assert!(buffers.raw_len <= payload.raw_len, "{value_type:?}: buffers are part of the stream"); + let (buffers, payload) = (format!("{:?}", read(buffers)), format!("{:?}", read(payload))); + assert_eq!(buffers, payload, "{value_type:?}, {rows} rows"); + } + } + } + #[test] fn every_value_type_writes_its_canonical_schema() { for value_type in [ @@ -1212,7 +1392,7 @@ mod tests { #[test] fn options_accept_only_the_three_levels() { assert!(BinaryIngestOptions::default().validate().is_ok()); - assert_eq!(BinaryIngestOptions::default().zstd_level, 9); + assert_eq!(BinaryIngestOptions::default().zstd_level, 3); assert!(BinaryIngestOptions::new().zstd_level(3).validate().is_ok()); assert!(BinaryIngestOptions::new().zstd_level(0).validate().is_err()); assert!(BinaryIngestOptions::new().zstd_level(22).validate().is_err()); diff --git a/src/timeseries/binary_read.rs b/src/timeseries/binary_read.rs new file mode 100644 index 0000000..ac1c55e --- /dev/null +++ b/src/timeseries/binary_read.rs @@ -0,0 +1,1196 @@ +//! The binary datapoint read behind `POST /timeseries/data/list/binary`. +//! +//! The request is the JSON retrieve body; the answer is the frames `POST /timeseries/data/binary` +//! takes, one series per frame, concatenated, and ended by the 4-byte marker `DHDE`. A body without +//! the marker was cut short, whether inside a frame or between two, and is an error rather than a +//! shorter answer. Codec 2 frames carry bucketed aggregates instead of values. A read's frames are +//! compression 2, each record batch compressing its own buffers, but compression 1 is read too. +//! +//! [`TimeSeriesService::stream_datapoints_binary`] hands the frames over as they arrive, decoded +//! on blocking threads while the rest download; [`TimeSeriesService::retrieve_datapoints_binary`] +//! collects that stream into one [`BinaryDatapoints`] per series. [`parse_read_response`] does the +//! same for a body already in memory. + +use std::io::Cursor; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow_array::{ + Array, Decimal128Array, Float32Array, Float64Array, Int64Array, RecordBatch, StringArray, + TimestampMillisecondArray, +}; +use arrow_ipc::reader::StreamReader; +use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use bytes::{Buf, Bytes, BytesMut}; +use futures::stream::{FuturesOrdered, Stream, StreamExt}; +use tokio::sync::mpsc; +use oauth2::http::StatusCode; + +use crate::generic::{ApiServiceProvider, DataWrapper, RetrieveFilter}; +use crate::http::ResponseError; +use crate::timeseries::binary::{ + DatapointValueType, CODEC_ARROW_IPC, COMPRESSION_ARROW_ZSTD, COMPRESSION_ZSTD, HEADER_BYTES, + MAGIC, MAX_FRAME_RAW_BYTES, MEDIA_TYPE, VERSION, +}; +use crate::timeseries::TimeSeriesService; + +/// Ends a binary read's response, after the last frame. +pub const END_OF_RESPONSE: &[u8; 4] = b"DHDE"; +/// A response frame whose payload holds bucketed aggregates in [`aggregates_schema`]. +pub const CODEC_ARROW_IPC_AGGREGATES: u8 = 2; +/// Buckets per aggregate frame, whatever the series' type. +pub const MAX_AGGREGATE_ROWS: usize = 50_000; + +/// The schema of an aggregate frame: the bucket start and four Float64s, for every value type. +pub fn aggregates_schema() -> SchemaRef { + let float = |name: &str| Field::new(name, DataType::Float64, false); + Arc::new(Schema::new(vec![ + Field::new("timeseries_id", DataType::Int64, false), + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + false, + ), + float("avg"), + float("min"), + float("max"), + float("sum"), + ])) +} + +/// One series' rows from a binary read, a column per field rather than an object per point. +#[derive(Debug, Clone, PartialEq)] +pub struct BinaryDatapoints { + pub id: u64, + pub external_id: String, + pub value_type: DatapointValueType, + /// Epoch milliseconds, oldest first; for aggregates, the start of each bucket. + pub timestamps: Vec, + pub values: BinaryValues, +} + +impl BinaryDatapoints { + pub fn len(&self) -> usize { + self.timestamps.len() + } + + pub fn is_empty(&self) -> bool { + self.timestamps.is_empty() + } + + pub fn is_aggregates(&self) -> bool { + matches!(self.values, BinaryValues::Aggregates { .. }) + } +} + +/// The value column of a [`BinaryDatapoints`], typed by the series' value type. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum BinaryValues { + Bigint(Vec), + Float(Vec), + Float32(Vec), + /// `numeric` and `decimal32`, exact: each value is `unscaled / 10^scale`. + Decimal { unscaled: Vec, scale: i8 }, + Text(Vec), + /// Each row sets exactly one side. + Mixed { + numeric: Vec>, + text: Vec>, + }, + /// Time-weighted `avg`, `min`, `max` and `sum` per bucket. All four come back whichever the + /// request named. + Aggregates { + avg: Vec, + min: Vec, + max: Vec, + sum: Vec, + }, +} + +impl BinaryValues { + fn reserve_exact(&mut self, additional: usize) { + match self { + Self::Bigint(v) => v.reserve_exact(additional), + Self::Float(v) => v.reserve_exact(additional), + Self::Float32(v) => v.reserve_exact(additional), + Self::Decimal { unscaled, .. } => unscaled.reserve_exact(additional), + Self::Text(v) => v.reserve_exact(additional), + Self::Mixed { numeric, text } => { + numeric.reserve_exact(additional); + text.reserve_exact(additional); + } + Self::Aggregates { avg, min, max, sum } => { + for v in [avg, min, max, sum] { + v.reserve_exact(additional); + } + } + } + } + + /// Appends `other`, which must be the same variant; false if it is not. + fn append(&mut self, other: BinaryValues) -> bool { + match (self, other) { + (Self::Bigint(a), Self::Bigint(b)) => a.extend(b), + (Self::Float(a), Self::Float(b)) => a.extend(b), + (Self::Float32(a), Self::Float32(b)) => a.extend(b), + ( + Self::Decimal { unscaled: a, scale: sa }, + Self::Decimal { unscaled: b, scale: sb }, + ) if *sa == sb => a.extend(b), + (Self::Text(a), Self::Text(b)) => a.extend(b), + ( + Self::Mixed { numeric: an, text: at }, + Self::Mixed { numeric: bn, text: bt }, + ) => { + an.extend(bn); + at.extend(bt); + } + ( + Self::Aggregates { avg: aa, min: ai, max: ax, sum: asum }, + Self::Aggregates { avg: ba, min: bi, max: bx, sum: bsum }, + ) => { + aa.extend(ba); + ai.extend(bi); + ax.extend(bx); + asum.extend(bsum); + } + _ => return false, + } + true + } +} + +fn u32_at(body: &[u8], offset: usize) -> usize { + u32::from_le_bytes(body[offset..offset + 4].try_into().unwrap()) as usize +} + +fn read_varint(body: &[u8], pos: &mut usize) -> Option { + let mut value = 0usize; + for shift in (0..35).step_by(7) { + let byte = *body.get(*pos)?; + *pos += 1; + value |= ((byte & 0x7F) as usize) << shift; + if byte & 0x80 == 0 { + return Some(value); + } + } + None +} + +/// Parses a binary read's whole response body: one entry per series in the order the frames came, +/// which is the order of the request's items, with consecutive frames of one series merged. +/// +/// The envelopes are walked first, which is cheap and finds a body that was cut short; the frames +/// are then decompressed and decoded on several threads, since zstd is most of the cost. +/// +/// Errors on anything short of a complete, well-formed answer: a missing end marker, bytes after +/// it, a frame that breaks the envelope, or a payload outside its canon. +pub fn parse_read_response(body: &[u8]) -> Result, String> { + let mut assembler = FrameAssembler::new(); + assembler.push(body); + let mut frames = Vec::new(); + while let Some(frame) = assembler.next_frame()? { + frames.push(frame); + } + assembler.finish()?; + let decoded = decode_frames(&frames)?; + merge_series(decoded.into_iter().flatten()) +} + +/// Merges consecutive parts of one series, the way they come off the frames, into one entry each. +fn merge_series( + parts: impl IntoIterator, +) -> Result, String> { + // Grouped before merging so each series is allocated once at its final size: extending part + // by part regrows and recopies a long series dozens of times. + let mut groups: Vec> = Vec::new(); + for part in parts { + let continues = groups.last().and_then(|g| g.last()).is_some_and(|last| { + last.id == part.id + && last.is_aggregates() == part.is_aggregates() + && last.timestamps.last() < part.timestamps.first() + }); + if continues { + groups.last_mut().unwrap().push(part); + } else { + groups.push(vec![part]); + } + } + groups.into_iter().map(merge).collect() +} + +/// One series from its consecutive parts, in order. +fn merge(parts: Vec) -> Result { + let rows: usize = parts.iter().map(BinaryDatapoints::len).sum(); + let mut parts = parts.into_iter(); + let mut series = parts.next().expect("a group has at least one part"); + series.timestamps.reserve_exact(rows - series.len()); + series.values.reserve_exact(rows - series.len()); + for part in parts { + let BinaryDatapoints { timestamps, values, .. } = part; + series.timestamps.extend(timestamps); + if !series.values.append(values) { + return Err(format!("series {} changes shape between frames", series.id)); + } + } + Ok(series) +} + +fn cut_short(index: usize) -> String { + format!("the response ends without its end marker, at frame {index}; it was cut short") +} + +/// Cuts a response into frames as its bytes arrive, in chunks of any size: a frame's header, its +/// directory and its payload may each be split anywhere. +struct FrameAssembler { + buf: BytesMut, + index: usize, + ended: bool, +} + +impl FrameAssembler { + fn new() -> Self { + FrameAssembler { + buf: BytesMut::new(), + index: 0, + ended: false, + } + } + + fn push(&mut self, chunk: &[u8]) { + self.buf.extend_from_slice(chunk); + } + + /// The next complete frame; `None` until more bytes arrive, and after the end marker. + fn next_frame(&mut self) -> Result, String> { + if !self.ended && self.buf.starts_with(END_OF_RESPONSE) { + self.ended = true; + self.buf.advance(END_OF_RESPONSE.len()); + } + if self.ended { + if !self.buf.is_empty() { + return Err(format!("{} bytes after the end marker", self.buf.len())); + } + return Ok(None); + } + if self.buf.len() >= MAGIC.len() && !self.buf.starts_with(MAGIC) { + return Err(format!("frame {}: does not start with DHDP", self.index)); + } + if self.buf.len() < HEADER_BYTES { + return Ok(None); + } + let total = check_header(&self.buf[..HEADER_BYTES], self.index)?; + if self.buf.len() < total { + return Ok(None); + } + let frame = FrameView::parse(self.buf.split_to(total).freeze(), self.index)?; + self.index += 1; + Ok(Some(frame)) + } + + /// After the last byte: the body has to have ended with its marker and nothing after it. + fn finish(&mut self) -> Result<(), String> { + if self.next_frame()?.is_some() || !self.ended { + return Err(cut_short(self.index)); + } + Ok(()) + } +} + +/// Checks a frame's 28-byte header and returns the frame's whole length. +fn check_header(header: &[u8], index: usize) -> Result { + let fail = |message: String| format!("frame {index}: {message}"); + if header[4] != VERSION { + return Err(fail(format!("version {}", header[4]))); + } + let value_type = DatapointValueType::from_id(header[5]) + .ok_or_else(|| fail(format!("value type id {}", header[5])))?; + let aggregates = match header[6] { + CODEC_ARROW_IPC => false, + CODEC_ARROW_IPC_AGGREGATES => true, + codec => return Err(fail(format!("codec {codec}"))), + }; + if header[7] != COMPRESSION_ZSTD && header[7] != COMPRESSION_ARROW_ZSTD { + return Err(fail(format!("compression {}", header[7]))); + } + let rows = u32_at(header, 8); + let max_rows = if aggregates { + MAX_AGGREGATE_ROWS + } else { + value_type.max_rows() + }; + if rows < 1 || rows > max_rows { + return Err(fail(format!("{rows} rows, allowed 1 to {max_rows}"))); + } + let raw_len = u32_at(header, 24); + if raw_len > MAX_FRAME_RAW_BYTES { + return Err(fail(format!( + "payload is {raw_len} bytes raw, the cap is {MAX_FRAME_RAW_BYTES}" + ))); + } + Ok(HEADER_BYTES + u32_at(header, 16) + u32_at(header, 20)) +} + +/// One complete frame whose header passed [`check_header`], its payload still compressed. Owns its +/// bytes, so it can be decoded on another thread. +struct FrameView { + index: usize, + value_type: DatapointValueType, + aggregates: bool, + compression: u8, + rows: usize, + names: Vec<(i64, String)>, + bytes: Bytes, + payload_start: usize, + /// The payload decompressed at compression 1. At compression 2 it is the batches' buffers + /// uncompressed, which arrow-rs does not report, so only its cap is checked; each buffer's own + /// zstd frame and the row count still are. + raw_len: usize, +} + +impl FrameView { + fn parse(bytes: Bytes, index: usize) -> Result { + let fail = |message: String| format!("frame {index}: {message}"); + let series_count = u32_at(&bytes, 12); + let directory_len = u32_at(&bytes, 16); + let directory = &bytes[HEADER_BYTES..HEADER_BYTES + directory_len]; + let mut names: Vec<(i64, String)> = Vec::with_capacity(series_count); + let mut at = 0; + for _ in 0..series_count { + let id = directory + .get(at..at + 8) + .map(|b| i64::from_le_bytes(b.try_into().unwrap())) + .ok_or_else(|| fail("directory is shorter than its series count".to_string()))?; + at += 8; + let len = read_varint(directory, &mut at) + .ok_or_else(|| fail("directory is shorter than its series count".to_string()))?; + let name = directory + .get(at..at + len) + .ok_or_else(|| fail("directory is shorter than its series count".to_string()))?; + at += len; + let name = String::from_utf8(name.to_vec()) + .map_err(|_| fail(format!("external id of series {id} is not UTF-8")))?; + names.push((id, name)); + } + if at != directory_len { + return Err(fail(format!("{} bytes after the directory", directory_len - at))); + } + let aggregates = bytes[6] == CODEC_ARROW_IPC_AGGREGATES; + Ok(FrameView { + index, + value_type: DatapointValueType::from_id(bytes[5]).expect("checked with the header"), + aggregates, + compression: bytes[7], + rows: u32_at(&bytes, 8), + names, + payload_start: HEADER_BYTES + directory_len, + raw_len: u32_at(&bytes, 24), + bytes, + }) + } +} + +/// Decodes every frame, on as many threads as there are cores and frames, keeping their order. +/// On failure, the error of the earliest frame that failed. +fn decode_frames(frames: &[FrameView]) -> Result>, String> { + let threads = std::thread::available_parallelism() + .map_or(1, |n| n.get()) + .min(frames.len()); + if threads <= 1 { + return frames.iter().map(decode_frame).collect(); + } + let next = AtomicUsize::new(0); + let mut decoded: Vec<(usize, Result, String>)> = + std::thread::scope(|scope| { + let workers: Vec<_> = (0..threads) + .map(|_| { + scope.spawn(|| { + let mut done = Vec::new(); + loop { + let i = next.fetch_add(1, Ordering::Relaxed); + let Some(frame) = frames.get(i) else { return done }; + done.push((i, decode_frame(frame))); + } + }) + }) + .collect(); + workers + .into_iter() + .flat_map(|w| w.join().expect("a frame decoder panicked")) + .collect() + }); + decoded.sort_unstable_by_key(|(i, _)| *i); + decoded.into_iter().map(|(_, result)| result).collect() +} + +/// One frame's rows, split into one part per series it names. +fn decode_frame(frame: &FrameView) -> Result, String> { + let fail = |message: String| format!("frame {}: {message}", frame.index); + let payload = frame.bytes.slice(frame.payload_start..); + let payload = if frame.compression == COMPRESSION_ZSTD { + let raw = zstd::bulk::decompress(&payload, frame.raw_len) + .map_err(|e| fail(format!("zstd: {e}")))?; + if raw.len() != frame.raw_len { + return Err(fail(format!( + "payload decompressed to {} bytes, the envelope says {}", + raw.len(), + frame.raw_len + ))); + } + Bytes::from(raw) + } else { + payload + }; + let reader = StreamReader::try_new(Cursor::new(payload), None) + .map_err(|e| fail(format!("arrow: {e}")))?; + let expected = if frame.aggregates { + aggregates_schema() + } else { + frame.value_type.schema() + }; + let schema = reader.schema(); + let same_columns = schema.fields().len() == expected.fields().len() + && schema + .fields() + .iter() + .zip(expected.fields()) + .all(|(got, want)| got.name() == want.name() && got.data_type() == want.data_type()); + if !same_columns { + return Err(fail(format!("schema {schema:?} is not the canon {expected:?}"))); + } + + let mut parts: Vec = Vec::new(); + let mut read = 0; + for batch in reader { + let batch = batch.map_err(|e| fail(format!("arrow: {e}")))?; + read += batch.num_rows(); + split_batch(&batch, frame.value_type, frame.aggregates, &frame.names, &mut parts) + .map_err(fail)?; + } + if read != frame.rows { + return Err(fail(format!( + "{read} rows in the payload, the envelope says {}", + frame.rows + ))); + } + Ok(parts) +} + +type Item = Result; + +/// A binary read as it arrives: one [`BinaryDatapoints`] per frame, in the order the frames come, +/// which is the order of the request's items. Consecutive items may belong to the same series — +/// a series longer than one frame comes as several, oldest first — so an item is a *chunk* of a +/// series, not all of it. +/// +/// Frames are decoded on blocking threads, several at a time, while the rest are still +/// downloading. A refusal is not an item: it is the error [`TimeSeriesService::stream_datapoints_binary`] +/// returns before there is a stream. **A response that breaks off is an error item after the +/// items that did arrive**, so a consumer that acts on items as they come has to treat what it +/// already took as incomplete when one appears; the stream ends after it. Dropping the stream +/// closes the connection. +/// +/// Read it with [`futures::StreamExt::next`]. +pub struct BinaryDatapointStream { + items: mpsc::Receiver, + task: tokio::task::JoinHandle<()>, +} + +impl Stream for BinaryDatapointStream { + type Item = Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.items.poll_recv(cx) + } +} + +impl Drop for BinaryDatapointStream { + fn drop(&mut self) { + self.task.abort(); + } +} + +/// Frames decoded at once. Past a handful the decoders only compete for memory bandwidth. +fn decoders() -> usize { + std::thread::available_parallelism().map_or(1, |n| n.get()).clamp(1, 8) +} + +fn unreadable(path: &str, message: String) -> ResponseError { + ResponseError { + status: StatusCode::BAD_GATEWAY, + message: format!("POST {path} returned an unreadable answer: {message}"), + content_type: None, + } +} + +/// Reads the body, cuts it into frames, decodes up to [`decoders`] of them at once and sends their +/// parts on in frame order. Whatever fails, the parts of the frames before it are sent first. +async fn produce(response: reqwest::Response, path: String, items: mpsc::Sender) { + let limit = decoders(); + let mut body = response.bytes_stream(); + let mut assembler = FrameAssembler::new(); + let mut pending = FuturesOrdered::new(); + let mut failure: Option = None; + let mut eof = false; + loop { + while failure.is_none() && pending.len() < limit { + match assembler.next_frame() { + Ok(Some(frame)) => { + pending.push_back(tokio::task::spawn_blocking(move || decode_frame(&frame))) + } + Ok(None) => break, + Err(message) => failure = Some(unreadable(&path, message)), + } + } + let reading = !eof && failure.is_none() && pending.len() < limit; + if pending.is_empty() && !reading { + let last = match failure { + Some(error) => Some(error), + None => assembler.finish().err().map(|m| unreadable(&path, m)), + }; + if let Some(error) = last { + let _ = items.send(Err(error)).await; + } + return; + } + tokio::select! { + biased; + Some(decoded) = pending.next(), if !pending.is_empty() => match decoded { + Ok(Ok(parts)) => { + for part in parts { + if items.send(Ok(part)).await.is_err() { + return; + } + } + } + // The earliest frame still in flight, so its error comes before any other. + Ok(Err(message)) => { + failure = Some(unreadable(&path, message)); + pending = FuturesOrdered::new(); + } + Err(join) => { + failure = Some(ResponseError { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: format!("decoding a frame failed: {join}"), + content_type: None, + }); + pending = FuturesOrdered::new(); + } + }, + chunk = body.next(), if reading => match chunk { + Some(Ok(bytes)) => assembler.push(&bytes), + Some(Err(error)) => { + failure = Some(ResponseError { + status: StatusCode::BAD_GATEWAY, + message: format!("the response was cut short while it was being read: {error}"), + content_type: None, + }) + } + None => eof = true, + }, + } + } +} + +/// Appends a batch's rows to `parts`, one part per run of a series id. +fn split_batch( + batch: &RecordBatch, + value_type: DatapointValueType, + aggregates: bool, + names: &[(i64, String)], + parts: &mut Vec, +) -> Result<(), String> { + let ids = column::(batch, 0)?; + let timestamps = column::(batch, 1)?; + let values = batch_values(batch, value_type, aggregates)?; + let mut start = 0; + while start < batch.num_rows() { + let id = ids.value(start); + let mut end = start + 1; + while end < batch.num_rows() && ids.value(end) == id { + end += 1; + } + let external_id = names + .iter() + .find(|(named, _)| *named == id) + .map(|(_, name)| name.clone()) + .ok_or_else(|| format!("series {id} is not in the directory"))?; + let part = BinaryDatapoints { + id: id as u64, + external_id, + value_type, + timestamps: timestamps.values()[start..end].to_vec(), + values: slice_values(&values, start, end), + }; + match parts.last_mut() { + Some(last) if last.id == part.id => { + last.timestamps.extend(part.timestamps); + last.values.append(part.values); + } + _ => parts.push(part), + } + start = end; + } + Ok(()) +} + +fn column(batch: &RecordBatch, index: usize) -> Result<&T, String> { + batch + .column(index) + .as_any() + .downcast_ref::() + .ok_or_else(|| format!("column {index} is not of its canonical type")) +} + +fn batch_values( + batch: &RecordBatch, + value_type: DatapointValueType, + aggregates: bool, +) -> Result { + let floats = |index| column::(batch, index).map(|a| a.values().to_vec()); + if aggregates { + return Ok(BinaryValues::Aggregates { + avg: floats(2)?, + min: floats(3)?, + max: floats(4)?, + sum: floats(5)?, + }); + } + Ok(match value_type { + DatapointValueType::Bigint => { + BinaryValues::Bigint(column::(batch, 2)?.values().to_vec()) + } + DatapointValueType::Float => BinaryValues::Float(floats(2)?), + DatapointValueType::Float32 => { + BinaryValues::Float32(column::(batch, 2)?.values().to_vec()) + } + DatapointValueType::Numeric | DatapointValueType::Decimal32 => { + let array = column::(batch, 2)?; + BinaryValues::Decimal { + unscaled: array.values().to_vec(), + scale: array.scale(), + } + } + DatapointValueType::Text => { + let array = column::(batch, 2)?; + BinaryValues::Text(array.iter().map(|v| v.unwrap_or_default().to_string()).collect()) + } + DatapointValueType::Mixed => { + let numeric = column::(batch, 2)?; + let text = column::(batch, 3)?; + BinaryValues::Mixed { + numeric: numeric.iter().collect(), + text: text.iter().map(|v| v.map(str::to_string)).collect(), + } + } + }) +} + +fn slice_values(values: &BinaryValues, start: usize, end: usize) -> BinaryValues { + match values { + BinaryValues::Bigint(v) => BinaryValues::Bigint(v[start..end].to_vec()), + BinaryValues::Float(v) => BinaryValues::Float(v[start..end].to_vec()), + BinaryValues::Float32(v) => BinaryValues::Float32(v[start..end].to_vec()), + BinaryValues::Decimal { unscaled, scale } => BinaryValues::Decimal { + unscaled: unscaled[start..end].to_vec(), + scale: *scale, + }, + BinaryValues::Text(v) => BinaryValues::Text(v[start..end].to_vec()), + BinaryValues::Mixed { numeric, text } => BinaryValues::Mixed { + numeric: numeric[start..end].to_vec(), + text: text[start..end].to_vec(), + }, + BinaryValues::Aggregates { avg, min, max, sum } => BinaryValues::Aggregates { + avg: avg[start..end].to_vec(), + min: min[start..end].to_vec(), + max: max[start..end].to_vec(), + sum: sum[start..end].to_vec(), + }, + } +} + +impl TimeSeriesService { + /// Reads datapoints through `POST /timeseries/data/list/binary`, the high-throughput twin of + /// [`retrieve_datapoints`](Self::retrieve_datapoints): the same request, answered as + /// zstd-compressed Arrow frames instead of JSON, and returned as one columnar + /// [`BinaryDatapoints`] per series, in the order of the request's items. + /// + /// Each item needs a `start` or an `end`. With a `limit` (at most 100 000) it returns that many + /// of the oldest points from there; without one it returns **every** point in its window, + /// however many, so bound a long series with `end`: the whole answer is held in memory. A + /// series with no points in its window is left out of the answer rather than returned empty. + /// There is no cursor (sending one is a 400): to read on, send the request again with `start` + /// one millisecond past the last timestamp received. An item with `aggregates` and a + /// `granularity` comes back as [`BinaryValues::Aggregates`], carrying all four whichever were + /// named; text series have none. + /// + /// Unlike the JSON read, a series that does not exist or sits in a dataset the caller cannot + /// read is a 400 for the whole request, not an empty result. A response the server could not + /// finish is a 502 whose message says it was cut short, never fewer points; sending the + /// request again is safe. + /// + /// The frames are decoded as they arrive, so this costs little more than the download; for + /// rows as they come, or a window too large to hold, use + /// [`stream_datapoints_binary`](Self::stream_datapoints_binary). + pub async fn retrieve_datapoints_binary( + &self, + json: &DataWrapper, + ) -> Result, ResponseError> { + let path = format!("{}/data/list/binary", self.base_url); + let mut stream = self.stream_datapoints_binary(json).await?; + let mut parts = Vec::new(); + while let Some(part) = stream.next().await { + parts.push(part?); + } + tokio::task::spawn_blocking(move || merge_series(parts)) + .await + .map_err(|e| ResponseError { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: format!("merging the binary read failed: {e}"), + content_type: None, + })? + .map_err(|message| unreadable(&path, message)) + } + + /// The same read as [`retrieve_datapoints_binary`](Self::retrieve_datapoints_binary), handed + /// over frame by frame as the response arrives instead of collected: see + /// [`BinaryDatapointStream`] for what an item is and how a broken-off response shows up. + /// + /// The request is checked before there is a stream, so a refusal (a 400 for an unknown + /// series, a `cursor`, aggregates on a text series) is the `Err` here rather than an item. + /// Holding the stream holds the connection, and the server's query, open: a consumer that + /// stops reading for long can run into the server's timeouts. + pub async fn stream_datapoints_binary( + &self, + json: &DataWrapper, + ) -> Result { + let path = format!("{}/data/list/binary", self.base_url); + let response = self + .execute_post_request_for_stream(&path, json, MEDIA_TYPE) + .await?; + let (tx, items) = mpsc::channel(2 * decoders()); + let task = tokio::spawn(produce(response, path, tx)); + Ok(BinaryDatapointStream { items, task }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::timeseries::binary::FrameWriter; + use arrow_ipc::writer::StreamWriter; + + /// Written by the platform's own frame writer (`DatapointFrameWriter.relayStream`, the code that + /// answers a binary read), from `feat/binary-datapoint-read`: one series per value type, then + /// the end marker, every frame compression 2. What the Rust reader must accept from the Java writer. + const JAVA_VALUES: &[u8] = include_bytes!("fixtures/read_values.bin"); + /// The same series as compression 1 frames, which a read answered with before the relay. + const JAVA_VALUES_ZSTD: &[u8] = include_bytes!("fixtures/read_values_zstd.bin"); + /// One codec-2 frame of three hourly buckets for series 21, `rd_aggregates`. + const JAVA_AGGREGATES: &[u8] = include_bytes!("fixtures/read_aggregates.bin"); + const JAVA_AGGREGATES_ZSTD: &[u8] = include_bytes!("fixtures/read_aggregates_zstd.bin"); + /// Compression 2 with buffers zstd actually shrinks: two record batches of 3000 floats in one + /// frame for series 31, then 12 000 texts for series 32 that the relay cut into two frames. + const JAVA_BATCHES: &[u8] = include_bytes!("fixtures/read_batches.bin"); + + const T0: i64 = 1_700_000_000_000; + + fn response(frames: &[Vec]) -> Vec { + let mut body: Vec = frames.concat(); + body.extend_from_slice(END_OF_RESPONSE); + body + } + + fn float_frame(id: u64, external_id: &str, timestamps: std::ops::Range) -> Vec { + let mut writer = FrameWriter::new(DatapointValueType::Float); + writer.series(id, external_id).unwrap(); + for t in timestamps { + writer.add_float(id, t, t as f64 / 2.0).unwrap(); + } + writer.build(1).unwrap().bytes + } + + /// A frame around any batch, with any codec byte: how a test makes a codec-2 frame. + fn frame_around(value_type: DatapointValueType, codec: u8, id: i64, name: &str, batch: &RecordBatch) -> Vec { + let mut raw = Vec::new(); + let mut writer = StreamWriter::try_new(&mut raw, batch.schema().as_ref()).unwrap(); + writer.write(batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + let compressed = zstd::bulk::compress(&raw, 1).unwrap(); + let mut directory = id.to_le_bytes().to_vec(); + directory.push(name.len() as u8); + directory.extend_from_slice(name.as_bytes()); + let mut frame = MAGIC.to_vec(); + frame.extend_from_slice(&[VERSION, value_type.id(), codec, COMPRESSION_ZSTD]); + for n in [batch.num_rows(), 1, directory.len(), compressed.len(), raw.len()] { + frame.extend_from_slice(&(n as u32).to_le_bytes()); + } + frame.extend_from_slice(&directory); + frame.extend_from_slice(&compressed); + frame + } + + #[test] + fn reads_every_value_type_the_java_writer_answers_with() { + assert_eq!(JAVA_VALUES[7], COMPRESSION_ARROW_ZSTD); + assert_eq!(JAVA_VALUES_ZSTD[7], COMPRESSION_ZSTD); + for body in [JAVA_VALUES, JAVA_VALUES_ZSTD] { + reads_every_value_type(body); + } + } + + fn reads_every_value_type(body: &[u8]) { + let series = parse_read_response(body).unwrap(); + let names: Vec<&str> = series.iter().map(|s| s.external_id.as_str()).collect(); + assert_eq!( + names, + ["rd_bigint", "rd_float", "rd_float32", "rd_numeric", "rd_decimal32", "rd_text", "rd_mixed", "rd_mixed_numbers"] + ); + + let bigint = &series[0]; + assert_eq!((bigint.id, bigint.value_type), (11, DatapointValueType::Bigint)); + assert_eq!(bigint.timestamps, [T0, T0 + 1000, T0 + 2000]); + assert_eq!(bigint.values, BinaryValues::Bigint(vec![1, 42, -9_000_000_000])); + + assert_eq!(series[1].values, BinaryValues::Float(vec![1.5, -2.25])); + assert_eq!(series[2].values, BinaryValues::Float32(vec![0.5])); + assert_eq!( + series[3].values, + BinaryValues::Decimal { unscaled: vec![123_456_789_012_345_678, -1], scale: 6 } + ); + assert_eq!( + series[4].values, + BinaryValues::Decimal { unscaled: vec![123_456_789], scale: 4 } + ); + assert_eq!( + series[5].values, + BinaryValues::Text(vec!["open".to_string(), "blåbær".to_string()]) + ); + assert_eq!( + series[6].values, + BinaryValues::Mixed { + numeric: vec![Some(3.5), None, Some(-7.0)], + text: vec![None, Some("fault".to_string()), None], + } + ); + // No text at all: the writer leaves the validity bitmap out of the numeric column. + assert_eq!( + series[7].values, + BinaryValues::Mixed { numeric: vec![Some(1.0), Some(2.0)], text: vec![None, None] } + ); + } + + #[test] + fn reads_an_aggregate_frame_the_java_writer_answers_with() { + for body in [JAVA_AGGREGATES, JAVA_AGGREGATES_ZSTD] { + reads_an_aggregate_frame(body); + } + } + + fn reads_an_aggregate_frame(body: &[u8]) { + let series = parse_read_response(body).unwrap(); + assert_eq!(series.len(), 1); + let s = &series[0]; + assert_eq!((s.id, s.external_id.as_str()), (21, "rd_aggregates")); + assert_eq!(s.value_type, DatapointValueType::Float, "the series' own type, not the canon's"); + assert!(s.is_aggregates()); + assert_eq!(s.timestamps, [T0, T0 + 3_600_000, T0 + 7_200_000]); + assert_eq!( + s.values, + BinaryValues::Aggregates { + avg: vec![1.5, 2.5, 3.5], + min: vec![1.0, 2.0, 3.0], + max: vec![2.0, 3.0, 4.0], + sum: vec![6.0, 10.0, 14.0], + } + ); + } + + #[test] + fn reads_compressed_batches_several_to_a_frame_and_a_series_cut_across_frames() { + let series = parse_read_response(JAVA_BATCHES).unwrap(); + assert_eq!(series.len(), 2); + let floats = &series[0]; + assert_eq!((floats.id, floats.external_id.as_str()), (31, "rd_float_batches")); + assert_eq!(floats.timestamps, (0..6000).map(|i| T0 + i * 1000).collect::>()); + assert_eq!(floats.values, BinaryValues::Float((0..6000).map(|i| i as f64 * 0.5).collect())); + let texts = &series[1]; + assert_eq!((texts.id, texts.external_id.as_str()), (32, "rd_text_cut")); + assert_eq!(texts.timestamps, (0..12_000).map(|i| T0 + i * 1000).collect::>()); + assert_eq!( + texts.values, + BinaryValues::Text((0..12_000).map(|i| format!("v{}", i % 10)).collect()) + ); + } + + #[test] + fn consecutive_frames_of_one_series_merge_in_order() { + let body = response(&[ + float_frame(1, "a", 0..3), + float_frame(1, "a", 3..5), + float_frame(2, "b", 0..2), + ]); + let series = parse_read_response(&body).unwrap(); + assert_eq!(series.len(), 2); + assert_eq!(series[0].timestamps, [0, 1, 2, 3, 4]); + assert_eq!(series[0].values, BinaryValues::Float(vec![0.0, 0.5, 1.0, 1.5, 2.0])); + assert_eq!(series[1].external_id, "b"); + } + + /// Decoded on several threads, merged in the order the frames came, and a failure reported for + /// the earliest frame that has one, whichever thread reached it first. + #[test] + fn many_frames_decode_in_parallel_and_keep_their_order() { + let mut frames: Vec> = Vec::new(); + for id in 1..=64u64 { + let start = id as i64 * 1_000; + frames.push(float_frame(id, &format!("s{id}"), start..start + 3)); + frames.push(float_frame(id, &format!("s{id}"), start + 3..start + 5)); + } + let series = parse_read_response(&response(&frames)).unwrap(); + assert_eq!(series.len(), 64); + for (i, s) in series.iter().enumerate() { + let id = i as u64 + 1; + assert_eq!((s.id, s.external_id.as_str()), (id, format!("s{id}").as_str())); + let start = id as i64 * 1_000; + assert_eq!(s.timestamps, (start..start + 5).collect::>()); + } + + let batch = RecordBatch::try_new( + DatapointValueType::Float.schema(), + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(TimestampMillisecondArray::from(vec![0]).with_timezone("UTC")), + Arc::new(Float64Array::from(vec![1.0])), + ], + ) + .unwrap(); + let bad = frame_around(DatapointValueType::Float, CODEC_ARROW_IPC_AGGREGATES, 1, "a", &batch); + frames[100] = bad.clone(); + frames[17] = bad; + let error = parse_read_response(&response(&frames)).unwrap_err(); + assert!(error.starts_with("frame 17:"), "{error}"); + } + + fn assemble_in_chunks(body: &[u8], chunk: usize) -> Result, String> { + let mut assembler = FrameAssembler::new(); + let mut frames = Vec::new(); + for piece in body.chunks(chunk) { + assembler.push(piece); + while let Some(frame) = assembler.next_frame()? { + frames.push(frame); + } + } + assembler.finish()?; + merge_series(decode_frames(&frames)?.into_iter().flatten()) + } + + #[test] + fn a_frame_may_be_split_anywhere() { + for body in [JAVA_VALUES, JAVA_VALUES_ZSTD, JAVA_AGGREGATES, JAVA_BATCHES] { + let whole = parse_read_response(body).unwrap(); + for chunk in [1, 3, 7, 28, 1000, 65536] { + assert_eq!(assemble_in_chunks(body, chunk).unwrap(), whole, "chunks of {chunk}"); + } + } + // Cut inside the marker itself: "DH" then nothing. + let mut body = response(&[float_frame(1, "a", 0..2)]); + body.truncate(body.len() - 2); + assert!(assemble_in_chunks(&body, 1).unwrap_err().contains("cut short")); + } + + /// A local HTTP server answering with `chunks` as a chunked body, each written and flushed on + /// its own; the chunk at `gate_after` waits for a permit on `permits`. + async fn serve(chunks: Vec>, permits: tokio::sync::mpsc::Receiver<()>, gate_after: usize) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let mut permits = permits; + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 4096]; + let _ = socket.read(&mut request).await; + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Type: application/vnd.intellistream.datapoint-block\r\nTransfer-Encoding: chunked\r\n\r\n") + .await + .unwrap(); + for (i, chunk) in chunks.iter().enumerate() { + if i == gate_after { + permits.recv().await; + } + socket.write_all(format!("{:x}\r\n", chunk.len()).as_bytes()).await.unwrap(); + socket.write_all(chunk).await.unwrap(); + socket.write_all(b"\r\n").await.unwrap(); + socket.flush().await.unwrap(); + } + let _ = socket.write_all(b"0\r\n\r\n").await; + let _ = socket.shutdown().await; + }); + format!("http://{address}/") + } + + async fn stream_of(url: &str) -> BinaryDatapointStream { + let response = reqwest::get(url).await.unwrap(); + let (tx, items) = mpsc::channel(4); + let task = tokio::spawn(produce(response, "test".to_string(), tx)); + BinaryDatapointStream { items, task } + } + + #[tokio::test] + async fn items_arrive_while_the_body_is_still_coming() { + let first = float_frame(1, "a", 0..3); + let rest = [float_frame(1, "a", 3..5), float_frame(2, "b", 0..2), END_OF_RESPONSE.to_vec()].concat(); + let (release, permits) = tokio::sync::mpsc::channel(1); + let url = serve(vec![first, rest], permits, 1).await; + let mut stream = stream_of(&url).await; + + // The server sends nothing more until this item is in hand. + let item = tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("the first frame did not arrive before the rest of the body") + .unwrap() + .unwrap(); + assert_eq!((item.id, item.timestamps.clone()), (1, vec![0, 1, 2])); + release.send(()).await.unwrap(); + + let rest: Vec = stream.map(|i| i.unwrap()).collect().await; + assert_eq!(rest.iter().map(|s| s.id).collect::>(), [1, 2], "one item per frame"); + let merged = merge_series(std::iter::once(item).chain(rest)).unwrap(); + assert_eq!(merged.len(), 2); + assert_eq!(merged[0].timestamps, [0, 1, 2, 3, 4]); + } + + #[tokio::test] + async fn a_stream_cut_short_ends_with_an_error_after_what_arrived() { + let frames = vec![float_frame(1, "a", 0..2), float_frame(2, "b", 0..2)]; + // Every frame intact and the chunked body properly ended, but no end marker. + let (_release, permits) = tokio::sync::mpsc::channel(1); + let url = serve(frames, permits, usize::MAX).await; + let items: Vec = stream_of(&url).await.collect().await; + assert_eq!(items.len(), 3); + assert!(items[0].is_ok() && items[1].is_ok()); + let error = items[2].as_ref().unwrap_err(); + assert_eq!(error.get_status(), StatusCode::BAD_GATEWAY); + assert!(error.get_message().contains("cut short"), "{}", error.get_message()); + } + + #[tokio::test] + async fn a_bad_frame_ends_the_stream_after_the_frames_before_it() { + let batch = RecordBatch::try_new( + DatapointValueType::Float.schema(), + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(TimestampMillisecondArray::from(vec![0]).with_timezone("UTC")), + Arc::new(Float64Array::from(vec![1.0])), + ], + ) + .unwrap(); + let bad = frame_around(DatapointValueType::Float, CODEC_ARROW_IPC_AGGREGATES, 2, "b", &batch); + let body = response(&[float_frame(1, "a", 0..2), bad, float_frame(3, "c", 0..2)]); + let (_release, permits) = tokio::sync::mpsc::channel(1); + let url = serve(body.chunks(5).map(<[u8]>::to_vec).collect(), permits, usize::MAX).await; + let items: Vec = stream_of(&url).await.collect().await; + assert_eq!(items.len(), 2, "the good frame, then the error; nothing after it"); + assert_eq!(items[0].as_ref().unwrap().id, 1); + assert!(items[1].as_ref().unwrap_err().get_message().contains("frame 1:")); + } + + /// The same series asked for twice, the second window starting before the first ended: two + /// entries, since merging them would put the timestamps out of order. + #[test] + fn a_series_whose_timestamps_restart_is_a_second_entry() { + let body = response(&[float_frame(1, "a", 0..3), float_frame(1, "a", 1..2)]); + let series = parse_read_response(&body).unwrap(); + assert_eq!(series.len(), 2); + assert_eq!(series[1].timestamps, [1]); + } + + #[test] + fn the_end_marker_alone_is_no_points() { + assert!(parse_read_response(END_OF_RESPONSE).unwrap().is_empty()); + } + + #[test] + fn a_response_cut_short_is_an_error_not_fewer_points() { + let first = float_frame(1, "a", 0..3); + let full = response(&[first.clone(), float_frame(2, "b", 0..3)]); + + // Between two frames: every byte that arrived parses, only the marker is missing. + let error = parse_read_response(&full[..first.len()]).unwrap_err(); + assert!(error.contains("cut short"), "{error}"); + // Inside a frame, in the header and in the payload. + for cut in [first.len() + 10, first.len() + HEADER_BYTES + 12, full.len() - 5] { + let error = parse_read_response(&full[..cut]).unwrap_err(); + assert!(error.contains("cut short"), "cut at {cut}: {error}"); + } + // Nothing at all. + assert!(parse_read_response(&[]).unwrap_err().contains("cut short")); + } + + #[test] + fn bytes_after_the_end_marker_are_refused() { + let mut body = response(&[float_frame(1, "a", 0..2)]); + body.push(0); + assert!(parse_read_response(&body).unwrap_err().contains("after the end marker")); + } + + #[test] + fn a_payload_outside_its_canon_is_refused() { + // Values in a codec-2 frame: right envelope, wrong schema. + let batch = RecordBatch::try_new( + DatapointValueType::Float.schema(), + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(TimestampMillisecondArray::from(vec![0]).with_timezone("UTC")), + Arc::new(Float64Array::from(vec![1.0])), + ], + ) + .unwrap(); + let body = response(&[frame_around(DatapointValueType::Float, CODEC_ARROW_IPC_AGGREGATES, 1, "a", &batch)]); + assert!(parse_read_response(&body).unwrap_err().contains("canon")); + + let body = response(&[frame_around(DatapointValueType::Float, 9, 1, "a", &batch)]); + assert!(parse_read_response(&body).unwrap_err().contains("codec 9")); + } + + #[test] + fn a_row_not_named_in_the_directory_is_refused() { + let batch = RecordBatch::try_new( + DatapointValueType::Float.schema(), + vec![ + Arc::new(Int64Array::from(vec![2])), + Arc::new(TimestampMillisecondArray::from(vec![0]).with_timezone("UTC")), + Arc::new(Float64Array::from(vec![1.0])), + ], + ) + .unwrap(); + let body = response(&[frame_around(DatapointValueType::Float, CODEC_ARROW_IPC, 1, "a", &batch)]); + assert!(parse_read_response(&body).unwrap_err().contains("not in the directory")); + } + + #[test] + fn the_aggregate_row_cap_is_its_own() { + let n = MAX_AGGREGATE_ROWS; + let floats = || Arc::new(Float64Array::from(vec![1.0; n])) as Arc; + let batch = RecordBatch::try_new( + aggregates_schema(), + vec![ + Arc::new(Int64Array::from(vec![1; n])), + Arc::new(TimestampMillisecondArray::from((0..n as i64).collect::>()).with_timezone("UTC")), + floats(), + floats(), + floats(), + floats(), + ], + ) + .unwrap(); + // Text's value cap is 10 000; an aggregate frame of a mixed series may still hold 50 000. + let body = response(&[frame_around(DatapointValueType::Mixed, CODEC_ARROW_IPC_AGGREGATES, 1, "a", &batch)]); + let series = parse_read_response(&body).unwrap(); + assert_eq!(series[0].len(), n); + } +} diff --git a/src/timeseries/binary_table.rs b/src/timeseries/binary_table.rs new file mode 100644 index 0000000..3295cab --- /dev/null +++ b/src/timeseries/binary_table.rs @@ -0,0 +1,715 @@ +//! Binary ingest straight from Arrow: a wide table, one timestamp column and one column per +//! series named by its external id, as polars, pyarrow and pandas all hand it over. +//! +//! The columns are used as they are wherever they already have the canon's type (a `Float64` +//! column for a `float` series is sliced, never copied); timestamps get one pass to epoch +//! milliseconds, the per-row `timeseries_id` column is materialised per frame, and the IPC +//! encoding and zstd are the serialisation. What this path does not do is make an object, or a +//! string, per value, which is what the string path costs. +//! +//! Values reach the series as the string path would store them: decimals round half-up at the +//! series' scale and `decimal32` clamps; text is held to the same limits. Two differences are +//! deliberate. A null is "no point at this timestamp", which is what makes a ragged wide table +//! work. And a `mixed` series takes a typed column as it is typed: a numeric column is +//! `value_numeric`, a string column `value_text`, where the string path classifies each string. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{ + Decimal128Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, + TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, + TimestampSecondType, UInt16Type, UInt32Type, UInt64Type, UInt8Type, +}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, Decimal128Array, Float32Array, Float64Array, Int64Array, + RecordBatch, StringArray, TimestampMillisecondArray, UInt32Array, +}; +use arrow_schema::{DataType, TimeUnit}; +use oauth2::http::StatusCode; + +use crate::http::ResponseError; +use crate::timeseries::binary::{ + assemble_frame, checked_text, is_stale_series_rejection, parse_scaled, + unprocessable, BinaryIngestOptions, DatapointValueType, Frame, FrameCompression, ResolvedSeries, + DECIMAL32_UNSCALED_MAX, MAX_SERIES_PER_FRAME, NUMERIC_UNSCALED_MAX, +}; +use crate::timeseries::TimeSeriesService; + +/// What [`TimeSeriesService::insert_table_binary`] sent. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct TableIngest { + /// Points sent: non-null cells, after a repeated timestamp in one series kept its last value. + pub points: usize, + pub series: usize, + pub frames: usize, + /// `decimal32` values clamped to the type's range. + pub clamped: usize, +} + +/// One series' points from one batch, in the canon: timestamps ascending and unique, values in the +/// column type(s) the series' frames carry. +struct Points { + series: usize, + timestamps: ArrayRef, + values: Vec, +} + +impl Points { + fn len(&self) -> usize { + self.timestamps.len() + } +} + +/// Epoch milliseconds from a timestamp column. Copy-free when it is already in milliseconds. +fn timestamps_ms(name: &str, column: &ArrayRef) -> Result { + if column.null_count() > 0 { + return Err(format!("timestamp column {name} has nulls")); + } + let exact = |values: &[i64], per_ms: i64| -> Result { + let mut out = Vec::with_capacity(values.len()); + for &v in values { + if v % per_ms != 0 { + return Err(format!( + "timestamp column {name} has a value finer than a millisecond; \ + the series store milliseconds, so round it first" + )); + } + out.push(v / per_ms); + } + Ok(Arc::new(TimestampMillisecondArray::from(out).with_timezone("UTC"))) + }; + match column.data_type() { + DataType::Timestamp(_, None) => Err(format!( + "timestamp column {name} has no time zone; give it one (e.g. UTC) rather than leave \ + the instant to be guessed" + )), + DataType::Timestamp(TimeUnit::Millisecond, Some(_)) => Ok(Arc::new( + TimestampMillisecondArray::new(column.as_primitive::().values().clone(), None) + .with_timezone("UTC"), + )), + DataType::Timestamp(TimeUnit::Microsecond, Some(_)) => { + exact(column.as_primitive::().values(), 1_000) + } + DataType::Timestamp(TimeUnit::Nanosecond, Some(_)) => { + exact(column.as_primitive::().values(), 1_000_000) + } + DataType::Timestamp(TimeUnit::Second, Some(_)) => Ok(Arc::new( + TimestampMillisecondArray::from_iter_values( + column.as_primitive::().values().iter().map(|s| s * 1_000), + ) + .with_timezone("UTC"), + )), + DataType::Int64 => Ok(Arc::new( + TimestampMillisecondArray::new(column.as_primitive::().values().clone(), None) + .with_timezone("UTC"), + )), + other => Err(format!( + "timestamp column {name} is {other}; it has to be a timestamp with a time zone, or \ + Int64 epoch milliseconds" + )), + } +} + +/// A numeric column as `f64`, for `float` and `mixed`. `None` when the column is not numeric. +fn as_f64(column: &ArrayRef) -> Option { + macro_rules! widen { + ($t:ty) => { + Float64Array::from_iter_values(column.as_primitive::<$t>().values().iter().map(|&v| v as f64)) + }; + } + Some(match column.data_type() { + DataType::Float64 => column.as_primitive::().clone(), + DataType::Float32 => widen!(Float32Type), + DataType::Int8 => widen!(Int8Type), + DataType::Int16 => widen!(Int16Type), + DataType::Int32 => widen!(Int32Type), + DataType::Int64 => widen!(Int64Type), + DataType::UInt8 => widen!(UInt8Type), + DataType::UInt16 => widen!(UInt16Type), + DataType::UInt32 => widen!(UInt32Type), + DataType::UInt64 => widen!(UInt64Type), + DataType::Decimal128(_, scale) => { + let divisor = 10f64.powi(*scale as i32); + Float64Array::from_iter_values( + column.as_primitive::().values().iter().map(|&v| v as f64 / divisor), + ) + } + _ => return None, + }) +} + +/// A string column of any of the three string layouts, each value held to the text limits. +fn as_text(name: &str, column: &ArrayRef) -> Option> { + let checked = |values: &mut dyn Iterator>| -> Result { + values + .map(|v| checked_text(v.unwrap_or_default()).map_err(|e| format!("column {name}: {e}"))) + .collect::, String>>() + .map(StringArray::from) + }; + Some(match column.data_type() { + DataType::Utf8 => checked(&mut column.as_string::().iter()), + DataType::LargeUtf8 => checked(&mut column.as_string::().iter()), + DataType::Utf8View => checked(&mut column.as_string_view().iter()), + _ => return None, + }) +} + +/// Unscaled values at `scale` decimals, rounded half-up like the string path's `parse_scaled`. +fn as_unscaled(name: &str, column: &ArrayRef, scale: i8) -> Result, String> { + let rescale = |v: i128, from: i8| -> i128 { + if from <= scale { + v * 10i128.pow((scale - from) as u32) + } else { + let divisor = 10i128.pow((from - scale) as u32); + let (quotient, remainder) = (v / divisor, v % divisor); + if remainder.abs() * 2 >= divisor { + quotient + v.signum() + } else { + quotient + } + } + }; + macro_rules! ints { + ($t:ty) => { + column.as_primitive::<$t>().values().iter().map(|&v| rescale(v as i128, 0)).collect() + }; + } + let floats = |values: &mut dyn Iterator| -> Result, String> { + values + .map(|v| { + parse_scaled(&v.to_string(), scale as u32) + .ok_or_else(|| format!("column {name}: {v} is not a decimal number")) + }) + .collect() + }; + Ok(match column.data_type() { + DataType::Decimal128(_, from) => { + let from = *from; + column.as_primitive::().values().iter().map(|&v| rescale(v, from)).collect() + } + DataType::Int8 => ints!(Int8Type), + DataType::Int16 => ints!(Int16Type), + DataType::Int32 => ints!(Int32Type), + DataType::Int64 => ints!(Int64Type), + DataType::UInt8 => ints!(UInt8Type), + DataType::UInt16 => ints!(UInt16Type), + DataType::UInt32 => ints!(UInt32Type), + DataType::UInt64 => ints!(UInt64Type), + DataType::Float64 => floats(&mut column.as_primitive::().values().iter().copied())?, + DataType::Float32 => floats(&mut column.as_primitive::().values().iter().map(|&v| v as f64))?, + other => return Err(format!("column {name} is {other}, which is not a number")), + }) +} + +/// A column with no nulls, in the canon's value type(s) for `value_type`. +fn to_canon( + name: &str, + value_type: DatapointValueType, + column: &ArrayRef, + clamped: &mut usize, +) -> Result, String> { + let wrong = || format!("column {name} is {}, which a {value_type:?} series cannot hold", column.data_type()); + Ok(match value_type { + DatapointValueType::Float => vec![Arc::new(as_f64(column).ok_or_else(wrong)?) as ArrayRef], + DatapointValueType::Float32 => { + let values: Float32Array = match column.data_type() { + DataType::Float32 => column.as_primitive::().clone(), + _ => { + let wide = as_f64(column).ok_or_else(wrong)?; + Float32Array::from_iter_values(wide.values().iter().map(|&v| v as f32)) + } + }; + vec![Arc::new(values)] + } + DatapointValueType::Bigint => { + macro_rules! widen { + ($t:ty) => { + Int64Array::from_iter_values(column.as_primitive::<$t>().values().iter().map(|&v| v as i64)) + }; + } + let values = match column.data_type() { + DataType::Int64 => column.as_primitive::().clone(), + DataType::Int8 => widen!(Int8Type), + DataType::Int16 => widen!(Int16Type), + DataType::Int32 => widen!(Int32Type), + DataType::UInt8 => widen!(UInt8Type), + DataType::UInt16 => widen!(UInt16Type), + DataType::UInt32 => widen!(UInt32Type), + DataType::UInt64 => { + let values = column.as_primitive::().values(); + if values.iter().any(|&v| v > i64::MAX as u64) { + return Err(format!("column {name} has a value past the bigint range")); + } + widen!(UInt64Type) + } + _ => return Err(wrong()), + }; + vec![Arc::new(values)] + } + DatapointValueType::Numeric | DatapointValueType::Decimal32 => { + let (precision, scale) = if value_type == DatapointValueType::Numeric { (18, 6) } else { (9, 4) }; + let mut unscaled = as_unscaled(name, column, scale)?; + if value_type == DatapointValueType::Numeric { + if let Some(v) = unscaled.iter().find(|v| v.abs() > NUMERIC_UNSCALED_MAX) { + return Err(format!( + "column {name}: {} does not fit Decimal(18, 6)", + *v as f64 / 1e6 + )); + } + } else { + for v in unscaled.iter_mut() { + let bounded = (*v).clamp(-DECIMAL32_UNSCALED_MAX, DECIMAL32_UNSCALED_MAX); + if bounded != *v { + *clamped += 1; + *v = bounded; + } + } + } + let array = Decimal128Array::from(unscaled) + .with_precision_and_scale(precision, scale) + .map_err(|e| e.to_string())?; + vec![Arc::new(array)] + } + DatapointValueType::Text => vec![Arc::new(as_text(name, column).ok_or_else(wrong)??)], + DatapointValueType::Mixed => { + let rows = column.len(); + if let Some(numeric) = as_f64(column) { + vec![Arc::new(numeric), Arc::new(StringArray::new_null(rows))] + } else if let Some(text) = as_text(name, column) { + vec![Arc::new(Float64Array::new_null(rows)), Arc::new(text?)] + } else { + return Err(wrong()); + } + } + }) +} + +/// One series' column of a batch as its points: nulls dropped, values in the canon, timestamps +/// ascending with a repeated one keeping its last value, as the string path keeps it. +fn points( + index: usize, + series: &ResolvedSeries, + name: &str, + timestamps: &ArrayRef, + column: &ArrayRef, + clamped: &mut usize, +) -> Result { + let (timestamps, column) = if column.null_count() == 0 { + (timestamps.clone(), column.clone()) + } else { + let present = BooleanArray::from_iter((0..column.len()).map(|i| Some(column.is_valid(i)))); + let filter = |a: &ArrayRef| arrow_select::filter::filter(a.as_ref(), &present).map_err(|e| e.to_string()); + (filter(timestamps)?, filter(column)?) + }; + let mut values = to_canon(name, series.value_type, &column, clamped)?; + let mut timestamps = timestamps; + + let ts = timestamps.as_primitive::().values(); + if !ts.windows(2).all(|w| w[0] < w[1]) { + let mut order: Vec = (0..ts.len() as u32).collect(); + order.sort_by_key(|&i| ts[i as usize]); + // Stable, so equal timestamps stay in input order and the last of each run wins. + let kept: Vec = order + .iter() + .enumerate() + .filter(|(at, &i)| order.get(at + 1).is_none_or(|&next| ts[next as usize] != ts[i as usize])) + .map(|(_, &i)| i) + .collect(); + let indices = UInt32Array::from(kept); + let take = |a: &ArrayRef| arrow_select::take::take(a.as_ref(), &indices, None).map_err(|e| e.to_string()); + timestamps = take(×tamps)?; + values = values.iter().map(take).collect::>()?; + } + Ok(Points { series: index, timestamps, values }) +} + +/// One frame's worth of points: slices of series of one value type, in ascending id order. +struct FrameSpec { + value_type: DatapointValueType, + pieces: Vec<(usize, usize, usize)>, +} + +/// Cuts the points into frames: each value type on its own, series in ascending id order, a frame +/// filled to its row cap, and a series longer than the room left continued in the next frame. +fn plan_frames(points: &[Points], series: &[ResolvedSeries]) -> Vec { + let mut order: Vec = (0..points.len()).filter(|&p| points[p].len() > 0).collect(); + order.sort_by_key(|&p| (series[points[p].series].value_type.id(), series[points[p].series].id)); + let mut specs: Vec = Vec::new(); + for p in order { + let value_type = series[points[p].series].value_type; + let max_rows = value_type.max_rows(); + let mut start = 0; + while start < points[p].len() { + let open = specs.last().is_some_and(|f| { + f.value_type == value_type + && f.pieces.iter().map(|(_, _, n)| n).sum::() < max_rows + && f.pieces.len() < MAX_SERIES_PER_FRAME + }); + if !open { + specs.push(FrameSpec { value_type, pieces: Vec::new() }); + } + let frame = specs.last_mut().unwrap(); + let room = max_rows - frame.pieces.iter().map(|(_, _, n)| n).sum::(); + let take = room.min(points[p].len() - start); + frame.pieces.push((p, start, take)); + start += take; + } + } + specs +} + +fn build_frame( + spec: &FrameSpec, + points: &[Points], + series: &[ResolvedSeries], + compression: FrameCompression, + zstd_level: i32, +) -> Result { + let rows: usize = spec.pieces.iter().map(|(_, _, n)| n).sum(); + let mut ids: Vec = Vec::with_capacity(rows); + let mut directory: Vec<(i64, &str)> = Vec::with_capacity(spec.pieces.len()); + for &(p, _, n) in &spec.pieces { + let s = &series[points[p].series]; + ids.extend(std::iter::repeat_n(s.id as i64, n)); + directory.push((s.id as i64, s.external_id.as_str())); + } + // A frame of one slice is the series' own arrays, sliced; several are concatenated. + let column = |pick: &dyn Fn(&Points) -> &ArrayRef| -> Result { + if let [(p, start, n)] = spec.pieces[..] { + return Ok(pick(&points[p]).slice(start, n)); + } + let slices: Vec = spec.pieces.iter().map(|&(p, start, n)| pick(&points[p]).slice(start, n)).collect(); + let refs: Vec<&dyn Array> = slices.iter().map(|a| a.as_ref()).collect(); + arrow_select::concat::concat(&refs).map_err(|e| e.to_string()) + }; + let mut columns: Vec = vec![Arc::new(Int64Array::from(ids)), column(&|p| &p.timestamps)?]; + let width = points[spec.pieces[0].0].values.len(); + for v in 0..width { + columns.push(column(&|p| &p.values[v])?); + } + let batch = RecordBatch::try_new(spec.value_type.schema(), columns).map_err(|e| e.to_string())?; + assemble_frame(spec.value_type, rows, &directory, &batch, rows * 32 + 2048, compression, zstd_level) +} + +/// A batch's frames, compressed on as many threads as there are cores and frames. +fn table_frames( + batch: &RecordBatch, + timestamp: &str, + series: &[ResolvedSeries], + compression: FrameCompression, + zstd_level: i32, +) -> Result<(Vec, TableIngest), String> { + let schema = batch.schema(); + let ts_index = schema.index_of(timestamp).map_err(|_| format!("no column named {timestamp}"))?; + let timestamps = timestamps_ms(timestamp, batch.column(ts_index))?; + let mut clamped = 0; + let mut all = Vec::with_capacity(series.len()); + let mut s = 0; + for (c, field) in schema.fields().iter().enumerate() { + if c == ts_index { + continue; + } + all.push(points(s, &series[s], field.name(), ×tamps, batch.column(c), &mut clamped)?); + s += 1; + } + let specs = plan_frames(&all, series); + let next = AtomicUsize::new(0); + let threads = std::thread::available_parallelism().map_or(1, |n| n.get()).min(specs.len()).max(1); + let mut built: Vec<(usize, Result)> = std::thread::scope(|scope| { + let workers: Vec<_> = (0..threads) + .map(|_| { + scope.spawn(|| { + let mut done = Vec::new(); + loop { + let i = next.fetch_add(1, Ordering::Relaxed); + let Some(spec) = specs.get(i) else { return done }; + done.push((i, build_frame(spec, &all, series, compression, zstd_level))); + } + }) + }) + .collect(); + workers.into_iter().flat_map(|w| w.join().expect("a frame builder panicked")).collect() + }); + built.sort_unstable_by_key(|(i, _)| *i); + let frames = built.into_iter().map(|(_, f)| f).collect::, _>>()?; + let summary = TableIngest { + points: all.iter().map(Points::len).sum(), + series: all.iter().filter(|p| p.len() > 0).count(), + frames: frames.len(), + clamped, + }; + Ok((frames, summary)) +} + +impl TimeSeriesService { + /// Inserts a wide Arrow table through `POST /timeseries/data/binary`: the column named + /// `timestamp` holds the timestamps, and every other column is a series, named by its + /// external id. The batches are taken as they are, one after another; polars, pyarrow and + /// pandas hand a table over this way through the Arrow C stream. + /// + /// The timestamp column is a timestamp with a time zone (any unit, whole milliseconds), or + /// `Int64` epoch milliseconds; a naive one, or one with nulls, is refused. In a series column + /// a null is "no point at this timestamp", so a ragged table works; NaN is a value. A column + /// is converted to its series' value type as the string path would store it (see the module + /// documentation); one that cannot be, such as a float column for a `bigint` series, is a + /// 422 before anything is sent. Rows need not be sorted, and a timestamp repeated within a + /// series keeps its last value. + /// + /// Series are resolved by external id as on the string path: a name that is not a series, + /// or one you cannot read, is a 404 naming it before anything is sent. Each batch is posted + /// before the next is converted, so memory holds one batch's frames. A table larger than one + /// request is not atomic: a refusal can leave the earlier requests stored. + pub async fn insert_table_binary( + &self, + batches: &[RecordBatch], + timestamp: &str, + options: &BinaryIngestOptions, + ) -> Result { + options.validate()?; + let Some(first) = batches.first() else { + return Ok(TableIngest::default()); + }; + let names: Vec = first + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .filter(|n| n != timestamp) + .collect(); + if first.schema().index_of(timestamp).is_err() { + return Err(ResponseError::bad_request(format!( + "no column named {timestamp}; the columns are {names:?}" + ))); + } + let mut total = TableIngest::default(); + for batch in batches { + if batch.schema().fields() != first.schema().fields() { + return Err(ResponseError::bad_request( + "the batches do not share one schema".to_string(), + )); + } + let sent = match self.insert_table_batch(batch, timestamp, &names, options).await { + Err(error) if is_stale_series_rejection(&error) => { + self.evict_binary_series_cache(); + self.insert_table_batch(batch, timestamp, &names, options).await + } + other => other, + }?; + total.points += sent.points; + total.series = total.series.max(sent.series); + total.frames += sent.frames; + total.clamped += sent.clamped; + } + Ok(total) + } + + async fn insert_table_batch( + &self, + batch: &RecordBatch, + timestamp: &str, + names: &[String], + options: &BinaryIngestOptions, + ) -> Result { + let wanted: Vec<(Option, Option<&str>)> = names.iter().map(|n| (None, Some(n.as_str()))).collect(); + let series = self.resolve_series(&wanted).await?; + let (batch, timestamp) = (batch.clone(), timestamp.to_string()); + let (compression, level) = (options.compression, options.zstd_level); + let (frames, summary) = + tokio::task::spawn_blocking(move || table_frames(&batch, ×tamp, &series, compression, level)) + .await + .map_err(|e| ResponseError { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: format!("building frames failed: {e}"), + content_type: None, + })? + .map_err(unprocessable)?; + if !frames.is_empty() { + self.post_frame_requests(frames, options).await?; + } + Ok(summary) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::timeseries::binary::FrameWriter; + use crate::timeseries::binary_read::{parse_read_response, BinaryDatapoints, BinaryValues, END_OF_RESPONSE}; + use arrow_array::{LargeStringArray, StringViewArray, TimestampMicrosecondArray}; + use arrow_schema::{Field, Schema}; + + fn series(id: u64, external_id: &str, value_type: DatapointValueType) -> ResolvedSeries { + ResolvedSeries { id, external_id: external_id.to_string(), value_type } + } + + fn table(timestamps: ArrayRef, columns: Vec<(&str, ArrayRef)>) -> RecordBatch { + let mut fields = vec![Field::new("timestamp", timestamps.data_type().clone(), false)]; + let mut arrays = vec![timestamps]; + for (name, array) in columns { + fields.push(Field::new(name, array.data_type().clone(), true)); + arrays.push(array); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap() + } + + fn ms(values: &[i64]) -> ArrayRef { + Arc::new(TimestampMillisecondArray::from(values.to_vec()).with_timezone("UTC")) + } + + /// What a body of these frames reads back as. + fn decode(frames: &[Frame]) -> Vec { + let mut body: Vec = frames.iter().flat_map(|f| f.bytes.clone()).collect(); + body.extend_from_slice(END_OF_RESPONSE); + parse_read_response(&body).unwrap() + } + + /// The same points through the string path's `FrameWriter`. + fn through_strings(s: &ResolvedSeries, points: &[(i64, &str)]) -> Vec { + let mut writer = FrameWriter::new(s.value_type); + writer.series(s.id, &s.external_id).unwrap(); + for (t, v) in points { + writer.add(s.id, *t, v).unwrap(); + } + decode(&[writer.build(1).unwrap()]) + } + + /// Built both ways, which must read back the same. + fn through_table(batch: &RecordBatch, resolved: &[ResolvedSeries]) -> Vec { + let payload = decode(&table_frames(batch, "timestamp", resolved, FrameCompression::Payload, 1).unwrap().0); + let buffers = decode(&table_frames(batch, "timestamp", resolved, FrameCompression::Buffers, 1).unwrap().0); + assert_eq!(format!("{payload:?}"), format!("{buffers:?}"), "compression 1 and 2 differ"); + payload + } + + #[test] + fn a_float_column_stores_what_the_string_path_stores() { + let s = [series(7, "pump", DatapointValueType::Float)]; + // Unsorted, a repeated timestamp (the last value wins), a null (no point), microseconds. + let ts = Arc::new( + TimestampMicrosecondArray::from(vec![3_000_000, 1_000_000, 2_000_000, 1_000_000, 4_000_000]) + .with_timezone("UTC"), + ); + let values = Arc::new(Float64Array::from(vec![Some(3.5), Some(1.0), None, Some(1.25), Some(f64::NAN)])); + let got = through_table(&table(ts, vec![("pump", values)]), &s); + let want = through_strings(&s[0], &[(3_000, "3.5"), (1_000, "1.0"), (1_000, "1.25"), (4_000, "NaN")]); + assert_eq!(got[0].timestamps, want[0].timestamps); + assert_eq!(got[0].timestamps, [1_000, 3_000, 4_000]); + assert_eq!(format!("{:?}", got[0].values), format!("{:?}", want[0].values)); + } + + #[test] + fn decimals_round_and_clamp_as_the_string_path_does() { + let s = [ + series(1, "numeric_from_float", DatapointValueType::Numeric), + series(2, "decimal32_from_decimal", DatapointValueType::Decimal32), + series(3, "numeric_from_int", DatapointValueType::Numeric), + ]; + let floats = Arc::new(Float64Array::from(vec![0.0000005, -1.2345675, 123.456789])); + let decimals = Arc::new( + Decimal128Array::from(vec![123_456_789_i128, -12_345_678_950, 99_999_999_995]) + .with_precision_and_scale(20, 6) + .unwrap(), + ); + let ints = Arc::new(arrow_array::Int32Array::from(vec![7, -3, 0])); + let batch = table(ms(&[1, 2, 3]), vec![("numeric_from_float", floats), ("decimal32_from_decimal", decimals), ("numeric_from_int", ints)]); + let (frames, summary) = table_frames(&batch, "timestamp", &s, FrameCompression::Payload, 1).unwrap(); + let got = decode(&frames); + + let strings: [&[&str]; 3] = [ + &["0.0000005", "-1.2345675", "123.456789"], + &["123.456789", "-12345.67895", "99999.999995"], + &["7", "-3", "0"], + ]; + for (i, values) in strings.iter().enumerate() { + let points: Vec<(i64, &str)> = values.iter().enumerate().map(|(t, v)| (t as i64 + 1, *v)).collect(); + let want = through_strings(&s[i], &points); + let got = got.iter().find(|g| g.id == s[i].id).unwrap(); + assert_eq!(got.values, want[0].values, "{}", s[i].external_id); + } + // 99999.999995 rounds past decimal32's range and clamps; the string path counts the same. + assert_eq!(summary.clamped, 1); + } + + #[test] + fn text_comes_from_every_string_layout_and_mixed_from_typed_columns() { + let s = [ + series(1, "a", DatapointValueType::Text), + series(2, "b", DatapointValueType::Text), + series(3, "c", DatapointValueType::Mixed), + series(4, "d", DatapointValueType::Mixed), + ]; + let batch = table( + ms(&[1, 2]), + vec![ + ("a", Arc::new(StringViewArray::from(vec!["open", "blåbær"]))), + ("b", Arc::new(LargeStringArray::from(vec![Some("x"), None]))), + ("c", Arc::new(Float64Array::from(vec![1.5, 2.5]))), + ("d", Arc::new(StringArray::from(vec!["3.5", "fault"]))), + ], + ); + let got = through_table(&batch, &s); + assert_eq!(got[0].values, BinaryValues::Text(vec!["open".into(), "blåbær".into()])); + assert_eq!(got[1].values, BinaryValues::Text(vec!["x".into()])); + assert_eq!( + got[2].values, + BinaryValues::Mixed { numeric: vec![Some(1.5), Some(2.5)], text: vec![None, None] } + ); + // A string column of a mixed series is text, "3.5" included: the column is typed. + assert_eq!( + got[3].values, + BinaryValues::Mixed { numeric: vec![None, None], text: vec![Some("3.5".into()), Some("fault".into())] } + ); + } + + #[test] + fn long_series_are_cut_at_the_row_cap_and_short_ones_share_a_frame() { + let s = [ + series(9, "long", DatapointValueType::Float), + series(2, "short_b", DatapointValueType::Float), + series(1, "short_a", DatapointValueType::Float), + ]; + let n = 250_000; + let ts: Vec = (0..n as i64).collect(); + let long = Arc::new(Float64Array::from_iter_values((0..n).map(|i| i as f64))); + let sparse = |every: usize| -> ArrayRef { + Arc::new(Float64Array::from_iter((0..n).map(|i| (i % every == 0).then_some(1.0)))) + }; + let batch = table(ms(&ts), vec![("long", long), ("short_b", sparse(100_000)), ("short_a", sparse(50_000))]); + let (frames, summary) = table_frames(&batch, "timestamp", &s, FrameCompression::Payload, 1).unwrap(); + + assert_eq!(summary.points, n + 3 + 5); + // Ascending ids: the two short series (5 + 3 rows) open the first frame, and the long one + // fills the rest of it and runs on over two more. + let rows: Vec = frames.iter().map(|f| f.rows).collect(); + assert_eq!(rows, [100_000, 100_000, 50_008]); + assert_eq!(frames[0].series, 3, "short_a, short_b and the start of long"); + let read = decode(&frames); + assert_eq!(read.iter().map(|s| (s.id, s.len())).collect::>(), [(1, 5), (2, 3), (9, n)]); + } + + #[test] + fn what_the_table_cannot_say_is_refused() { + let float = [series(1, "v", DatapointValueType::Float)]; + let naive = Arc::new(TimestampMillisecondArray::from(vec![1])); + assert!(table_frames(&table(naive, vec![("v", Arc::new(Float64Array::from(vec![1.0])))]), "timestamp", &float, FrameCompression::Payload, 1) + .unwrap_err() + .contains("no time zone")); + + let sub_ms = Arc::new(TimestampMicrosecondArray::from(vec![1_500]).with_timezone("UTC")); + assert!(table_frames(&table(sub_ms, vec![("v", Arc::new(Float64Array::from(vec![1.0])))]), "timestamp", &float, FrameCompression::Payload, 1) + .unwrap_err() + .contains("finer than a millisecond")); + + let bigint = [series(1, "v", DatapointValueType::Bigint)]; + assert!(table_frames(&table(ms(&[1]), vec![("v", Arc::new(Float64Array::from(vec![1.0])))]), "timestamp", &bigint, FrameCompression::Payload, 1) + .unwrap_err() + .contains("cannot hold")); + + let text = [series(1, "v", DatapointValueType::Text)]; + assert!(table_frames(&table(ms(&[1]), vec![("v", Arc::new(StringArray::from(vec!["x".repeat(65)])))]), "timestamp", &text, FrameCompression::Payload, 1) + .unwrap_err() + .contains("characters")); + } +} diff --git a/src/timeseries/fixtures/GenerateReadFixtures.java b/src/timeseries/fixtures/GenerateReadFixtures.java new file mode 100644 index 0000000..4c31343 --- /dev/null +++ b/src/timeseries/fixtures/GenerateReadFixtures.java @@ -0,0 +1,165 @@ +// Regenerates read_values.bin, read_aggregates.bin and read_batches.bin with the platform's own +// DatapointFrameWriter.relayStream, the code that answers POST /timeseries/data/list/binary: Arrow +// streams compressed in-band, as ClickHouse answers, relayed into compression 2 frames. Run from a +// platform checkout that has built datahub-api-model, with arrow-format, flatbuffers-java, +// zero-allocation-hashing, zstd-jni and slf4j-api on the classpath: +// +// java -cp ":" \ +// GenerateReadFixtures.java src/timeseries/fixtures +// +// The Rust tests in binary_read.rs assert on the exact values written here. read_*_zstd.bin are +// the same series as compression 1 frames, written by the writer's former `wrap`, which is gone. +import ai.intellistream.datahub.api.binary.*; +import ai.intellistream.datahub.api.binary.ArrowIpc.*; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; + +/** Writes binary-read response bodies with the platform's own writer, for the Rust reader to parse. */ +public class GenerateReadFixtures { + static final long T0 = 1_700_000_000_000L; + static final PayloadCodec CODEC = new ZstdPayloadCodec(3); + + /** An ingest frame's batches as ClickHouse would answer them, relayed the way the API answers a read. */ + static byte[] response(DatapointValueType type, long id, String ext, DatapointFrameWriter w) throws Exception { + byte[] ingest = w.build(CODEC); + int dirLen = ByteBuffer.wrap(ingest).order(ByteOrder.LITTLE_ENDIAN).getInt(16); + int payLen = ByteBuffer.wrap(ingest).order(ByteOrder.LITTLE_ENDIAN).getInt(20); + int rawLen = ByteBuffer.wrap(ingest).order(ByteOrder.LITTLE_ENDIAN).getInt(24); + byte[] payload = java.util.Arrays.copyOfRange(ingest, 28 + dirLen, 28 + dirLen + payLen); + byte[] raw = CODEC.decompress(payload, rawLen); + List data = new ArrayList<>(); + for (Batch b : ArrowIpc.readStream(ArrowSchemaCanon.of(type), raw, 0, 100_000)) { + data.add(b.toData()); + } + return relay(type, false, id, ext, ArrowIpc.writeStream(ArrowSchemaCanon.of(type), data, CODEC)); + } + + static byte[] relay(DatapointValueType type, boolean aggregates, long id, String ext, byte[] stream) throws Exception { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + DatapointFrameWriter.relayStream(type, aggregates, id, ext, new ByteArrayInputStream(stream), CODEC, out); + return out.toByteArray(); + } + + /** {@code rows} rows of series {@code id} from {@code from}, a second apart, values from {@code values}. */ + static BatchData batch(long id, int from, int rows, java.util.function.IntFunction values) { + long[] ids = new long[rows]; + long[] stamps = new long[rows]; + for (int r = 0; r < rows; r++) { + ids[r] = id; + stamps[r] = T0 + (from + r) * 1000L; + } + return new BatchData(rows, new ColumnData[]{ + new ColumnData(null, 0, null, longs(ids)), + new ColumnData(null, 0, null, longs(stamps)), + values.apply(from), + }); + } + + static ByteBuffer longs(long... v) { + ByteBuffer b = ByteBuffer.allocate(v.length * 8).order(ByteOrder.LITTLE_ENDIAN); + for (long x : v) b.putLong(x); + return b.flip(); + } + + static ByteBuffer doubles(double... v) { + ByteBuffer b = ByteBuffer.allocate(v.length * 8).order(ByteOrder.LITTLE_ENDIAN); + for (double x : v) b.putDouble(x); + return b.flip(); + } + + public static void main(String[] args) throws Exception { + Path out = Path.of(args[0]); + ByteArrayOutputStream values = new ByteArrayOutputStream(); + + var bigint = DatapointFrameWriter.forType(DatapointValueType.BIGINT).series(11, "rd_bigint"); + bigint.addBigint(11, T0 + 2000, -9_000_000_000L); + bigint.addBigint(11, T0, 1); + bigint.addBigint(11, T0 + 1000, 42); + values.writeBytes(response(DatapointValueType.BIGINT, 11, "rd_bigint", bigint)); + + var flt = DatapointFrameWriter.forType(DatapointValueType.FLOAT).series(12, "rd_float"); + flt.addFloat(12, T0, 1.5); + flt.addFloat(12, T0 + 1000, -2.25); + values.writeBytes(response(DatapointValueType.FLOAT, 12, "rd_float", flt)); + + var f32 = DatapointFrameWriter.forType(DatapointValueType.FLOAT32).series(13, "rd_float32"); + f32.addFloat32(13, T0, 0.5f); + values.writeBytes(response(DatapointValueType.FLOAT32, 13, "rd_float32", f32)); + + var numeric = DatapointFrameWriter.forType(DatapointValueType.NUMERIC).series(14, "rd_numeric"); + numeric.addNumeric(14, T0, new BigDecimal("123456789012.345678")); + numeric.addNumeric(14, T0 + 1000, new BigDecimal("-0.000001")); + values.writeBytes(response(DatapointValueType.NUMERIC, 14, "rd_numeric", numeric)); + + var dec32 = DatapointFrameWriter.forType(DatapointValueType.DECIMAL32).series(15, "rd_decimal32"); + dec32.addDecimal32(15, T0, new BigDecimal("12345.6789")); + values.writeBytes(response(DatapointValueType.DECIMAL32, 15, "rd_decimal32", dec32)); + + var text = DatapointFrameWriter.forType(DatapointValueType.TEXT).series(16, "rd_text"); + text.addText(16, T0, "open"); + text.addText(16, T0 + 1000, "blåbær"); + values.writeBytes(response(DatapointValueType.TEXT, 16, "rd_text", text)); + + var mixed = DatapointFrameWriter.forType(DatapointValueType.MIXED).series(17, "rd_mixed"); + mixed.addMixed(17, T0, "3.5"); + mixed.addMixed(17, T0 + 1000, "fault"); + mixed.addMixed(17, T0 + 2000, "-7"); + values.writeBytes(response(DatapointValueType.MIXED, 17, "rd_mixed", mixed)); + + // Only numbers: the API leaves the validity bitmap out of a column with no nulls. + var mixedNumbers = DatapointFrameWriter.forType(DatapointValueType.MIXED).series(18, "rd_mixed_numbers"); + mixedNumbers.addMixed(18, T0, "1"); + mixedNumbers.addMixed(18, T0 + 1000, "2"); + values.writeBytes(response(DatapointValueType.MIXED, 18, "rd_mixed_numbers", mixedNumbers)); + + values.writeBytes(FrameLimits.END_OF_RESPONSE); + Files.write(out.resolve("read_values.bin"), values.toByteArray()); + + ByteArrayOutputStream aggregates = new ByteArrayOutputStream(); + BatchData buckets = new BatchData(3, new ColumnData[]{ + new ColumnData(null, 0, null, longs(21, 21, 21)), + new ColumnData(null, 0, null, longs(T0, T0 + 3_600_000, T0 + 7_200_000)), + new ColumnData(null, 0, null, doubles(1.5, 2.5, 3.5)), + new ColumnData(null, 0, null, doubles(1, 2, 3)), + new ColumnData(null, 0, null, doubles(2, 3, 4)), + new ColumnData(null, 0, null, doubles(6, 10, 14)), + }); + aggregates.writeBytes(relay(DatapointValueType.FLOAT, true, 21, "rd_aggregates", + ArrowIpc.writeStream(ArrowSchemaCanon.aggregates(), List.of(buckets), CODEC))); + aggregates.writeBytes(FrameLimits.END_OF_RESPONSE); + Files.write(out.resolve("read_aggregates.bin"), aggregates.toByteArray()); + + // Batches big enough that zstd shrinks their buffers, several to a frame, and one stream + // that came uncompressed and longer than a frame, which the relay cuts and compresses. + ByteArrayOutputStream batches = new ByteArrayOutputStream(); + int n = 3000; + List floats = List.of( + batch(31, 0, n, from -> new ColumnData(null, 0, null, doubles(java.util.stream.IntStream.range(from, from + n).mapToDouble(i -> i * 0.5).toArray()))), + batch(31, n, n, from -> new ColumnData(null, 0, null, doubles(java.util.stream.IntStream.range(from, from + n).mapToDouble(i -> i * 0.5).toArray())))); + batches.writeBytes(relay(DatapointValueType.FLOAT, false, 31, "rd_float_batches", + ArrowIpc.writeStream(ArrowSchemaCanon.of(DatapointValueType.FLOAT), floats, CODEC))); + int t = 12_000; + BatchData texts = batch(32, 0, t, from -> { + ByteBuffer offsets = ByteBuffer.allocate((t + 1) * 4).order(ByteOrder.LITTLE_ENDIAN); + ByteArrayOutputStream chars = new ByteArrayOutputStream(); + for (int i = 0; i < t; i++) { + offsets.putInt(chars.size()); + chars.writeBytes(("v" + (i % 10)).getBytes(java.nio.charset.StandardCharsets.UTF_8)); + } + offsets.putInt(chars.size()); + return new ColumnData(null, 0, offsets.flip(), ByteBuffer.wrap(chars.toByteArray())); + }); + batches.writeBytes(relay(DatapointValueType.TEXT, false, 32, "rd_text_cut", + ArrowIpc.writeStream(ArrowSchemaCanon.of(DatapointValueType.TEXT), List.of(texts)))); + batches.writeBytes(FrameLimits.END_OF_RESPONSE); + Files.write(out.resolve("read_batches.bin"), batches.toByteArray()); + } +} diff --git a/src/timeseries/fixtures/read_aggregates.bin b/src/timeseries/fixtures/read_aggregates.bin new file mode 100644 index 0000000..0f2d458 Binary files /dev/null and b/src/timeseries/fixtures/read_aggregates.bin differ diff --git a/src/timeseries/fixtures/read_aggregates_zstd.bin b/src/timeseries/fixtures/read_aggregates_zstd.bin new file mode 100644 index 0000000..adab85f Binary files /dev/null and b/src/timeseries/fixtures/read_aggregates_zstd.bin differ diff --git a/src/timeseries/fixtures/read_batches.bin b/src/timeseries/fixtures/read_batches.bin new file mode 100644 index 0000000..fc5ad7f Binary files /dev/null and b/src/timeseries/fixtures/read_batches.bin differ diff --git a/src/timeseries/fixtures/read_values.bin b/src/timeseries/fixtures/read_values.bin new file mode 100644 index 0000000..3a4d999 Binary files /dev/null and b/src/timeseries/fixtures/read_values.bin differ diff --git a/src/timeseries/fixtures/read_values_zstd.bin b/src/timeseries/fixtures/read_values_zstd.bin new file mode 100644 index 0000000..a61144c Binary files /dev/null and b/src/timeseries/fixtures/read_values_zstd.bin differ diff --git a/src/timeseries/mod.rs b/src/timeseries/mod.rs index e34231f..73c905b 100644 --- a/src/timeseries/mod.rs +++ b/src/timeseries/mod.rs @@ -1,8 +1,14 @@ pub mod binary; +pub mod binary_read; +pub mod binary_table; pub mod datapoint_listen; mod test; -pub use binary::{BinaryIngestOptions, DatapointValueType, Frame, FrameWriter, ResolvedSeries}; +pub use binary::{ + BinaryIngestOptions, DatapointValueType, Frame, FrameCompression, FrameWriter, ResolvedSeries, +}; +pub use binary_read::{BinaryDatapointStream, BinaryDatapoints, BinaryValues}; +pub use binary_table::TableIngest; pub use datapoint_listen::{DatapointListener, LiveDatapoint}; use crate::buffer::DurableSpool; diff --git a/src/timeseries/test.rs b/src/timeseries/test.rs index 2f7bb38..8550af7 100644 --- a/src/timeseries/test.rs +++ b/src/timeseries/test.rs @@ -12,7 +12,8 @@ mod tests { use crate::generic::{DataWrapper, Datapoint, DatapointString, DatapointsCollection, DeleteFilter, IdAndExtId, RetrieveFilter}; use crate::http::ResponseError; use crate::timeseries::{TimeSeries, TimeSeriesFilter, TimeSeriesFilterForm, TimeSeriesUpdate, TimeSeriesUpdateCollection, TimeSeriesUpdateFields}; - use crate::timeseries::binary::BinaryIngestOptions; + use crate::timeseries::binary::{BinaryIngestOptions, DatapointValueType, FrameCompression}; + use crate::timeseries::binary_read::{BinaryDatapoints, BinaryValues}; use crate::tests::cleanup::cleanup_timeseries; use crate::tests::ids::{unique_id, unique_token}; use crate::tests::polling::poll_until_for; @@ -712,12 +713,185 @@ mod tests { Ok(()) } + /// `POST /timeseries/data/list/binary` against the JSON read of the same series: every point + /// of a float series past one frame's row cap, hourly aggregates, a text series, an empty + /// window, and the two refusals the JSON read does not have. Needs a backend that serves the + /// binary read. + #[tokio::test] + #[ignore] + async fn test_datapoints_binary_read() -> Result<(), Box> { + let api_service = create_api_service(); + let float_ext_id = unique_id("bin_read_float"); + let text_ext_id = unique_id("bin_read_text"); + let mut ts_collection = DataWrapper::new(); + for (external_id, value_type) in [(&float_ext_id, "float"), (&text_ext_id, "text")] { + ts_collection.add_item( + TimeSeries::builder() + .set_external_id(external_id) + .set_name(external_id) + .set_unit("celsius") + .set_value_type(value_type) + .clone(), + ); + } + let mut cleanup = cleanup_timeseries(vec![float_ext_id.clone(), text_ext_id.clone()]); + api_service.time_series.create(&ts_collection).await.expect("could not create the series"); + + // 150 000 one-second floats: past the 100 000-row numeric frame cap and past the 100 000 + // an item's `limit` may ask for, so they come back whole only from an item without one, + // as two frames. 25 000 texts: past the 10 000-row text frame cap, so three frames. + let start = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap(); + let count: i64 = 150_000; + let value = |i: i64| (i % 3600) as f64 * 0.25 - 100.0; + let mut floats = DatapointsCollection::from_external_id(&float_ext_id); + floats.datapoints = (0..count) + .map(|i| DatapointString::new(&(start.timestamp_millis() + i * 1000).to_string(), &value(i).to_string())) + .collect(); + let mut texts = DatapointsCollection::from_external_id(&text_ext_id); + let text_count: i64 = 25_000; + let text = |i: i64| format!("blåbær_{}", i % 97); + texts.datapoints = (0..text_count) + .map(|i| DatapointString::new(&(start.timestamp_millis() + i * 1000).to_string(), &text(i))) + .collect(); + api_service.time_series + .insert_datapoints_binary(&DataWrapper::from_vec(vec![floats, texts]), &BinaryIngestOptions::default()) + .await + .unwrap_or_else(|e| panic!("binary insert failed with {}: {}", e.get_status(), e.get_message())); + let item = |external_id: &str, limit: u64| { + let mut rf = RetrieveFilter::new(); + rf.set_external_id(external_id); + rf.set_start(start); + rf.set_limit(limit); + rf + }; + poll_all_points(&api_service, &float_ext_id, count as usize).await; + // Not poll_all_points: the JSON `Datapoint` holds only numbers and cannot read text. + let text_request = DataWrapper::from_vec(vec![item(&text_ext_id, text_count as u64)]); + poll_until_for( + std::time::Duration::from_secs(60), + || api_service.time_series.retrieve_datapoints_binary(&text_request), + |read| matches!(read, Ok(series) if series.first().is_some_and(|s| s.len() == text_count as usize)), + ) + .await; + + println!("Every point, in order, across frames..."); + let mut whole_window = RetrieveFilter::new(); + whole_window.set_external_id(&float_ext_id); + whole_window.set_start(start); + whole_window.set_end(start + Duration::seconds(count)); + let request = DataWrapper::from_vec(vec![whole_window, item(&text_ext_id, text_count as u64)]); + let first = api_service.time_series + .retrieve_datapoints_binary(&request) + .await + .unwrap_or_else(|e| panic!("binary read failed with {}: {}", e.get_status(), e.get_message())); + assert_eq!(first.len(), 2, "one entry per series, the frames of one merged"); + let float = first[0].clone(); + assert_eq!(float.external_id, float_ext_id); + assert_eq!(float.value_type, DatapointValueType::Float); + let want_timestamps: Vec = (0..count).map(|i| start.timestamp_millis() + i * 1000).collect(); + assert_eq!(float.timestamps, want_timestamps); + let want_values: Vec = (0..count).map(value).collect(); + assert_eq!(float.values, BinaryValues::Float(want_values)); + + println!("The same read, streamed frame by frame..."); + let mut stream = api_service.time_series.stream_datapoints_binary(&request).await?; + let (mut frames, mut streamed) = (Vec::new(), Vec::new()); + while let Some(chunk) = futures::StreamExt::next(&mut stream).await { + let chunk = chunk?; + if chunk.id == float.id { + streamed.extend_from_slice(&chunk.timestamps); + } + frames.push((chunk.len(), chunk.external_id)); + } + assert_eq!(frames.len(), 5, "two float frames and three text frames, one item each: {frames:?}"); + assert_eq!(streamed, want_timestamps); + + let texts = &first[1]; + assert_eq!(texts.external_id, text_ext_id); + assert_eq!(texts.timestamps, want_timestamps[..text_count as usize]); + assert_eq!(texts.values, BinaryValues::Text((0..text_count).map(text).collect())); + + println!("The same window as the JSON read..."); + let json = api_service.time_series + .retrieve_datapoints(&DataWrapper::from_vec(vec![item(&float_ext_id, 1000)])) + .await?; + let json_points = &json.get_items()[0].datapoints; + assert_eq!(json_points.len(), 1000); + for (i, dp) in json_points.iter().enumerate() { + assert_eq!(dp.timestamp.timestamp_millis(), float.timestamps[i]); + assert_eq!(dp.value, Some(want_value_at(&float, i))); + } + + println!("Hourly aggregates against the JSON aggregates..."); + let hourly = |mut rf: RetrieveFilter| { + rf.set_aggregates(vec!["avg".to_string(), "min".to_string(), "max".to_string(), "sum".to_string()]); + rf.set_granularity("1h"); + rf.set_end(start + Duration::seconds(count)); + DataWrapper::from_vec(vec![rf]) + }; + let mut rf = RetrieveFilter::new(); + rf.set_external_id(&float_ext_id); + rf.set_start(start); + let binary_aggs = api_service.time_series.retrieve_datapoints_binary(&hourly(rf.clone())).await?; + let json_aggs = api_service.time_series.retrieve_datapoints(&hourly(rf)).await?; + let json_buckets = &json_aggs.get_items()[0].datapoints; + let BinaryValues::Aggregates { avg, min, max, sum } = &binary_aggs[0].values else { + panic!("an aggregate read came back as {:?}", binary_aggs[0].values); + }; + assert_eq!(binary_aggs[0].len(), json_buckets.len(), "the two reads answered a different number of buckets"); + assert_eq!(json_buckets.len(), (count as usize).div_ceil(3600)); + for (i, bucket) in json_buckets.iter().enumerate() { + assert_eq!(binary_aggs[0].timestamps[i], bucket.timestamp.timestamp_millis()); + assert_eq!((Some(min[i]), Some(max[i])), (bucket.min, bucket.max), "bucket {i}"); + assert!((avg[i] - bucket.average.unwrap()).abs() < AVG_TOLERANCE, "bucket {i}: avg"); + assert!((sum[i] - bucket.sum.unwrap()).abs() < 1e-6 * sum[i].abs().max(1.0), "bucket {i}: sum"); + } + + println!("An empty window is no entry..."); + let mut empty = item(&float_ext_id, 10); + empty.set_start(start - Duration::days(10)); + empty.set_end(start - Duration::days(9)); + let none = api_service.time_series.retrieve_datapoints_binary(&DataWrapper::from_vec(vec![empty])).await?; + assert!(none.is_empty(), "{none:?}"); + + println!("Refusals..."); + let mut with_cursor = item(&float_ext_id, 10); + with_cursor.cursor = Some("anything".to_string()); + let error = api_service.time_series + .retrieve_datapoints_binary(&DataWrapper::from_vec(vec![with_cursor])) + .await + .unwrap_err(); + assert_eq!(error.get_status(), StatusCode::BAD_REQUEST, "{}", error.get_message()); + assert!(error.problem().is_some(), "a refusal is a problem document: {}", error.get_message()); + + let unknown = unique_id("bin_read_missing"); + let error = api_service.time_series + .retrieve_datapoints_binary(&DataWrapper::from_vec(vec![item(&unknown, 10)])) + .await + .unwrap_err(); + assert_eq!(error.get_status(), StatusCode::BAD_REQUEST, "{}", error.get_message()); + assert!(error.get_message().contains(&unknown), "{}", error.get_message()); + + delete_timeseries(&api_service, &[&float_ext_id, &text_ext_id]).await; + cleanup.disarm(); + Ok(()) + } + + fn want_value_at(series: &BinaryDatapoints, i: usize) -> f64 { + match &series.values { + BinaryValues::Float(values) => values[i], + other => panic!("not a float series: {other:?}"), + } + } + /// JSON against binary ingest from Rust, reporting the same columns as the Java benchmark in /// the platform's `datahub-e2e` module and the Python one in `python_tests`, so the three /// clients can be compared. /// /// Sized by `DATAHUB_BENCH_POINTS` (default 10 million) across `DATAHUB_BENCH_SERIES` - /// series. The api must have its daily quota and rate limiter off for a run of any size: + /// series. `DATAHUB_BENCH_ZSTD_LEVEL` and `DATAHUB_BENCH_COMPRESSION` (`payload` or `buffers`) + /// override the binary path's defaults. The api must have its daily quota and rate limiter off + /// for a run of any size: /// `-Ddatahub.limits.quota.enabled=false -Ddatahub.limits.rate.enabled=false`. /// /// **Run it with `--release`.** `cargo test` builds with the `dev` profile at @@ -740,8 +914,16 @@ mod tests { .ok().and_then(|v| v.parse().ok()).unwrap_or(100); let chunk: usize = std::env::var("DATAHUB_BENCH_CHUNK") .ok().and_then(|v| v.parse().ok()).unwrap_or(1_000_000); + let zstd_level: i32 = std::env::var("DATAHUB_BENCH_ZSTD_LEVEL") + .ok().and_then(|v| v.parse().ok()).unwrap_or(BinaryIngestOptions::default().zstd_level); + let compression = match std::env::var("DATAHUB_BENCH_COMPRESSION").as_deref() { + Ok("buffers") => FrameCompression::Buffers, + Ok("payload") => FrameCompression::Payload, + Err(_) => BinaryIngestOptions::default().compression, + Ok(other) => panic!("DATAHUB_BENCH_COMPRESSION is payload or buffers, not {other}"), + }; let api_service = create_api_service(); - println!("\n=== {total} points across {series_count} series, float32 ==="); + println!("\n=== {total} points across {series_count} series, float32, binary zstd {zstd_level} {compression:?} ==="); // Which paths to run. A sweep over request sizes only needs the binary one, and the JSON // path is slow enough that running it every time makes the sweep the expensive part. @@ -795,7 +977,10 @@ mod tests { api_service.time_series .insert_datapoints_binary( &request, - &BinaryIngestOptions::default().request_concurrency(concurrency), + &BinaryIngestOptions::default() + .zstd_level(zstd_level) + .compression(compression) + .request_concurrency(concurrency), ) .await .unwrap_or_else(|e| panic!("{label} insert failed: {}", e.get_message())); @@ -834,7 +1019,7 @@ mod tests { } println!("\n=== datapoint ingest from Rust ==="); - println!("{total} points across {series_count} series, {chunk} per call, float32\n"); + println!("{total} points across {series_count} series, {chunk} per call, float32, binary zstd {zstd_level} {compression:?}\n"); print!("{:<28}", "metric"); for r in &results { print!("{:>18}", r.0); } println!();