diff --git a/asap-precompute-rs/Cargo.toml b/asap-precompute-rs/Cargo.toml index 83e5881..f93de8a 100644 --- a/asap-precompute-rs/Cargo.toml +++ b/asap-precompute-rs/Cargo.toml @@ -39,3 +39,7 @@ default = [] # (decode_batch / encode_batch over `arrow::RecordBatch`, plus the # `AsapSketchesPlugin` Tokio runtime). otap = ["dep:arrow-array", "dep:arrow-schema", "dep:tokio", "dep:futures"] + +[[example]] +name = "sketch_pipeline_demo" +required-features = ["otap"] diff --git a/asap-precompute-rs/examples/sketch_pipeline_demo.rs b/asap-precompute-rs/examples/sketch_pipeline_demo.rs new file mode 100644 index 0000000..094372c --- /dev/null +++ b/asap-precompute-rs/examples/sketch_pipeline_demo.rs @@ -0,0 +1,206 @@ +//! Runnable end-to-end demo of the `docs/data_model.md` SCHEMA / +//! DICTIONARY / RECORD wire shape, three "processors" wired together +//! in one binary: +//! +//! 1. **Sketch creation processor** (`run_producer`) — feeds synthetic +//! latency samples into a [`PrecomputeImpl`], closes a window at a +//! time, and encodes the closed-window envelopes against a +//! [`SeriesDictionary`] — the SCHEMA/DICTIONARY/RECORD codec from +//! `otap::dictionary`. +//! 2. **Receive processor** (`run_receiver`) — decodes each +//! [`SketchStreamBatch`] via a [`SeriesDictionaryDecoder`], merges +//! the reconstructed envelopes into its own [`PrecomputeImpl`] +//! (`Precompute::observe_envelope` — bytes merged as sketch state, +//! never expanded to samples), then *queries* that merged sketch by +//! draining it with `transmit_sketch = false`: the runtime's +//! estimate-mode path (`Sketch::estimate`) turns the merged DDSketch +//! into a p99 gauge. +//! 3. **Prometheus backend** — stands in as a `println!` of the +//! gauge in Prometheus text exposition format; swap +//! [`format_prometheus_gauge`]'s call site for a real HTTP +//! `/metrics` handler to serve it for real. +//! +//! Producer and receiver run as two Tokio tasks connected by an +//! in-process `mpsc` channel carrying [`SketchStreamBatch`] — no +//! network transport, no serialization — so what crosses the "wire" +//! here is literally the same four `RecordBatch`es a real inter-node +//! hop would carry. Watch the printed row counts: window 0 carries +//! `SCHEMA`+`DICTIONARY`+`LABELS` rows (the series is new), every +//! later window carries only `RECORD` rows — that's the whole point of +//! this doc's design, made visible. +//! +//! Run with: +//! ```text +//! cargo run --example sketch_pipeline_demo --features otap +//! ``` + +use std::time::Duration; + +use asap_precompute_rs::envelope::SketchEnvelope; +use asap_precompute_rs::observation::{KeyValue, Observation, ObservationValue}; +use asap_precompute_rs::otap::config::{resolve, PluginConfig}; +use asap_precompute_rs::otap::{SeriesDictionary, SeriesDictionaryDecoder, SketchStreamBatch}; +use asap_precompute_rs::precompute::{Precompute, PrecomputeImpl}; +use tokio::sync::mpsc; + +/// Both processors run one aggregation plan, so they share this +/// join key — in a real deployment the receiver would learn it out of +/// band (control plane), same as `docs/data_model.md`'s "Open design +/// questions" section describes. +const AGG_ID: u64 = 1; +/// How many synthetic windows the producer closes before shutting +/// down. Each window's samples are drawn from a slowly climbing +/// distribution so the printed p99 visibly moves window to window. +const NUM_WINDOWS: usize = 4; +/// Purely for demo pacing (so the two tasks visibly interleave in the +/// terminal) — window *closing* itself is driven by explicit `drain()` +/// calls below, not wall-clock alignment, so this value doesn't affect +/// correctness. +const PACING: Duration = Duration::from_millis(150); + +#[tokio::main] +async fn main() { + let (tx, rx) = mpsc::unbounded_channel::(); + + let producer = tokio::spawn(run_producer(tx)); + let receiver = tokio::spawn(run_receiver(rx)); + + let _ = tokio::join!(producer, receiver); +} + +/// Sketch creation processor: observes synthetic samples, closes a +/// window at a time, encodes against a [`SeriesDictionary`], sends the +/// result downstream. +async fn run_producer(tx: mpsc::UnboundedSender) { + let plugin_cfg = PluginConfig { + sketch_type: "ddsketch".into(), + window_size: Duration::from_secs(10), // never naturally fires — see drain() below. + output_metric_name: "http_request_duration_ms".into(), + agg_id: AGG_ID, + sketch_params: [("relative_accuracy".to_string(), 0.01)] + .into_iter() + .collect(), + ..Default::default() + }; + let (pcfg, dispatch) = resolve(&plugin_cfg).expect("producer config"); + let precompute = PrecomputeImpl::new( + Some(pcfg.clone()), + Some(dispatch.factory), + Some(dispatch.observer), + ); + let mut dictionary = SeriesDictionary::new(); + + for window_idx in 0..NUM_WINDOWS { + // One series (path=/api), latency drifting upward window to + // window so the receiver's printed p99 visibly changes. + for i in 0..200u64 { + let base = 10.0 + (window_idx as f64) * 8.0; + let latency = base + (i % 25) as f64; + let obs = Observation::new( + now_ms(), + "http_request_duration_ms", + Vec::new(), + vec![KeyValue::new("path", "/api")], + ObservationValue::float(latency), + ); + precompute.observe(&obs).expect("observe"); + } + + // Force this window closed regardless of wall clock — see the + // module doc's "PACING" note. + let envelopes = precompute.drain(); + if envelopes.is_empty() { + continue; + } + let batch = dictionary + .encode(&envelopes, Some(&pcfg)) + .expect("encode window"); + println!( + "[producer] window {window_idx}: schema={} dictionary={} labels={} record={} row(s)", + batch.schema.num_rows(), + batch.dictionary.num_rows(), + batch.labels.num_rows(), + batch.record.num_rows(), + ); + if tx.send(batch).is_err() { + break; // receiver gone. + } + tokio::time::sleep(PACING).await; + } + // Dropping `tx` here signals the receiver's `rx.recv()` to return + // `None` once the channel drains. +} + +/// Receive processor: decodes each batch, merges the reconstructed +/// envelopes into its own window, queries the merge by draining in +/// estimate mode, and hands the resulting gauge to +/// [`format_prometheus_gauge`]. +async fn run_receiver(mut rx: mpsc::UnboundedReceiver) { + let plugin_cfg = PluginConfig { + sketch_type: "ddsketch".into(), + window_size: Duration::from_secs(10), // same "driven by drain(), not wall clock" note as the producer. + output_metric_name: "http_request_duration_ms_p99".into(), + agg_id: AGG_ID, + transmit_sketch: false, // query mode: drain() yields quantile estimates, not sketch bytes. + quantiles: vec![0.99], + ..Default::default() + }; + let (pcfg, dispatch) = resolve(&plugin_cfg).expect("receiver config"); + let precompute = + PrecomputeImpl::new(Some(pcfg), Some(dispatch.factory), Some(dispatch.observer)); + let mut decoder = SeriesDictionaryDecoder::new(); + + while let Some(batch) = rx.recv().await { + let envelopes = decoder.decode(&batch).expect("decode stream batch"); + println!( + "[receiver] decoded {} envelope(s) from the batch", + envelopes.len() + ); + for env in &envelopes { + // Merge only — the runtime never expands envelope bytes + // back into scalar samples (the bandwidth invariant). + precompute.observe_envelope(env).expect("observe_envelope"); + } + + // Query: force this window's merged sketch to close now, in + // estimate mode, so the runtime's own estimate machinery + // (Sketch::estimate) does the quantile math for us. + for estimate in precompute.drain() { + print!("{}", format_prometheus_gauge(&estimate)); + } + } +} + +/// Formats one estimate-mode [`SketchEnvelope`] (`payload` empty, +/// `value` set — see `docs/data_model.md`'s `RECORD.value`) as a +/// Prometheus text-exposition gauge sample. Stands in for a real +/// `/metrics` HTTP handler. +fn format_prometheus_gauge(env: &SketchEnvelope) -> String { + let mut labels: Vec = env + .labels + .iter() + .map(|kv| format!("{}=\"{}\"", kv.key, kv.value)) + .collect(); + labels.sort(); + let label_str = if labels.is_empty() { + String::new() + } else { + format!("{{{}}}", labels.join(",")) + }; + format!( + "# HELP {name} sketch-derived quantile estimate\n\ + # TYPE {name} gauge\n\ + {name}{label_str} {value} {ts}\n", + name = env.metric_name, + value = env.value, + ts = env.window_end_ms, + ) +} + +fn now_ms() -> u64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before Unix epoch") + .as_millis() as u64 +} diff --git a/asap-precompute-rs/src/config.rs b/asap-precompute-rs/src/config.rs index a9f06f2..501cf85 100644 --- a/asap-precompute-rs/src/config.rs +++ b/asap-precompute-rs/src/config.rs @@ -140,6 +140,35 @@ pub fn sketch_param_get(params: &SketchParams, key: &str, default: f64) -> f64 { params.get(key).copied().unwrap_or(default) } +/// Renders `sketch_type`'s primary size/accuracy parameter(s) out of +/// `params` as a human-readable string — the `SCHEMA.sketch_size` +/// field in `docs/data_model.md`'s field reference, matching that +/// doc's own worked examples (`0.01`, `200`, `12`, `2048 x 4`) rather +/// than forcing every algorithm's parameter shape into one numeric +/// type (CountSketch / CountMinSketch need two numbers, not one). +/// +/// Returns `None` when the relevant key(s) for `sketch_type` aren't +/// present in `params` — matches the field's "optional" status. +pub fn sketch_size_string(sketch_type: SketchType, params: &SketchParams) -> Option { + match sketch_type { + SketchType::DDSketch => params.get("relative_accuracy").map(f64::to_string), + SketchType::KLLSketch => params.get("k").map(|v| (*v as u64).to_string()), + SketchType::HLLSketch => params.get("precision").map(|v| (*v as u64).to_string()), + SketchType::CountSketch => match (params.get("width"), params.get("depth")) { + (Some(w), Some(d)) => Some(format!("{} x {}", *w as u64, *d as u64)), + _ => match (params.get("epsilon"), params.get("delta")) { + (Some(e), Some(d)) => Some(format!("epsilon={e}, delta={d}")), + _ => None, + }, + }, + SketchType::CountMinSketch => match (params.get("rows"), params.get("columns")) { + (Some(r), Some(c)) => Some(format!("{} x {}", *r as u64, *c as u64)), + _ => None, + }, + SketchType::Unspecified => None, + } +} + /// Host-neutral form of today's per-OTel-processor `Config` struct, /// plus `max_series` / `on_overflow`. #[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] @@ -428,4 +457,45 @@ mod tests { assert_eq!(sketch_param_get(&p, "precision", 4.0), 12.0); assert_eq!(sketch_param_get(&p, "missing", 4.0), 4.0); } + + #[test] + fn sketch_size_string_matches_doc_worked_examples() { + let mut dd = SketchParams::new(); + dd.insert("relative_accuracy".into(), 0.01); + assert_eq!( + sketch_size_string(SketchType::DDSketch, &dd), + Some("0.01".to_string()) + ); + + let mut kll = SketchParams::new(); + kll.insert("k".into(), 200.0); + assert_eq!( + sketch_size_string(SketchType::KLLSketch, &kll), + Some("200".to_string()) + ); + + let mut hll = SketchParams::new(); + hll.insert("precision".into(), 12.0); + assert_eq!( + sketch_size_string(SketchType::HLLSketch, &hll), + Some("12".to_string()) + ); + + let mut cs = SketchParams::new(); + cs.insert("width".into(), 2048.0); + cs.insert("depth".into(), 4.0); + assert_eq!( + sketch_size_string(SketchType::CountSketch, &cs), + Some("2048 x 4".to_string()) + ); + + assert_eq!( + sketch_size_string(SketchType::Unspecified, &SketchParams::new()), + None + ); + assert_eq!( + sketch_size_string(SketchType::DDSketch, &SketchParams::new()), + None + ); + } } diff --git a/asap-precompute-rs/src/otap/decode.rs b/asap-precompute-rs/src/otap/decode.rs index cabbf00..a57ff05 100644 --- a/asap-precompute-rs/src/otap/decode.rs +++ b/asap-precompute-rs/src/otap/decode.rs @@ -70,6 +70,42 @@ pub enum OtapDecodeError { /// Raw string value observed. value: String, }, + + /// A `RECORD` row referenced a `series_id` with no matching + /// `DICTIONARY` entry in [`super::dictionary::SeriesDictionaryDecoder`]'s + /// retained state. + /// + /// Per `docs/data_model.md`'s "Where the Schema/Dictionary + /// statefulness guarantee actually comes from": `DICTIONARY`'s + /// incremental-append economics assume a continuous-stream + /// contract (stable routing, no replica hand-off mid-stream). This + /// error surfaces a violation of that contract rather than + /// silently decoding a `RECORD` with an empty series identity. + #[error("otap decode: record references unknown series_id {series_id}")] + UnknownSeriesId { + /// The unresolved `series_id`. + series_id: u32, + }, + + /// A `DICTIONARY` row referenced an `agg_id` with no matching + /// `SCHEMA` entry in the decoder's retained state. Same + /// continuous-stream caveat as [`Self::UnknownSeriesId`]. + #[error("otap decode: dictionary entry references unknown agg_id {agg_id}")] + UnknownAggId { + /// The unresolved `agg_id`. + agg_id: u64, + }, + + /// A required column was missing entirely from one of + /// [`super::dictionary::SketchStreamBatch`]'s four sub-batches. + #[error("otap decode: batch {batch:?} missing required column {column:?}")] + MissingColumn { + /// Which sub-batch (`"schema"` / `"dictionary"` / `"labels"` / + /// `"record"`) is missing the column. + batch: &'static str, + /// Column name. + column: &'static str, + }, } /// Decode an OTAP `RecordBatch` into a `Vec`. @@ -225,7 +261,10 @@ fn build_envelope( }) } -fn parse_sketch_type(row: usize, raw: &str) -> Result { +/// Parses a canonical sketch-type string. `pub(super)` so +/// [`super::dictionary`] can reuse the same parsing (and error +/// variant) for `SCHEMA.sketch_type` rows. +pub(super) fn parse_sketch_type(row: usize, raw: &str) -> Result { match raw { "DDSketch" => Ok(SketchType::DDSketch), "KLLSketch" => Ok(SketchType::KLLSketch), @@ -240,7 +279,10 @@ fn parse_sketch_type(row: usize, raw: &str) -> Result Result { +/// Parses a canonical encoding string. `pub(super)` so +/// [`super::dictionary`] can reuse the same parsing (and error +/// variant) for `SCHEMA.encoding` rows. +pub(super) fn parse_encoding(row: usize, raw: &str) -> Result { match raw { "PROTO_FULL" => Ok(Encoding::ProtoFull), "PROTO_DELTA" => Ok(Encoding::ProtoDelta), diff --git a/asap-precompute-rs/src/otap/dictionary.rs b/asap-precompute-rs/src/otap/dictionary.rs new file mode 100644 index 0000000..ecd06fc --- /dev/null +++ b/asap-precompute-rs/src/otap/dictionary.rs @@ -0,0 +1,1091 @@ +//! `SeriesDictionary` / `SeriesDictionaryDecoder` — the Schema / +//! Dictionary / Record tiering from +//! `docs/data_model.md#schema--dictionary--record-as-entities`, +//! implemented as ASAP's own inter-node sketch-stream wire shape. +//! +//! This is deliberately a *different* codec from [`super::encode_batch`] +//! / [`super::decode_batch`] / [`super::records`], not a replacement: +//! those exist to disguise a sketch envelope as an OTAP-Metrics-shaped +//! payload (one self-contained row per envelope, `_asap_*` attributes +//! lifted onto the per-row attribute child batch) so it can transit an +//! OTAP pipeline hop that only knows how to move Logs/Metrics/Traces +//! payloads. `docs/data_model.md` is about a narrower, different hop — +//! its very first line scopes it to "sketch state cross[ing] a node or +//! network boundary between `asap_sketches` processor instances" — an +//! ASAP-aware sender talking to an ASAP-aware receiver, where there's +//! no need to *look like* a generic OTLP metric. That's the hop this +//! module implements: `SCHEMA` is sent once per distinct `agg_id`, +//! `DICTIONARY` (+ `LABELS`) once per distinct series, and `RECORD` +//! carries only what's genuinely unique per window — `series_id`, +//! window bounds, and `envelope`/`value`. No `metric` name or label is +//! ever repeated on a `RECORD` row. +//! +//! # Statefulness +//! +//! Per the doc's "Where the Schema/Dictionary statefulness guarantee +//! actually comes from": this only saves anything if the same +//! [`SeriesDictionary`] keeps encoding every batch for a given output +//! stream (so "already sent" state is meaningful), and the same +//! [`SeriesDictionaryDecoder`] keeps decoding every batch from that +//! stream in order (so retained `SCHEMA`/`DICTIONARY` state is there +//! to join against). A `RECORD` referencing a `series_id`/`agg_id` +//! the decoder never saw a `DICTIONARY`/`SCHEMA` row for is a hard +//! decode error, not a silent partial result — see +//! [`super::decode::OtapDecodeError::UnknownSeriesId`] / +//! [`OtapDecodeError::UnknownAggId`]. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use arrow_array::{ + Array, BinaryArray, Float64Array, RecordBatch, StringArray, UInt32Array, UInt64Array, +}; +use arrow_schema::{DataType, Field, Schema}; + +use crate::config::{sketch_size_string, AggId, PrecomputeConfig}; +use crate::envelope::{Encoding, SketchEnvelope, SketchType}; +use crate::observation::KeyValue; + +use super::decode::{parse_encoding, parse_sketch_type, OtapDecodeError}; +use super::encode::OtapEncodeError; +use super::schema::{ + DICT_COLUMN_METRIC, DICT_COLUMN_SERIES_ID, LABELS_COLUMN_KEY, LABELS_COLUMN_VALUE, + RECORD_COLUMN_ENVELOPE, RECORD_COLUMN_VALUE, RECORD_COLUMN_WINDOW_END_MS, + RECORD_COLUMN_WINDOW_START_MS, SCHEMA_COLUMN_AGG_ID, SCHEMA_COLUMN_ENCODING, + SCHEMA_COLUMN_HASH_FUNCTION, SCHEMA_COLUMN_HASH_SEED, SCHEMA_COLUMN_SCHEMA_VERSION, + SCHEMA_COLUMN_SKETCH_SIZE, SCHEMA_COLUMN_SKETCH_TYPE, +}; + +/// The four-batch family [`SeriesDictionary::encode`] produces and +/// [`SeriesDictionaryDecoder::decode`] consumes, mirroring +/// `docs/data_model.md`'s ER diagram one-for-one. +#[derive(Debug, Clone)] +pub struct SketchStreamBatch { + /// One row per `agg_id` first seen by the encoding + /// [`SeriesDictionary`]. Empty (schema-only) once every live + /// `agg_id` has already been sent. + pub schema: RecordBatch, + /// One row per series first seen by the encoding + /// [`SeriesDictionary`]. Empty once every live series has already + /// been sent. + pub dictionary: RecordBatch, + /// One row per label key, for series first seen this call (a + /// child of `dictionary` — same `series_id`s, zero or more rows + /// apiece). + pub labels: RecordBatch, + /// One row per envelope passed to [`SeriesDictionary::encode`], + /// always — this is the only batch whose row count scales with + /// observations. + pub record: RecordBatch, +} + +impl SketchStreamBatch { + /// True when every batch in the family has zero rows. + pub fn is_empty(&self) -> bool { + self.schema.num_rows() == 0 + && self.dictionary.num_rows() == 0 + && self.labels.num_rows() == 0 + && self.record.num_rows() == 0 + } +} + +/// Sender-side dictionary state: assigns stable `series_id`s and +/// tracks which `agg_id`s / series have already had a `SCHEMA` / +/// `DICTIONARY` row emitted, so repeat windows for the same series +/// cost only a `RECORD` row. +/// +/// One instance per **output stream** (i.e. per downstream receiver +/// this node is emitting to) — its whole value comes from persisting +/// across calls to [`Self::encode`], the same way +/// [`crate::snapshot_cache::SnapshotCache`] persists across window +/// rotations for delta encoding. +pub struct SeriesDictionary { + next_series_id: u32, + series_ids: HashMap, + known_series: HashSet, + known_schemas: HashSet, +} + +impl Default for SeriesDictionary { + fn default() -> Self { + Self::new() + } +} + +impl SeriesDictionary { + /// Constructs an empty dictionary — nothing sent yet. + pub fn new() -> Self { + Self { + next_series_id: 0, + series_ids: HashMap::new(), + known_series: HashSet::new(), + known_schemas: HashSet::new(), + } + } + + /// Canonical series identity: `(agg_id, metric_name, labels)`, + /// labels sorted by key. Deliberately independent of + /// [`crate::matchers::series_key`] (which is byte-parity-locked + /// for the snapshot-cache's unrelated purpose) — this + /// canonicalization is private to this dictionary and never + /// crosses a process boundary itself, only the `series_id` it + /// produces does. + fn identity_key(env: &SketchEnvelope) -> String { + let mut labels: Vec<&KeyValue> = env.labels.iter().collect(); + labels.sort_by(|a, b| a.key.cmp(&b.key)); + let mut buf = String::new(); + buf.push_str(&env.agg_id.to_string()); + buf.push('|'); + buf.push_str(&env.metric_name); + buf.push('|'); + for kv in labels { + buf.push_str(&kv.key); + buf.push('='); + buf.push_str(&kv.value); + buf.push(';'); + } + buf + } + + /// Returns `env`'s `series_id`, assigning a fresh one the first + /// time this identity is seen. + fn series_id_for(&mut self, env: &SketchEnvelope) -> u32 { + let key = Self::identity_key(env); + if let Some(id) = self.series_ids.get(&key) { + return *id; + } + let id = self.next_series_id; + self.next_series_id += 1; + self.series_ids.insert(key, id); + id + } + + /// Encodes one `Precompute::tick`/`drain` call's worth of + /// envelopes against this dictionary's accumulated state. + /// + /// `cfg` sources `SCHEMA`-tier facts that live on the config + /// rather than on `SketchEnvelope` (`sketch_params`, rendered via + /// [`sketch_size_string`]); pass `None` when unavailable — the + /// `sketch_size` column is simply left null for that row, matching + /// the field's optional status. Only used for envelopes whose + /// `agg_id` matches `cfg.agg_id`; irrelevant when every envelope + /// in `envelopes` shares one `agg_id` (the common case — see + /// `crate::precompute::Precompute`'s "one instance owns one + /// `agg_id`" contract). + /// + /// `schema` / `dictionary` / `labels` rows appear only for + /// `agg_id`s / series not already marked known; `record` always + /// carries one row per envelope. An empty `envelopes` slice + /// produces four empty batches. + pub fn encode( + &mut self, + envelopes: &[SketchEnvelope], + cfg: Option<&PrecomputeConfig>, + ) -> Result { + let mut schema_agg_id: Vec = Vec::new(); + let mut schema_sketch_type: Vec<&'static str> = Vec::new(); + let mut schema_sketch_size: Vec> = Vec::new(); + let mut schema_hash_seed: Vec> = Vec::new(); + let mut schema_hash_function: Vec> = Vec::new(); + let mut schema_encoding: Vec<&'static str> = Vec::new(); + let mut schema_version_col: Vec = Vec::new(); + + let mut dict_series_id: Vec = Vec::new(); + let mut dict_agg_id: Vec = Vec::new(); + let mut dict_metric: Vec = Vec::new(); + + let mut labels_series_id: Vec = Vec::new(); + let mut labels_key: Vec = Vec::new(); + let mut labels_value: Vec> = Vec::new(); + + let mut rec_series_id: Vec = Vec::new(); + let mut rec_window_start: Vec = Vec::new(); + let mut rec_window_end: Vec = Vec::new(); + let mut rec_envelope: Vec>> = Vec::new(); + let mut rec_value: Vec> = Vec::new(); + + for env in envelopes { + if self.known_schemas.insert(env.agg_id) { + schema_agg_id.push(env.agg_id); + schema_sketch_type.push(env.sketch_type.name()); + schema_sketch_size.push( + cfg.filter(|c| c.agg_id == env.agg_id) + .and_then(|c| sketch_size_string(env.sketch_type, &c.sketch_params)), + ); + let (seed, function) = resolve_hash_seed(env.hash_spec.as_ref()); + schema_hash_seed.push(seed); + schema_hash_function.push(function); + schema_encoding.push(env.encoding.name()); + schema_version_col.push(env.schema_version); + } + + let series_id = self.series_id_for(env); + if self.known_series.insert(series_id) { + dict_series_id.push(series_id); + dict_agg_id.push(env.agg_id); + dict_metric.push(env.metric_name.clone()); + for kv in &env.labels { + labels_series_id.push(series_id); + labels_key.push(kv.key.clone()); + labels_value.push(Some(kv.value.clone())); + } + } + + rec_series_id.push(series_id); + rec_window_start.push(env.window_start_ms); + rec_window_end.push(env.window_end_ms); + // `RECORD` carries envelope bytes xor an estimate value, + // never both — `Precompute::serialize_series` already + // enforces non-empty payload for sketch-mode envelopes + // (empty-payload sketch rows are dropped before reaching + // here), so `payload.is_empty()` is an unambiguous + // discriminator between the two modes. + if env.payload.is_empty() { + rec_envelope.push(None); + rec_value.push(Some(env.value)); + } else { + rec_envelope.push(Some(env.payload.clone())); + rec_value.push(None); + } + } + + let schema = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(SCHEMA_COLUMN_AGG_ID, DataType::UInt64, false), + Field::new(SCHEMA_COLUMN_SKETCH_TYPE, DataType::Utf8, false), + Field::new(SCHEMA_COLUMN_SKETCH_SIZE, DataType::Utf8, true), + Field::new(SCHEMA_COLUMN_HASH_SEED, DataType::UInt64, true), + Field::new(SCHEMA_COLUMN_HASH_FUNCTION, DataType::Utf8, true), + Field::new(SCHEMA_COLUMN_ENCODING, DataType::Utf8, false), + Field::new(SCHEMA_COLUMN_SCHEMA_VERSION, DataType::UInt32, false), + ])), + vec![ + Arc::new(UInt64Array::from(schema_agg_id)), + Arc::new(StringArray::from(schema_sketch_type)), + Arc::new(StringArray::from(schema_sketch_size)), + Arc::new(UInt64Array::from(schema_hash_seed)), + Arc::new(StringArray::from(schema_hash_function)), + Arc::new(StringArray::from(schema_encoding)), + Arc::new(UInt32Array::from(schema_version_col)), + ], + ) + .map_err(|e| OtapEncodeError::ArrowError(e.to_string()))?; + + let dictionary = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(DICT_COLUMN_SERIES_ID, DataType::UInt32, false), + Field::new(SCHEMA_COLUMN_AGG_ID, DataType::UInt64, false), + Field::new(DICT_COLUMN_METRIC, DataType::Utf8, false), + ])), + vec![ + Arc::new(UInt32Array::from(dict_series_id)), + Arc::new(UInt64Array::from(dict_agg_id)), + Arc::new(StringArray::from(dict_metric)), + ], + ) + .map_err(|e| OtapEncodeError::ArrowError(e.to_string()))?; + + let labels = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(DICT_COLUMN_SERIES_ID, DataType::UInt32, false), + Field::new(LABELS_COLUMN_KEY, DataType::Utf8, false), + Field::new(LABELS_COLUMN_VALUE, DataType::Utf8, true), + ])), + vec![ + Arc::new(UInt32Array::from(labels_series_id)), + Arc::new(StringArray::from(labels_key)), + Arc::new(StringArray::from(labels_value)), + ], + ) + .map_err(|e| OtapEncodeError::ArrowError(e.to_string()))?; + + let record = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new(DICT_COLUMN_SERIES_ID, DataType::UInt32, false), + Field::new(RECORD_COLUMN_WINDOW_START_MS, DataType::UInt64, false), + Field::new(RECORD_COLUMN_WINDOW_END_MS, DataType::UInt64, false), + Field::new(RECORD_COLUMN_ENVELOPE, DataType::Binary, true), + Field::new(RECORD_COLUMN_VALUE, DataType::Float64, true), + ])), + vec![ + Arc::new(UInt32Array::from(rec_series_id)), + Arc::new(UInt64Array::from(rec_window_start)), + Arc::new(UInt64Array::from(rec_window_end)), + Arc::new(BinaryArray::from_opt_vec( + rec_envelope.iter().map(|o| o.as_deref()).collect(), + )), + Arc::new(Float64Array::from(rec_value)), + ], + ) + .map_err(|e| OtapEncodeError::ArrowError(e.to_string()))?; + + Ok(SketchStreamBatch { + schema, + dictionary, + labels, + record, + }) + } +} + +/// Resolves a [`asap_sketchlib::proto::sketchlib::HashSpec`] down to +/// the *one* seed this envelope's sketch actually used, plus which +/// algorithm it hashed with. +/// +/// `asap_sketchlib`'s own self-describing wire format +/// (`docs/asapv1_wire_format.md`) inlines the full 20-entry +/// `seed_list` plus several per-family index fields (canonical / +/// matrix / hydra / …) so a receiver can reconstruct *any* of the +/// hasher's seeds from the bytes alone — necessary there because one +/// producer process's `HashProfile` backs several concurrently-running +/// sketch families at once. `SCHEMA_COLUMN_HASH_SEED` doesn't need +/// that generality: one `SCHEMA` row already describes exactly one +/// `agg_id`'s one `sketch_type`, so there's exactly one seed position +/// that matters — `canonical_seed_index` is the one both libraries use +/// by default, and the field `HashSpec` actually exposes on this proto +/// message (`docs/data_model.md`'s `hash_seed` field is deliberately +/// this single resolved value, not the whole table). +/// +/// Returns `(None, None)` when `spec` is absent (nothing upstream +/// populates [`SketchEnvelope::hash_spec`] yet) or when +/// `canonical_seed_index` is out of bounds for `seed_list` (a +/// malformed spec — better to omit the seed than fabricate one). +fn resolve_hash_seed( + spec: Option<&asap_sketchlib::proto::sketchlib::HashSpec>, +) -> (Option, Option) { + let Some(spec) = spec else { + return (None, None); + }; + let seed = spec + .seed_list + .get(spec.canonical_seed_index as usize) + .copied(); + let function = asap_sketchlib::proto::sketchlib::HashAlgorithm::try_from(spec.algorithm) + .ok() + .map(|a| a.as_str_name().to_string()); + (seed, function) +} + +#[derive(Clone)] +struct SchemaFacts { + sketch_type: SketchType, + encoding: Encoding, + schema_version: u32, + hash_seed: Option, + hash_function: Option, +} + +/// Public snapshot of one `agg_id`'s retained `SCHEMA` facts, returned +/// by [`SeriesDictionaryDecoder::schema_for`]. Exists because +/// [`SketchEnvelope::hash_spec`] deliberately isn't reconstructed on +/// decode (see the comment in `build_records`) — a caller that needs +/// the resolved hash seed reads it from here instead. +#[derive(Clone, Debug, PartialEq)] +pub struct SchemaSnapshot { + /// Which sketch algorithm this `agg_id` runs. + pub sketch_type: SketchType, + /// Wire layout of this `agg_id`'s `RECORD.envelope` bytes. + pub encoding: Encoding, + /// Wire-schema version. + pub schema_version: u32, + /// The one resolved canonical hash seed, if this `agg_id`'s + /// sketch hashes at all — see `resolve_hash_seed`. + pub hash_seed: Option, + /// Which hash function `hash_seed` applies to (the proto + /// `HashAlgorithm`'s canonical name), if any. + pub hash_function: Option, +} + +#[derive(Clone)] +struct SeriesFacts { + agg_id: AggId, + metric: String, + labels: Vec, +} + +/// Receiver-side mirror of [`SeriesDictionary`]'s state — retains +/// every `SCHEMA` / `DICTIONARY` / `LABELS` row it has ever seen from +/// one continuous stream, so a bare `RECORD` row (just `series_id` + +/// window + envelope/value) can be joined back into a full +/// [`SketchEnvelope`]. +/// +/// One instance per **input stream** (i.e. per upstream sender this +/// node is receiving from), fed every [`SketchStreamBatch`] that +/// sender's [`SeriesDictionary`] produced, in order. See the module +/// doc's "Statefulness" section for the continuity contract this +/// assumes. +#[derive(Default)] +pub struct SeriesDictionaryDecoder { + schemas: HashMap, + series: HashMap, +} + +impl SeriesDictionaryDecoder { + /// Constructs a decoder with no retained state. + pub fn new() -> Self { + Self::default() + } + + /// Returns the retained `SCHEMA` facts for `agg_id`, or `None` if + /// this decoder has never ingested a `SCHEMA` row for it. + pub fn schema_for(&self, agg_id: AggId) -> Option { + self.schemas.get(&agg_id).map(|f| SchemaSnapshot { + sketch_type: f.sketch_type, + encoding: f.encoding, + schema_version: f.schema_version, + hash_seed: f.hash_seed, + hash_function: f.hash_function.clone(), + }) + } + + /// Ingests one [`SketchStreamBatch`], updating retained + /// `SCHEMA`/`DICTIONARY`/`LABELS` state from any new rows, and + /// reconstructs a full [`SketchEnvelope`] for every `RECORD` row + /// by joining back to that state (freshly-arrived or previously + /// retained). + /// + /// Returns [`OtapDecodeError::UnknownSeriesId`] / + /// [`OtapDecodeError::UnknownAggId`] if a `RECORD` (or + /// `DICTIONARY`) row references an identity this decoder has + /// never seen a defining row for — a continuity-contract + /// violation rather than a value to silently paper over. + pub fn decode( + &mut self, + batch: &SketchStreamBatch, + ) -> Result, OtapDecodeError> { + self.ingest_schema(&batch.schema)?; + self.ingest_dictionary(&batch.dictionary)?; + self.ingest_labels(&batch.labels)?; + self.build_records(&batch.record) + } + + fn ingest_schema(&mut self, batch: &RecordBatch) -> Result<(), OtapDecodeError> { + if batch.num_rows() == 0 { + return Ok(()); + } + let agg_id = col_u64(batch, "schema", SCHEMA_COLUMN_AGG_ID)?; + let sketch_type = col_str(batch, "schema", SCHEMA_COLUMN_SKETCH_TYPE)?; + let encoding = col_str(batch, "schema", SCHEMA_COLUMN_ENCODING)?; + let schema_version = col_u32(batch, "schema", SCHEMA_COLUMN_SCHEMA_VERSION)?; + let hash_seed = opt_u64(batch, SCHEMA_COLUMN_HASH_SEED)?; + let hash_function = opt_str(batch, SCHEMA_COLUMN_HASH_FUNCTION)?; + for row in 0..batch.num_rows() { + let id = agg_id.value(row); + let st = parse_sketch_type(row, sketch_type.value(row))?; + let enc = parse_encoding(row, encoding.value(row))?; + let seed = hash_seed + .as_ref() + .filter(|arr| !arr.is_null(row)) + .map(|arr| arr.value(row)); + let function = hash_function + .as_ref() + .filter(|arr| !arr.is_null(row)) + .map(|arr| arr.value(row).to_string()); + self.schemas.insert( + id, + SchemaFacts { + sketch_type: st, + encoding: enc, + schema_version: schema_version.value(row), + hash_seed: seed, + hash_function: function, + }, + ); + } + Ok(()) + } + + fn ingest_dictionary(&mut self, batch: &RecordBatch) -> Result<(), OtapDecodeError> { + if batch.num_rows() == 0 { + return Ok(()); + } + let series_id = col_u32(batch, "dictionary", DICT_COLUMN_SERIES_ID)?; + let agg_id = col_u64(batch, "dictionary", SCHEMA_COLUMN_AGG_ID)?; + let metric = col_str(batch, "dictionary", DICT_COLUMN_METRIC)?; + for row in 0..batch.num_rows() { + let sid = series_id.value(row); + self.series.insert( + sid, + SeriesFacts { + agg_id: agg_id.value(row), + metric: metric.value(row).to_string(), + labels: Vec::new(), + }, + ); + } + Ok(()) + } + + fn ingest_labels(&mut self, batch: &RecordBatch) -> Result<(), OtapDecodeError> { + if batch.num_rows() == 0 { + return Ok(()); + } + let series_id = col_u32(batch, "labels", DICT_COLUMN_SERIES_ID)?; + let key = col_str(batch, "labels", LABELS_COLUMN_KEY)?; + let value = opt_str(batch, LABELS_COLUMN_VALUE)?; + for row in 0..batch.num_rows() { + let sid = series_id.value(row); + let v = value + .as_ref() + .filter(|arr| !arr.is_null(row)) + .map(|arr| arr.value(row).to_string()) + .unwrap_or_default(); + // A LABELS row for a series_id this batch's own + // DICTIONARY didn't define (and no earlier batch did + // either) is dropped here; build_records still fails + // loudly for that series_id since it never resolves. + if let Some(entry) = self.series.get_mut(&sid) { + entry + .labels + .push(KeyValue::new(key.value(row).to_string(), v)); + } + } + Ok(()) + } + + fn build_records(&self, batch: &RecordBatch) -> Result, OtapDecodeError> { + let series_id = col_u32(batch, "record", DICT_COLUMN_SERIES_ID)?; + let window_start = col_u64(batch, "record", RECORD_COLUMN_WINDOW_START_MS)?; + let window_end = col_u64(batch, "record", RECORD_COLUMN_WINDOW_END_MS)?; + let envelope = opt_binary(batch, RECORD_COLUMN_ENVELOPE)?; + let value = opt_f64(batch, RECORD_COLUMN_VALUE)?; + + let mut out = Vec::with_capacity(batch.num_rows()); + for row in 0..batch.num_rows() { + let sid = series_id.value(row); + let series = self + .series + .get(&sid) + .ok_or(OtapDecodeError::UnknownSeriesId { series_id: sid })?; + let schema = self + .schemas + .get(&series.agg_id) + .ok_or(OtapDecodeError::UnknownAggId { + agg_id: series.agg_id, + })?; + + let payload = envelope + .as_ref() + .filter(|arr| !arr.is_null(row)) + .map(|arr| arr.value(row).to_vec()) + .unwrap_or_default(); + let est_value = value + .as_ref() + .filter(|arr| !arr.is_null(row)) + .map(|arr| arr.value(row)) + .unwrap_or(0.0); + + out.push(SketchEnvelope { + schema_version: schema.schema_version, + sketch_type: schema.sketch_type, + agg_id: series.agg_id, + resource_labels: Vec::new(), + labels: series.labels.clone(), + window_start_ms: window_start.value(row), + window_end_ms: window_end.value(row), + encoding: schema.encoding, + payload, + // Deliberately not reconstructed: `resolve_hash_seed` + // only carries the one resolved canonical seed across + // the wire, not `asap_sketchlib`'s full `HashSpec` + // (algorithm + 20-entry seed_list + seed_derivation) — + // synthesizing a fake one-entry `HashSpec` here would + // be more misleading than omitting it. A receiver that + // needs the resolved seed calls + // `Self::schema_for(series.agg_id)` instead of + // expecting it to round-trip through this field. + hash_spec: None, + metric_name: series.metric.clone(), + count: 0, + aggregation_temporality: 0, + value: est_value, + }); + } + Ok(out) + } +} + +// -- Small typed-column accessors ------------------------------------------- +// +// Deliberately local rather than shared with `records.rs` / `decode.rs`'s +// own downcast helpers — the batch shapes here are simple and fixed, and +// duplicating a handful of one-line downcasts is cheaper to read than a +// shared generic accessor would be. + +fn col_u64<'a>( + batch: &'a RecordBatch, + label: &'static str, + name: &'static str, +) -> Result<&'a UInt64Array, OtapDecodeError> { + let col = batch + .column_by_name(name) + .ok_or(OtapDecodeError::MissingColumn { + batch: label, + column: name, + })?; + col.as_any() + .downcast_ref::() + .ok_or_else(|| OtapDecodeError::WrongColumnType { + column: name.to_string(), + expected: "UInt64", + actual: col.data_type().clone(), + }) +} + +fn col_u32<'a>( + batch: &'a RecordBatch, + label: &'static str, + name: &'static str, +) -> Result<&'a UInt32Array, OtapDecodeError> { + let col = batch + .column_by_name(name) + .ok_or(OtapDecodeError::MissingColumn { + batch: label, + column: name, + })?; + col.as_any() + .downcast_ref::() + .ok_or_else(|| OtapDecodeError::WrongColumnType { + column: name.to_string(), + expected: "UInt32", + actual: col.data_type().clone(), + }) +} + +fn col_str<'a>( + batch: &'a RecordBatch, + label: &'static str, + name: &'static str, +) -> Result<&'a StringArray, OtapDecodeError> { + let col = batch + .column_by_name(name) + .ok_or(OtapDecodeError::MissingColumn { + batch: label, + column: name, + })?; + col.as_any() + .downcast_ref::() + .ok_or_else(|| OtapDecodeError::WrongColumnType { + column: name.to_string(), + expected: "Utf8", + actual: col.data_type().clone(), + }) +} + +fn opt_str<'a>( + batch: &'a RecordBatch, + name: &'static str, +) -> Result, OtapDecodeError> { + match batch.column_by_name(name) { + None => Ok(None), + Some(col) => Ok(Some( + col.as_any().downcast_ref::().ok_or_else(|| { + OtapDecodeError::WrongColumnType { + column: name.to_string(), + expected: "Utf8", + actual: col.data_type().clone(), + } + })?, + )), + } +} + +fn opt_u64<'a>( + batch: &'a RecordBatch, + name: &'static str, +) -> Result, OtapDecodeError> { + match batch.column_by_name(name) { + None => Ok(None), + Some(col) => Ok(Some( + col.as_any().downcast_ref::().ok_or_else(|| { + OtapDecodeError::WrongColumnType { + column: name.to_string(), + expected: "UInt64", + actual: col.data_type().clone(), + } + })?, + )), + } +} + +fn opt_binary<'a>( + batch: &'a RecordBatch, + name: &'static str, +) -> Result, OtapDecodeError> { + match batch.column_by_name(name) { + None => Ok(None), + Some(col) => Ok(Some( + col.as_any().downcast_ref::().ok_or_else(|| { + OtapDecodeError::WrongColumnType { + column: name.to_string(), + expected: "Binary", + actual: col.data_type().clone(), + } + })?, + )), + } +} + +fn opt_f64<'a>( + batch: &'a RecordBatch, + name: &'static str, +) -> Result, OtapDecodeError> { + match batch.column_by_name(name) { + None => Ok(None), + Some(col) => Ok(Some( + col.as_any().downcast_ref::().ok_or_else(|| { + OtapDecodeError::WrongColumnType { + column: name.to_string(), + expected: "Float64", + actual: col.data_type().clone(), + } + })?, + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::observation::KeyValue; + + fn envelope( + agg_id: u64, + metric: &str, + labels: Vec, + window: [u64; 2], + ) -> SketchEnvelope { + SketchEnvelope { + schema_version: 1, + sketch_type: SketchType::DDSketch, + agg_id, + resource_labels: Vec::new(), + labels, + window_start_ms: window[0], + window_end_ms: window[1], + encoding: Encoding::ProtoFull, + payload: vec![1, 2, 3, 4], + hash_spec: None, + metric_name: metric.to_string(), + count: 10, + aggregation_temporality: 1, + value: 0.0, + } + } + + #[test] + fn first_window_emits_schema_dictionary_labels_record() { + let mut dict = SeriesDictionary::new(); + let env = envelope( + 7, + "http_request_duration", + vec![ + KeyValue::new("path", "/api"), + KeyValue::new("region", "us-east"), + ], + [1_000, 11_000], + ); + let batch = dict + .encode(std::slice::from_ref(&env), None) + .expect("encode"); + + assert_eq!(batch.schema.num_rows(), 1, "new agg_id gets a SCHEMA row"); + assert_eq!( + batch.dictionary.num_rows(), + 1, + "new series gets a DICTIONARY row" + ); + assert_eq!( + batch.labels.num_rows(), + 2, + "two label keys on the new series" + ); + assert_eq!(batch.record.num_rows(), 1, "one RECORD row per envelope"); + assert!(!batch.is_empty()); + } + + #[test] + fn repeat_window_for_same_series_emits_record_only() { + let mut dict = SeriesDictionary::new(); + let env1 = envelope( + 7, + "http_request_duration", + vec![KeyValue::new("path", "/api")], + [1_000, 11_000], + ); + let _ = dict + .encode(std::slice::from_ref(&env1), None) + .expect("first window"); + + // Same series (same agg_id/metric/labels), next window. + let env2 = envelope( + 7, + "http_request_duration", + vec![KeyValue::new("path", "/api")], + [11_000, 21_000], + ); + let batch2 = dict + .encode(std::slice::from_ref(&env2), None) + .expect("second window"); + + assert_eq!( + batch2.schema.num_rows(), + 0, + "agg_id already known — no SCHEMA row" + ); + assert_eq!( + batch2.dictionary.num_rows(), + 0, + "series already known — no DICTIONARY row" + ); + assert_eq!( + batch2.labels.num_rows(), + 0, + "series already known — no LABELS rows" + ); + assert_eq!( + batch2.record.num_rows(), + 1, + "RECORD is still emitted every window" + ); + + // The series_id assigned in window 1 is reused in window 2. + let record_series_id = batch2 + .record + .column_by_name(DICT_COLUMN_SERIES_ID) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(record_series_id, 0); + } + + #[test] + fn distinct_label_combinations_get_distinct_series_ids() { + let mut dict = SeriesDictionary::new(); + let envs = vec![ + envelope(7, "m", vec![KeyValue::new("path", "/api")], [0, 10]), + envelope(7, "m", vec![KeyValue::new("path", "/login")], [0, 10]), + ]; + let batch = dict.encode(&envs, None).expect("encode"); + assert_eq!(batch.dictionary.num_rows(), 2, "two distinct series"); + let ids = batch + .dictionary + .column_by_name(DICT_COLUMN_SERIES_ID) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_ne!(ids.value(0), ids.value(1)); + } + + #[test] + fn empty_envelopes_produce_four_empty_batches() { + let mut dict = SeriesDictionary::new(); + let batch = dict.encode(&[], None).expect("encode empty"); + assert!(batch.is_empty()); + } + + #[test] + fn hash_seed_resolves_to_the_one_canonical_position() { + use asap_sketchlib::proto::sketchlib::{HashAlgorithm, HashSpec, SeedDerivation}; + + let spec = HashSpec { + algorithm: HashAlgorithm::Xxh364 as i32, + canonical_seed_index: 5, + seed_list: (0..20).map(|i| 1000 + i as u64).collect(), + seed_derivation: SeedDerivation::AdditiveOffset as i32, + }; + let mut dict = SeriesDictionary::new(); + let env = SketchEnvelope { + hash_spec: Some(spec), + sketch_type: SketchType::HLLSketch, + ..envelope(7, "m", vec![], [0, 10]) + }; + let batch = dict + .encode(std::slice::from_ref(&env), None) + .expect("encode"); + + let seeds = batch + .schema + .column_by_name(SCHEMA_COLUMN_HASH_SEED) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + // canonical_seed_index = 5 -> seed_list[5] = 1005, NOT the + // whole 20-entry table. + assert_eq!(seeds.value(0), 1005); + + let functions = batch + .schema + .column_by_name(SCHEMA_COLUMN_HASH_FUNCTION) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(functions.value(0), "HASH_ALGORITHM_XXH3_64"); + } + + #[test] + fn decoder_exposes_resolved_hash_seed_via_schema_for() { + use asap_sketchlib::proto::sketchlib::{HashAlgorithm, HashSpec, SeedDerivation}; + + let spec = HashSpec { + algorithm: HashAlgorithm::Xxh364 as i32, + canonical_seed_index: 5, + seed_list: (0..20).map(|i| 1000 + i as u64).collect(), + seed_derivation: SeedDerivation::AdditiveOffset as i32, + }; + let mut dict = SeriesDictionary::new(); + let mut decoder = SeriesDictionaryDecoder::new(); + let env = SketchEnvelope { + hash_spec: Some(spec), + sketch_type: SketchType::HLLSketch, + ..envelope(7, "m", vec![], [0, 10]) + }; + let batch = dict + .encode(std::slice::from_ref(&env), None) + .expect("encode"); + let decoded = decoder.decode(&batch).expect("decode"); + + // The reconstructed envelope itself doesn't carry a + // (necessarily lossy) HashSpec back... + assert!(decoded[0].hash_spec.is_none()); + // ...but the decoder retains the resolved seed for direct + // lookup by agg_id. + let schema = decoder.schema_for(7).expect("schema retained"); + assert_eq!(schema.hash_seed, Some(1005)); + assert_eq!( + schema.hash_function.as_deref(), + Some("HASH_ALGORITHM_XXH3_64") + ); + assert_eq!(schema.sketch_type, SketchType::HLLSketch); + } + + #[test] + fn hash_seed_is_null_without_a_hash_spec() { + let mut dict = SeriesDictionary::new(); + let env = envelope(7, "m", vec![], [0, 10]); // hash_spec: None (default) + let batch = dict + .encode(std::slice::from_ref(&env), None) + .expect("encode"); + let seeds = batch + .schema + .column_by_name(SCHEMA_COLUMN_HASH_SEED) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(seeds.is_null(0)); + } + + #[test] + fn sketch_size_is_populated_from_config_when_agg_id_matches() { + use crate::config::{PrecomputeConfig, SketchParams}; + let mut params = SketchParams::new(); + params.insert("relative_accuracy".into(), 0.01); + let cfg = PrecomputeConfig { + agg_id: 7, + sketch_params: params, + ..Default::default() + }; + let mut dict = SeriesDictionary::new(); + let env = envelope(7, "m", vec![], [0, 10]); + let batch = dict + .encode(std::slice::from_ref(&env), Some(&cfg)) + .expect("encode"); + let sizes = batch + .schema + .column_by_name(SCHEMA_COLUMN_SKETCH_SIZE) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(sizes.value(0), "0.01"); + } + + #[test] + fn round_trip_two_windows_preserves_envelopes() { + let mut dict = SeriesDictionary::new(); + let mut decoder = SeriesDictionaryDecoder::new(); + + let env1 = envelope( + 7, + "http_request_duration", + vec![ + KeyValue::new("path", "/api"), + KeyValue::new("region", "us-east"), + ], + [1_000, 11_000], + ); + let batch1 = dict + .encode(std::slice::from_ref(&env1), None) + .expect("encode 1"); + let decoded1 = decoder.decode(&batch1).expect("decode 1"); + assert_eq!(decoded1.len(), 1); + assert_eq!(decoded1[0].payload, env1.payload); + assert_eq!(decoded1[0].metric_name, env1.metric_name); + assert_eq!(decoded1[0].labels, env1.labels); + assert_eq!(decoded1[0].sketch_type, env1.sketch_type); + assert_eq!(decoded1[0].window_start_ms, env1.window_start_ms); + assert_eq!(decoded1[0].window_end_ms, env1.window_end_ms); + + // Second window: sender's DICTIONARY/SCHEMA batches are empty + // (already known), so the encoded bytes genuinely shrink — but + // the decoder must still reconstruct the full envelope by + // joining the bare RECORD row against retained state. + let env2 = envelope( + 7, + "http_request_duration", + vec![ + KeyValue::new("path", "/api"), + KeyValue::new("region", "us-east"), + ], + [11_000, 21_000], + ); + let batch2 = dict + .encode(std::slice::from_ref(&env2), None) + .expect("encode 2"); + assert!(batch2.schema.num_rows() == 0 && batch2.dictionary.num_rows() == 0); + let decoded2 = decoder.decode(&batch2).expect("decode 2"); + assert_eq!(decoded2.len(), 1); + assert_eq!(decoded2[0].labels, env2.labels); + assert_eq!(decoded2[0].metric_name, env2.metric_name); + assert_eq!(decoded2[0].window_start_ms, 11_000); + assert_eq!(decoded2[0].window_end_ms, 21_000); + } + + #[test] + fn decode_unknown_series_id_is_a_hard_error() { + // A RECORD batch referencing a series_id with no prior + // DICTIONARY row (fresh decoder, no schema/dictionary/labels + // ingested first) must fail loudly, not synthesize a + // half-empty envelope. + let mut dict = SeriesDictionary::new(); + let env = envelope(7, "m", vec![], [0, 10]); + let batch = dict + .encode(std::slice::from_ref(&env), None) + .expect("encode"); + + // Fresh decoder that never saw batch's SCHEMA/DICTIONARY rows + // — simulate by decoding only the RECORD-bearing part via a + // hand-built batch with empty schema/dictionary/labels. + let empty_schema = batch.schema.slice(0, 0); + let empty_dict = batch.dictionary.slice(0, 0); + let empty_labels = batch.labels.slice(0, 0); + let record_only = SketchStreamBatch { + schema: empty_schema, + dictionary: empty_dict, + labels: empty_labels, + record: batch.record, + }; + let mut decoder = SeriesDictionaryDecoder::new(); + let err = decoder.decode(&record_only).expect_err("should fail"); + assert!(matches!( + err, + OtapDecodeError::UnknownSeriesId { series_id: 0 } + )); + } +} diff --git a/asap-precompute-rs/src/otap/lifecycle.rs b/asap-precompute-rs/src/otap/lifecycle.rs index 561e8f2..85165ca 100644 --- a/asap-precompute-rs/src/otap/lifecycle.rs +++ b/asap-precompute-rs/src/otap/lifecycle.rs @@ -20,13 +20,19 @@ //! //! 2. **Flush ticker** — modelled on OTAP's `NodeControlMsg::Wakeup` //! (a Tokio `interval(window_size)` here). Each tick calls -//! `Precompute::tick(now_ms)`, encodes the resulting envelopes -//! via [`super::encode_batch`], lifts the Strategy-B carrier -//! columns onto the per-row attribute child batch via -//! [`super::records::lift`], and pushes the -//! [`super::records::OtapMetricRecords`] family onto the emit -//! channel. (Phase D wires this to OTAP's -//! `effect_handler.send_message`.) +//! `Precompute::tick(now_ms)` and encodes the resulting envelopes +//! against this plugin's [`super::dictionary::SeriesDictionary`] +//! state, pushing the resulting +//! [`super::dictionary::SketchStreamBatch`] onto the emit channel — +//! this is the node-to-node sketch-stream hop +//! `docs/data_model.md` describes, so `SCHEMA`/`DICTIONARY` rows +//! only ride the wire the first time an `agg_id`/series is seen. +//! (Phase D wires this to OTAP's `effect_handler.send_message`.) +//! The input task below is unrelated and keeps using the flat +//! OTAP-Metrics-shaped codec — it accepts arbitrary upstream OTAP +//! producers (raw telemetry, not necessarily another +//! `asap_sketches` node), which is exactly the compatibility +//! `encode_batch`/`decode_batch`/[`super::records`] exist for. //! //! 3. **Control-channel task** — polls the //! [`crate::control_channel::ControlChannel`] every @@ -52,8 +58,9 @@ use crate::envelope::SketchEnvelope; use crate::precompute::{Precompute, PrecomputeError, PrecomputeImpl, StatsSnapshot}; use super::config::{resolve, ConfigError, PluginConfig}; -use super::records::{flatten, lift, OtapMetricRecords, OtapRecordsError}; -use super::{decode_batch, encode_batch, OtapDecodeError, OtapEncodeError}; +use super::dictionary::{SeriesDictionary, SketchStreamBatch}; +use super::records::{flatten, OtapMetricRecords, OtapRecordsError}; +use super::{decode_batch, OtapDecodeError, OtapEncodeError}; /// Default poll cadence for the control-channel task. Fast enough /// that operators see plan @@ -93,11 +100,11 @@ pub enum PluginError { /// Convenience type — the emit channel sender shared by the flush /// ticker and the drain path. -pub type EmitSender = mpsc::UnboundedSender; +pub type EmitSender = mpsc::UnboundedSender; /// Convenience type — the emit channel receiver returned to the /// caller (i.e. tests + the Phase D OTAP shell). -pub type EmitReceiver = mpsc::UnboundedReceiver; +pub type EmitReceiver = mpsc::UnboundedReceiver; /// `AsapSketchesPlugin` — the Layer-4 plugin lifecycle. Replaces /// Phase B's `StubPlugin

` with a real Tokio-based runtime around @@ -117,6 +124,13 @@ pub type EmitReceiver = mpsc::UnboundedReceiver; pub struct AsapSketchesPlugin { inner: Arc, window_size: Duration, + /// Outbound `SCHEMA`/`DICTIONARY` state for this plugin's emit + /// stream — persists for the plugin's whole lifetime (across every + /// tick and the final drain) so repeat windows for an already-known + /// series cost only a `RECORD` row. `tokio::sync::Mutex` because + /// it's shared between the ticker task and the supervisor's final + /// drain. + dictionary: Arc>, } impl AsapSketchesPlugin { @@ -134,6 +148,7 @@ impl AsapSketchesPlugin { Ok(Self { inner: Arc::new(pc), window_size: pcfg.window.size, + dictionary: Arc::new(Mutex::new(SeriesDictionary::new())), }) } @@ -148,6 +163,7 @@ impl AsapSketchesPlugin { Self { inner: precompute, window_size, + dictionary: Arc::new(Mutex::new(SeriesDictionary::new())), } } @@ -193,11 +209,13 @@ impl AsapSketchesPlugin { let precompute = self.inner.clone(); let window_size = self.window_size; + let dictionary = self.dictionary.clone(); let opts = Arc::new(opts); let input_task = spawn_input_task(precompute.clone(), input, cancellation.clone()); let ticker_task = spawn_ticker_task( precompute.clone(), + dictionary.clone(), window_size, emit_tx.clone(), cancellation.clone(), @@ -225,7 +243,9 @@ impl AsapSketchesPlugin { // Final drain — flush any in-flight window before exit. let envs = precompute.drain(); if !envs.is_empty() { - let _ = emit_drain(&emit_tx, &envs); + let cfg = precompute.active_config(); + let mut dict = dictionary.lock().await; + let _ = emit_drain(&emit_tx, &envs, &mut dict, cfg.as_ref()); } }); @@ -345,6 +365,7 @@ fn ingest_one_batch( fn spawn_ticker_task( precompute: Arc, + dictionary: Arc>, window_size: Duration, emit_tx: EmitSender, cancel: Cancellation, @@ -365,7 +386,9 @@ fn spawn_ticker_task( if envs.is_empty() { continue; } - let _ = emit_envelopes(&emit_tx, &envs); + let cfg = precompute.active_config(); + let mut dict = dictionary.lock().await; + let _ = emit_envelopes(&emit_tx, &envs, &mut dict, cfg.as_ref()); } } } @@ -398,15 +421,24 @@ fn spawn_control_task( }) } -fn emit_envelopes(emit_tx: &EmitSender, envelopes: &[SketchEnvelope]) -> Result<(), PluginError> { - let flat = encode_batch(envelopes)?; - let lifted = lift(&flat)?; - let _ = emit_tx.send(lifted); // receiver may have been dropped on shutdown +fn emit_envelopes( + emit_tx: &EmitSender, + envelopes: &[SketchEnvelope], + dictionary: &mut SeriesDictionary, + cfg: Option<&PrecomputeConfig>, +) -> Result<(), PluginError> { + let batch = dictionary.encode(envelopes, cfg)?; + let _ = emit_tx.send(batch); // receiver may have been dropped on shutdown Ok(()) } -fn emit_drain(emit_tx: &EmitSender, envelopes: &[SketchEnvelope]) -> Result<(), PluginError> { - emit_envelopes(emit_tx, envelopes) +fn emit_drain( + emit_tx: &EmitSender, + envelopes: &[SketchEnvelope], + dictionary: &mut SeriesDictionary, + cfg: Option<&PrecomputeConfig>, +) -> Result<(), PluginError> { + emit_envelopes(emit_tx, envelopes, dictionary, cfg) } /// Wall-clock millisecond timestamp. Wraps `SystemTime::now()` so diff --git a/asap-precompute-rs/src/otap/mod.rs b/asap-precompute-rs/src/otap/mod.rs index 3190a96..fb96837 100644 --- a/asap-precompute-rs/src/otap/mod.rs +++ b/asap-precompute-rs/src/otap/mod.rs @@ -50,6 +50,36 @@ //! codec's flat shape is also the easiest to round-trip in a unit //! test. //! +//! # Schema / Dictionary / Record stream ([`dictionary`]) +//! +//! [`encode_batch`] / [`decode_batch`] above are a *different* codec +//! from [`dictionary::SeriesDictionary`] / [`dictionary::SeriesDictionaryDecoder`], +//! not an earlier draft of it — they solve different problems: +//! +//! - `encode_batch` / `decode_batch` (+ [`records::flatten`] / +//! [`records::lift`]) make a `SketchEnvelope` **look like** a +//! generic OTAP-Metrics payload — one self-contained row per +//! envelope, `_asap_*` carrier keys lifted onto the per-row +//! attribute child batch — so it can transit an OTAP pipeline hop +//! that only knows how to move Logs/Metrics/Traces payloads. +//! - [`dictionary::SeriesDictionary`] / [`dictionary::SeriesDictionaryDecoder`] +//! implement `docs/data_model.md`'s `SCHEMA` / `DICTIONARY` / `RECORD` +//! tiering for the hop that doc actually describes: sketch state +//! crossing a node boundary between two `asap_sketches` processor +//! instances (an ASAP-aware sender talking to an ASAP-aware +//! receiver — see [`Precompute::tick`](crate::precompute::Precompute::tick) / +//! `drain`, which is exactly that boundary). There, config-level +//! facts (`sketch_type`, `sketch_size`, `encoding`, …) are sent once +//! per `agg_id` and series identity (`metric` + labels) once per +//! distinct series — not repeated on every window's `RECORD` row — +//! because both ends are expected to retain that state across the +//! stream, the same way an Arrow IPC decoder retains Schema / +//! Dictionary state. [`AsapSketchesPlugin`] and [`StubPlugin`] both +//! use this codec for their tick/drain (encode) and inbound-envelope +//! (decode) paths; `encode_batch`/`decode_batch` remain for raw +//! (non-envelope) observation ingestion and for whatever still needs +//! OTAP-Metrics-payload compatibility. +//! //! # Phase C — full plugin lifecycle //! //! Phase B shipped the stateless codec (`decode_batch` / @@ -93,6 +123,7 @@ //! Phase B's tests passing as a regression backstop. mod decode; +mod dictionary; mod encode; mod plugin; mod schema; @@ -102,6 +133,9 @@ pub mod lifecycle; pub mod records; pub use decode::{decode_batch, OtapDecodeError}; +pub use dictionary::{ + SchemaSnapshot, SeriesDictionary, SeriesDictionaryDecoder, SketchStreamBatch, +}; pub use encode::{encode_batch, OtapEncodeError}; pub use plugin::StubPlugin; pub use schema::{ diff --git a/asap-precompute-rs/src/otap/plugin.rs b/asap-precompute-rs/src/otap/plugin.rs index 415868d..24556da 100644 --- a/asap-precompute-rs/src/otap/plugin.rs +++ b/asap-precompute-rs/src/otap/plugin.rs @@ -18,12 +18,15 @@ //! `otap-patch/plugins/asap_sketches/`. Comments below mark the seams //! Phase C will fill in. +use std::sync::Mutex; + use arrow_array::RecordBatch; use crate::envelope::SketchEnvelope; use crate::precompute::{Precompute, PrecomputeError}; -use super::{decode_batch, encode_batch, OtapDecodeError, OtapEncodeError}; +use super::dictionary::{SeriesDictionary, SeriesDictionaryDecoder, SketchStreamBatch}; +use super::{decode_batch, OtapDecodeError, OtapEncodeError}; /// Failure modes from the stub plugin's process-and-emit cycle. #[derive(Debug, thiserror::Error)] @@ -51,14 +54,29 @@ pub enum StubPluginError { /// Generic over `P: Precompute` so that downstream tests / Phase C's /// plugin can pass a real `PrecomputeImpl` (or any `Precompute` impl) /// without forcing a `Box` allocation. +/// +/// Owns a [`SeriesDictionary`] (outbound) and a +/// [`SeriesDictionaryDecoder`] (inbound) — see `otap/mod.rs`'s +/// "Schema / Dictionary / Record stream" section for why this plugin +/// uses that codec rather than `encode_batch`/`decode_batch` for the +/// node-to-node envelope hop. Both are per-instance state that must +/// persist across calls, matching the continuous-stream contract +/// `docs/data_model.md` assumes. pub struct StubPlugin { precompute: P, + dictionary: Mutex, + decoder: Mutex, } impl StubPlugin

{ - /// Construct a stub plugin around an existing `Precompute`. + /// Construct a stub plugin around an existing `Precompute`, with a + /// fresh (nothing-sent-yet) dictionary and decoder. pub fn new(precompute: P) -> Self { - Self { precompute } + Self { + precompute, + dictionary: Mutex::new(SeriesDictionary::new()), + decoder: Mutex::new(SeriesDictionaryDecoder::new()), + } } /// Borrow the wrapped `Precompute`. Tests use this to inspect @@ -73,6 +91,11 @@ impl StubPlugin

{ /// in `Precompute::observe` redirects pre-aggregated rows to /// `observe_envelope` — see `precompute.rs::observe`. /// + /// For **raw** (non-envelope) observations riding a flat, + /// OTAP-Metrics-shaped `RecordBatch`. Inbound pre-aggregated + /// envelopes from another `asap_sketches` node use + /// [`Self::ingest_stream`] instead. + /// /// **Phase C will replace this with the OTAP `Processor::process` /// method body** that consumes from the input /// `Stream` and pushes errors onto OTAP's @@ -86,25 +109,46 @@ impl StubPlugin

{ Ok(()) } + /// Decode a [`SketchStreamBatch`] from an upstream `asap_sketches` + /// node's [`SeriesDictionary`], reconstructing full envelopes via + /// this plugin's retained [`SeriesDictionaryDecoder`] state, and + /// route each through `Precompute::observe_envelope`. + pub fn ingest_stream(&self, batch: &SketchStreamBatch) -> Result<(), StubPluginError> { + let envelopes = { + let mut decoder = self.decoder.lock().expect("decoder lock poisoned"); + decoder.decode(batch)? + }; + for env in &envelopes { + self.precompute.observe_envelope(env)?; + } + Ok(()) + } + /// Force a `Precompute::tick` and encode the resulting envelopes - /// as a fresh `RecordBatch` ready for emit. + /// against this plugin's [`SeriesDictionary`] state. /// /// **Phase C will replace this with a `NodeControlMsg::Wakeup` /// handler driven by an `interval(window_size)` Tokio timer** — /// at which point this method becomes the body of the timer /// callback. For Phase B we expose it as an explicit method so a /// unit test can drive it deterministically. - pub fn tick(&self, now_ms: u64) -> Result { + pub fn tick(&self, now_ms: u64) -> Result { let envelopes: Vec = self.precompute.tick(now_ms); - encode_batch(&envelopes) + self.encode(&envelopes) } /// Force a `Precompute::drain` (graceful shutdown flush) and /// encode the result. Phase C's `NodeControlMsg::Shutdown` /// handler calls this once before dropping the plugin. - pub fn drain(&self) -> Result { + pub fn drain(&self) -> Result { let envelopes = self.precompute.drain(); - encode_batch(&envelopes) + self.encode(&envelopes) + } + + fn encode(&self, envelopes: &[SketchEnvelope]) -> Result { + let cfg = self.precompute.active_config(); + let mut dictionary = self.dictionary.lock().expect("dictionary lock poisoned"); + dictionary.encode(envelopes, cfg.as_ref()) } } @@ -113,6 +157,7 @@ mod tests { use super::*; use crate::config::{PrecomputeConfig, PrecomputeConfigSet, WindowSpec}; use crate::envelope::SketchType; + use crate::otap::encode_batch; use crate::precompute::PrecomputeImpl; use std::time::Duration; @@ -150,9 +195,43 @@ mod tests { }); let plugin = StubPlugin::new(precompute); - // No observations were ever pushed, so tick should return a - // schema-only batch with zero rows. + // No observations were ever pushed, so tick should return + // four empty batches. let out = plugin.tick(123_000).expect("tick"); - assert_eq!(out.num_rows(), 0); + assert!(out.is_empty()); + } + + #[test] + fn ingest_stream_round_trips_through_dictionary_into_observe_envelope() { + // Sender side: a PrecomputeImpl producing envelopes, encoded + // via a fresh SeriesDictionary. + let cfg = PrecomputeConfig { + agg_id: 7, + sketch_type: SketchType::DDSketch, + window: WindowSpec { + size: Duration::from_millis(1), + ..Default::default() + }, + ..Default::default() + }; + let sender = PrecomputeImpl::new(Some(cfg.clone()), None, None); + sender.update_config(&PrecomputeConfigSet { + version: 1, + configs: vec![cfg.clone()], + }); + let sender_plugin = StubPlugin::new(sender); + // Force a window rotation with no observations — exercises the + // empty-envelopes path deterministically without needing a + // real sketch factory. + let batch = sender_plugin.tick(u64::MAX).expect("tick"); + assert!(batch.is_empty()); + + // Receiver side: ingest_stream must accept the (empty) stream + // batch without error even though this receiver's Precompute + // has no config installed (observe_envelope is simply never + // called for zero envelopes). + let receiver = PrecomputeImpl::new(None, None, None); + let receiver_plugin = StubPlugin::new(receiver); + assert!(receiver_plugin.ingest_stream(&batch).is_ok()); } } diff --git a/asap-precompute-rs/src/otap/schema.rs b/asap-precompute-rs/src/otap/schema.rs index 82857d3..3024590 100644 --- a/asap-precompute-rs/src/otap/schema.rs +++ b/asap-precompute-rs/src/otap/schema.rs @@ -75,6 +75,95 @@ pub fn is_reserved_column(name: &str) -> bool { ) } +// ---------- Schema / Dictionary / Record stream columns ---------- +// +// Column names for the four-batch family produced by +// [`super::dictionary::SeriesDictionary::encode`], mirroring the ER +// diagram in `docs/data_model.md#schema--dictionary--record-as-entities`: +// `SCHEMA` is keyed by `agg_id`, `DICTIONARY` (+ its child `LABELS`) is +// keyed by `series_id`, and `RECORD` references a `DICTIONARY` entry by +// `series_id` instead of repeating `metric`/labels inline. Unlike the +// `ATTR_*` carrier keys above (which exist to smuggle these facts +// through a single OTAP-Metrics-shaped row), these four batches are +// ASAP's own inter-node sketch-stream wire shape — see the doc's +// opening line: "information carried when sketch state crosses a node +// or network boundary between `asap_sketches` processor instances." + +/// `SCHEMA.agg_id` / `DICTIONARY.agg_id` — controller-plan join key. +/// `UInt64` column. +pub const SCHEMA_COLUMN_AGG_ID: &str = "agg_id"; + +/// `SCHEMA.sketch_type`. `Utf8` column; same canonical names as +/// [`ATTR_SKETCH_TYPE`]. +pub const SCHEMA_COLUMN_SKETCH_TYPE: &str = "sketch_type"; + +/// `SCHEMA.sketch_size` — the algorithm's size/accuracy parameter +/// (relative accuracy, buffer size `k`, precision, or `width x depth` +/// — whichever `sketch_type` calls for), rendered as a string so one +/// column covers every algorithm's parameter shape without a lossy +/// numeric union. `Utf8`, optional. +pub const SCHEMA_COLUMN_SKETCH_SIZE: &str = "sketch_size"; + +/// `SCHEMA.hash_seed` — determinism contract for hash-based sketches: +/// the single seed value at +/// `hash_spec.seed_list[hash_spec.canonical_seed_index]`, resolved +/// from [`crate::envelope::SketchEnvelope::hash_spec`] by +/// `super::dictionary::resolve_hash_seed`. `UInt64`, optional — null +/// when the envelope carries no `hash_spec` (nothing upstream +/// populates it yet) or the sketch type doesn't hash at all (DDSketch, +/// KLL). Deliberately just the one resolved seed, not +/// `asap_sketchlib`'s full 20-entry `seed_list` — see +/// `resolve_hash_seed`'s doc comment for why one `SCHEMA` row (one +/// `agg_id`, one `sketch_type`) only ever needs one seed position. +pub const SCHEMA_COLUMN_HASH_SEED: &str = "hash_seed"; + +/// `SCHEMA.hash_function` — which hash function, for algorithms that +/// need one (`hash_spec.algorithm`'s canonical proto name, e.g. +/// `"HASH_ALGORITHM_XXH3_64"`). `Utf8`, optional. Same +/// null-when-no-`hash_spec` caveat as [`SCHEMA_COLUMN_HASH_SEED`]. +pub const SCHEMA_COLUMN_HASH_FUNCTION: &str = "hash_function"; + +/// `SCHEMA.encoding`. `Utf8` column; same canonical names as +/// [`ATTR_ENCODING`]. +pub const SCHEMA_COLUMN_ENCODING: &str = "encoding"; + +/// `SCHEMA.schema_version`. `UInt32` column. +pub const SCHEMA_COLUMN_SCHEMA_VERSION: &str = "schema_version"; + +/// `DICTIONARY.series_id` / `LABELS.series_id` / `RECORD.series_id` — +/// primary key of a `DICTIONARY` entry, and the join key `RECORD` uses +/// instead of repeating `metric`/labels inline. `UInt32` column. +pub const DICT_COLUMN_SERIES_ID: &str = "series_id"; + +/// `DICTIONARY.metric` — metric name. `Utf8` column. +pub const DICT_COLUMN_METRIC: &str = "metric"; + +/// `LABELS.key` — label key. `Utf8` column. +pub const LABELS_COLUMN_KEY: &str = "key"; + +/// `LABELS.value` — label value. `Utf8` column, optional. +pub const LABELS_COLUMN_VALUE: &str = "value"; + +/// `RECORD.window_start_ms` — inclusive lower bound of the window this +/// record summarizes (Unix milliseconds). `UInt64` column. +pub const RECORD_COLUMN_WINDOW_START_MS: &str = "window_start_ms"; + +/// `RECORD.window_end_ms` — exclusive upper bound of the window this +/// record summarizes (Unix milliseconds). `UInt64` column. +pub const RECORD_COLUMN_WINDOW_END_MS: &str = "window_end_ms"; + +/// `RECORD.envelope` — serialized sketch state or delta. `Binary` +/// column, optional (mutually exclusive with [`RECORD_COLUMN_VALUE`] +/// per record — a `RECORD` carries sketch state or an estimate, never +/// both). +pub const RECORD_COLUMN_ENVELOPE: &str = "envelope"; + +/// `RECORD.value` — estimate-mode scalar (quantile or cardinality +/// estimate), carried instead of [`RECORD_COLUMN_ENVELOPE`] when the +/// series emits estimates rather than sketch state. `Float64` column, +/// optional. +pub const RECORD_COLUMN_VALUE: &str = "value"; + #[cfg(test)] mod tests { use super::*; diff --git a/asap-precompute-rs/src/precompute.rs b/asap-precompute-rs/src/precompute.rs index c843617..7c3b1db 100644 --- a/asap-precompute-rs/src/precompute.rs +++ b/asap-precompute-rs/src/precompute.rs @@ -327,6 +327,17 @@ pub trait Precompute: Send + Sync { /// are). fn update_config(&self, cs: &PrecomputeConfigSet); + /// Returns a clone of the active config, or `None` if unconfigured. + /// + /// Used by the OTAP encode layer + /// ([`crate::otap::dictionary::SeriesDictionary`]) to source + /// `SCHEMA`-tier facts (`sketch_params`, and eventually the hash + /// determinism contract) that live on the config rather than being + /// re-derived per envelope — those facts change at config-time, + /// not per-record, so they don't belong on + /// [`crate::envelope::SketchEnvelope`] itself. + fn active_config(&self) -> Option; + /// Returns the live counters; safe to call concurrently. fn stats(&self) -> StatsSnapshot; @@ -731,6 +742,18 @@ impl Precompute for PrecomputeImpl { self.finish_rotate(closed, rng, rng[1]) } + fn active_config(&self) -> Option { + // Duplicated (rather than delegating to the inherent + // `PrecomputeImpl::active_config` above) so the trait method's + // body doesn't read as a same-name recursive call — inherent + // methods win method resolution over trait methods for the + // same receiver, so `self.active_config()` here would in fact + // call the inherent one and not recurse, but spelling out the + // two-line lock+clone directly is clearer than relying on that + // priority rule. + self.cfg.lock().expect("config lock poisoned").clone() + } + fn update_config(&self, cs: &PrecomputeConfigSet) { if cs.configs.is_empty() { return; diff --git a/asap-precompute-rs/tests/otap_lifecycle.rs b/asap-precompute-rs/tests/otap_lifecycle.rs index 86cddd1..099f3c1 100644 --- a/asap-precompute-rs/tests/otap_lifecycle.rs +++ b/asap-precompute-rs/tests/otap_lifecycle.rs @@ -33,18 +33,15 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Duration; -use arrow_array::{ - Array, BinaryArray, Float64Array, RecordBatch, StringArray, UInt32Array, UInt64Array, -}; +use arrow_array::{BinaryArray, Float64Array, RecordBatch, StringArray, UInt32Array, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; use asap_precompute_rs::config::{PrecomputeConfig, PrecomputeConfigSet, WindowSpec}; use asap_precompute_rs::control_channel::ControlChannel; -use asap_precompute_rs::envelope::{Encoding, SketchType}; +use asap_precompute_rs::envelope::{Encoding, SketchEnvelope, SketchType}; use asap_precompute_rs::otap::{ - AsapSketchesPlugin, OtapMetricRecords, PluginConfig, PluginHandle, StartOptions, ATTR_AGG_ID, - ATTR_ENCODING, ATTR_ENVELOPE, ATTR_SCHEMA_VERSION, ATTR_SKETCH_TYPE, ATTR_WINDOW_END_MS, - ATTR_WINDOW_START_MS, COLUMN_METRIC, COLUMN_TIME_UNIX_NANO, COLUMN_VALUE, + AsapSketchesPlugin, OtapMetricRecords, PluginConfig, PluginHandle, SeriesDictionaryDecoder, + SketchStreamBatch, StartOptions, COLUMN_METRIC, COLUMN_TIME_UNIX_NANO, COLUMN_VALUE, }; const PARENT_ID_COL: &str = "parent_id"; @@ -101,7 +98,7 @@ fn scalar_records(metric: &str, value: f64, timestamp_ms: u64, host: &str) -> Ot /// Drain an [`asap_precompute_rs::otap::EmitReceiver`] with a /// generous timeout. The lifecycle tasks emit eagerly on shutdown, /// so 5s is far more than needed in practice. -async fn drain_emit(rx: &mut asap_precompute_rs::otap::EmitReceiver) -> Vec { +async fn drain_emit(rx: &mut asap_precompute_rs::otap::EmitReceiver) -> Vec { let mut out = Vec::new(); let timeout = Duration::from_secs(5); let deadline = tokio::time::Instant::now() + timeout; @@ -111,7 +108,7 @@ async fn drain_emit(rx: &mut asap_precompute_rs::otap::EmitReceiver) -> Vec out.push(records), + Ok(Some(batch)) => out.push(batch), Ok(None) => break, Err(_) => break, } @@ -119,61 +116,36 @@ async fn drain_emit(rx: &mut asap_precompute_rs::otap::EmitReceiver) -> Vec Vec { - let attrs = &records.attributes; - let key_col = attrs - .column_by_name(ATTR_KEY_COL) - .expect("attr key column") - .as_any() - .downcast_ref::() - .expect("key Utf8"); - let bytes_col = attrs - .column_by_name(ATTR_BYTES_COL) - .expect("attr bytes column") - .as_any() - .downcast_ref::() - .expect("bytes Binary"); - let str_col = attrs - .column_by_name(ATTR_STR_COL) - .expect("attr str column") - .as_any() - .downcast_ref::() - .expect("str Utf8"); - let parent_col = attrs - .column_by_name(PARENT_ID_COL) - .expect("attr parent_id column") - .as_any() - .downcast_ref::() - .expect("parent_id UInt32"); - - // Pair parent_id with envelope payload + sketch_type by walking - // attribute rows and grouping by parent_id. - let mut payload_by_parent: std::collections::BTreeMap> = Default::default(); - let mut type_by_parent: std::collections::BTreeMap = Default::default(); - for row in 0..attrs.num_rows() { - let pid = parent_col.value(row); - let key = key_col.value(row); - match key { - ATTR_ENVELOPE if !bytes_col.is_null(row) => { - payload_by_parent.insert(pid, bytes_col.value(row).to_vec()); - } - ATTR_SKETCH_TYPE if !str_col.is_null(row) => { - type_by_parent.insert(pid, str_col.value(row).to_string()); - } - _ => {} - } +/// Decode every emitted [`SketchStreamBatch`] through one shared +/// [`SeriesDictionaryDecoder`], in order — matching the +/// continuous-stream contract `docs/data_model.md` assumes (a `RECORD` +/// row past the first window carries no `metric`/labels of its own, +/// only a `series_id` referencing a `DICTIONARY` row an earlier batch +/// in this same sequence carried). +fn decode_all(batches: &[SketchStreamBatch]) -> Vec { + let mut decoder = SeriesDictionaryDecoder::new(); + let mut out = Vec::new(); + for batch in batches { + out.extend(decoder.decode(batch).expect("decode stream batch")); } - let mut matched: Vec> = payload_by_parent - .into_iter() - .filter_map(|(pid, bytes)| { - type_by_parent - .get(&pid) - .filter(|t| t.as_str() == expected_type) - .map(|_| bytes) - }) + out +} + +/// Finds the one envelope of `expected_type` among `envelopes` and +/// returns its payload bytes; panics if not exactly one is found. +fn extract_envelope_payload(envelopes: &[SketchEnvelope], expected_type: &str) -> Vec { + let expected = match expected_type { + "DDSketch" => SketchType::DDSketch, + "KLLSketch" => SketchType::KLLSketch, + "HLLSketch" => SketchType::HLLSketch, + "CountSketch" => SketchType::CountSketch, + "CountMinSketch" => SketchType::CountMinSketch, + other => panic!("unknown sketch type in test helper: {other}"), + }; + let mut matched: Vec> = envelopes + .iter() + .filter(|e| e.sketch_type == expected && !e.payload.is_empty()) + .map(|e| e.payload.clone()) .collect(); assert_eq!( matched.len(), @@ -184,39 +156,15 @@ fn extract_envelope_payload(records: &OtapMetricRecords, expected_type: &str) -> matched.remove(0) } -/// Assert the metrics-side schema of the emitted records does NOT -/// carry any `_asap_*` top-level columns — this is the Strategy-B -/// attribute-lift contract. OTAP's strict validator -/// (`crates/pdata/src/schema/payloads.rs::check_match`) rejects -/// extension columns, so the lift step on emit must remove them -/// from the metrics batch. -fn assert_no_strategy_b_top_level_columns(records: &OtapMetricRecords) { - for name in [ - ATTR_ENVELOPE, - ATTR_SKETCH_TYPE, - ATTR_AGG_ID, - ATTR_SCHEMA_VERSION, - ATTR_WINDOW_START_MS, - ATTR_WINDOW_END_MS, - ATTR_ENCODING, - ] { - assert!( - records.metrics.column_by_name(name).is_none(), - "metrics batch must NOT carry top-level column {name}" - ); - } - // The lift step adds parent_id; sanity-check it is present. - assert!(records.metrics.column_by_name(PARENT_ID_COL).is_some()); -} - /// Run a full `Start → N inputs → Shutdown → drain` cycle for one -/// sketch type. Returns the emitted records batch (post-lift) so -/// the per-sketch test can introspect the envelope payload. +/// sketch type. Returns the emitted `SketchStreamBatch`es so the +/// per-sketch test can decode them (via [`decode_all`]) and introspect +/// the envelope payload. async fn run_lifecycle( sketch_type: &str, metric: &str, inputs: &[(f64, &str)], -) -> Vec { +) -> Vec { let cfg = PluginConfig { sketch_type: sketch_type.into(), // Long enough that the natural ticker doesn't fire during @@ -257,9 +205,8 @@ async fn lifecycle_ddsketch_emits_envelope_with_correct_sketch_type() { ) .await; assert!(!records.is_empty(), "no records emitted on drain"); - let last = records.last().expect("at least one batch"); - assert_no_strategy_b_top_level_columns(last); - let payload = extract_envelope_payload(last, "DDSketch"); + let envelopes = decode_all(&records); + let payload = extract_envelope_payload(&envelopes, "DDSketch"); assert!(!payload.is_empty(), "DDSketch payload must not be empty"); } @@ -271,9 +218,8 @@ async fn lifecycle_kll_emits_envelope_with_correct_sketch_type() { &[(10.0, "h1"), (20.0, "h1"), (30.0, "h1")], ) .await; - let last = records.last().expect("at least one batch"); - assert_no_strategy_b_top_level_columns(last); - let payload = extract_envelope_payload(last, "KLLSketch"); + let envelopes = decode_all(&records); + let payload = extract_envelope_payload(&envelopes, "KLLSketch"); assert!(!payload.is_empty()); } @@ -289,9 +235,8 @@ async fn lifecycle_hll_emits_envelope_with_correct_sketch_type() { &[(1.0, "h1"), (2.0, "h1"), (3.0, "h1"), (4.0, "h1")], ) .await; - let last = records.last().expect("at least one batch"); - assert_no_strategy_b_top_level_columns(last); - let payload = extract_envelope_payload(last, "HLLSketch"); + let envelopes = decode_all(&records); + let payload = extract_envelope_payload(&envelopes, "HLLSketch"); assert!(!payload.is_empty()); } @@ -306,9 +251,8 @@ async fn lifecycle_countsketch_emits_envelope_with_correct_sketch_type() { &[(1.0, "h1"), (1.0, "h1"), (1.0, "h1"), (1.0, "h1")], ) .await; - let last = records.last().expect("at least one batch"); - assert_no_strategy_b_top_level_columns(last); - let payload = extract_envelope_payload(last, "CountSketch"); + let envelopes = decode_all(&records); + let payload = extract_envelope_payload(&envelopes, "CountSketch"); assert!(!payload.is_empty()); } @@ -326,9 +270,8 @@ async fn lifecycle_countminsketch_emits_envelope_with_correct_sketch_type() { &[(1.0, "h1"), (1.0, "h1"), (1.0, "h1"), (1.0, "h1")], ) .await; - let last = records.last().expect("at least one batch"); - assert_no_strategy_b_top_level_columns(last); - let payload = extract_envelope_payload(last, "CountMinSketch"); + let envelopes = decode_all(&records); + let payload = extract_envelope_payload(&envelopes, "CountMinSketch"); assert!( !payload.is_empty(), "CountMinSketch payload must not be empty" @@ -347,8 +290,9 @@ async fn drain_flushes_in_flight_observations_before_window_boundary() { &[(1.0, "h1"), (2.0, "h1"), (3.0, "h1"), (4.0, "h1")], ) .await; - let last = records.last().expect("drain must emit at least one batch"); - let payload = extract_envelope_payload(last, "DDSketch"); + assert!(!records.is_empty(), "drain must emit at least one batch"); + let envelopes = decode_all(&records); + let payload = extract_envelope_payload(&envelopes, "DDSketch"); assert!(!payload.is_empty()); } @@ -452,17 +396,13 @@ async fn control_channel_plan_change_acks_after_apply() { handle.shutdown().await.expect("shutdown"); let records = drain_emit(&mut emit_rx).await; - let last = records.last().expect("drain emit"); - // Verify the metric_name column reflects the applied plan. - let metric_col = last - .metrics - .column_by_name(COLUMN_METRIC) - .expect("metric col") - .as_any() - .downcast_ref::() - .expect("Utf8"); + assert!(!records.is_empty(), "drain emit"); + // Verify every decoded envelope's metric_name reflects the applied + // plan (sourced from DICTIONARY.metric, not repeated per RECORD). + let envelopes = decode_all(&records); + assert!(!envelopes.is_empty(), "drain emit"); assert!( - (0..metric_col.len()).all(|i| !metric_col.is_null(i) && metric_col.value(i) == "after"), + envelopes.iter().all(|e| e.metric_name == "after"), "post-plan-change emit should carry metric_name=after" ); }