diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs index 05b34353..c2644f4d 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs @@ -273,10 +273,12 @@ pub enum AggregationType { HydraKLL, CountMinSketch, CountMinSketchWithHeap, + CountSketch, // ---------- cardinality / set tracking ---------- SetAggregator, DeltaSetAggregator, HLL, + DDSketch, // ---------- legacy config wrapper names ---------- SingleSubpopulation, MultipleSubpopulation, @@ -295,9 +297,11 @@ impl AggregationType { AggregationType::HydraKLL => "HydraKLL", AggregationType::CountMinSketch => "CountMinSketch", AggregationType::CountMinSketchWithHeap => "CountMinSketchWithHeap", + AggregationType::CountSketch => "CountSketch", AggregationType::SetAggregator => "SetAggregator", AggregationType::DeltaSetAggregator => "DeltaSetAggregator", AggregationType::HLL => "HLL", + AggregationType::DDSketch => "DDSketch", AggregationType::SingleSubpopulation => "SingleSubpopulation", AggregationType::MultipleSubpopulation => "MultipleSubpopulation", } @@ -313,6 +317,7 @@ impl AggregationType { | AggregationType::MultipleMinMax | AggregationType::CountMinSketch | AggregationType::CountMinSketchWithHeap + | AggregationType::CountSketch | AggregationType::HydraKLL ) } @@ -326,6 +331,7 @@ impl AggregationType { | AggregationType::MultipleIncrease | AggregationType::CountMinSketch | AggregationType::CountMinSketchWithHeap + | AggregationType::CountSketch ) } @@ -360,9 +366,11 @@ impl FromStr for AggregationType { "HydraKLL" => Ok(AggregationType::HydraKLL), "CountMinSketch" => Ok(AggregationType::CountMinSketch), "CountMinSketchWithHeap" => Ok(AggregationType::CountMinSketchWithHeap), + "CountSketch" => Ok(AggregationType::CountSketch), "SetAggregator" => Ok(AggregationType::SetAggregator), "DeltaSetAggregator" => Ok(AggregationType::DeltaSetAggregator), "HLL" | "HyperLogLog" => Ok(AggregationType::HLL), + "DDSketch" | "DdSketch" => Ok(AggregationType::DDSketch), "SingleSubpopulation" => Ok(AggregationType::SingleSubpopulation), "MultipleSubpopulation" => Ok(AggregationType::MultipleSubpopulation), // Legacy accumulator-suffixed aliases @@ -384,6 +392,9 @@ impl FromStr for AggregationType { Ok(AggregationType::CountMinSketch) } "CountMinSketchWithHeapAccumulator" => Ok(AggregationType::CountMinSketchWithHeap), + "CountSketchAccumulator" | "CS" | "cs" | "count_sketch" => { + Ok(AggregationType::CountSketch) + } "SetAggregatorAccumulator" => Ok(AggregationType::SetAggregator), "DeltaSetAggregatorAccumulator" => Ok(AggregationType::DeltaSetAggregator), _ => Err(format!("Unknown aggregation type: '{s}'")), diff --git a/asap-common/sketch-core/src/count_sketch.rs b/asap-common/sketch-core/src/count_sketch.rs new file mode 100644 index 00000000..ea2eb314 --- /dev/null +++ b/asap-common/sketch-core/src/count_sketch.rs @@ -0,0 +1,165 @@ +//! Count Sketch (a.k.a. Count-Min-style signed-counter sketch) — +//! element-wise mergeable frequency estimator. +//! +//! Parallel to `count_min::CountMinSketch` but with **signed** counters, +//! matching the `asap_sketchlib::proto::sketchlib::CountSketchState` wire +//! format that DataCollector's `countsketchprocessor` emits via the +//! modified OTLP `Metric.data = CountSketch{…}` variant. +//! +//! This is the minimal surface needed for PR C-CountSketch in the +//! modified-OTLP hot path: construct from a decoded proto state, merge +//! element-wise with another sketch, emit the matrix for queries and +//! serialization. The richer query semantics of Count Sketch (median- +//! of-estimators heavy-hitter tracking, `TopKState` integration, etc.) +//! are intentionally deferred to a follow-up — the wire format already +//! carries the matrix losslessly, so the merge/store round-trip works +//! with just a matrix today. + +use serde::{Deserialize, Serialize}; + +/// Minimal Count Sketch state — a flat `rows × cols` matrix of signed +/// counts. Element-wise mergeable (sum over aligned cells). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CountSketch { + pub row_num: usize, + pub col_num: usize, + /// Row-major matrix of signed counts. `matrix[r][c]` is the value of + /// hash row `r`, column `c`. + pub matrix: Vec>, +} + +impl CountSketch { + /// Construct an all-zero sketch with the given dimensions. + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + row_num, + col_num, + matrix: vec![vec![0.0; col_num]; row_num], + } + } + + /// Construct from a pre-built matrix (used by the modified-OTLP + /// proto-decode path). + pub fn from_legacy_matrix(matrix: Vec>, row_num: usize, col_num: usize) -> Self { + debug_assert_eq!(matrix.len(), row_num, "row count mismatch"); + debug_assert!( + matrix.iter().all(|r| r.len() == col_num), + "column count mismatch in at least one row" + ); + Self { + row_num, + col_num, + matrix, + } + } + + /// Borrow the inner matrix. + pub fn sketch(&self) -> &Vec> { + &self.matrix + } + + /// Merge one other sketch into self via element-wise addition. Both + /// operands must have identical dimensions. + pub fn merge( + &mut self, + other: &CountSketch, + ) -> Result<(), Box> { + if self.row_num != other.row_num || self.col_num != other.col_num { + return Err(format!( + "CountSketch dimension mismatch: self={}x{}, other={}x{}", + self.row_num, self.col_num, other.row_num, other.col_num + ) + .into()); + } + for r in 0..self.row_num { + for c in 0..self.col_num { + self.matrix[r][c] += other.matrix[r][c]; + } + } + Ok(()) + } + + /// Merge a slice of references into a single new sketch. All inputs + /// must share the same dimensions; returns `Err` on mismatch or an + /// empty input. + pub fn merge_refs( + inputs: &[&CountSketch], + ) -> Result> { + let first = inputs + .first() + .ok_or("CountSketch::merge_refs called with empty input")?; + let mut merged = CountSketch::new(first.row_num, first.col_num); + for cs in inputs { + merged.merge(cs)?; + } + Ok(merged) + } + + /// Serialize to MessagePack bytes (used by the legacy Arroyo path + /// and by PR I's `_ENCODING_MSGPACK` variant when that lands). + pub fn serialize_msgpack(&self) -> Vec { + rmp_serde::to_vec(self).unwrap_or_default() + } + + /// Deserialize from MessagePack bytes. + pub fn deserialize_msgpack( + buffer: &[u8], + ) -> Result> { + Ok(rmp_serde::from_slice(buffer)?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_empty() { + let cs = CountSketch::new(2, 3); + assert_eq!(cs.row_num, 2); + assert_eq!(cs.col_num, 3); + assert_eq!(cs.sketch(), &vec![vec![0.0, 0.0, 0.0], vec![0.0, 0.0, 0.0]]); + } + + #[test] + fn test_from_legacy_matrix() { + let m = vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]]; + let cs = CountSketch::from_legacy_matrix(m.clone(), 2, 3); + assert_eq!(cs.sketch(), &m); + } + + #[test] + fn test_merge_element_wise() { + let mut a = CountSketch::from_legacy_matrix(vec![vec![1.0, 2.0], vec![3.0, 4.0]], 2, 2); + let b = CountSketch::from_legacy_matrix(vec![vec![-1.0, -2.0], vec![-3.0, -4.0]], 2, 2); + a.merge(&b).unwrap(); + assert_eq!(a.sketch(), &vec![vec![0.0, 0.0], vec![0.0, 0.0]]); + } + + #[test] + fn test_merge_dimension_mismatch() { + let mut a = CountSketch::new(2, 3); + let b = CountSketch::new(3, 3); + assert!(a.merge(&b).is_err()); + } + + #[test] + fn test_merge_refs() { + let a = CountSketch::from_legacy_matrix(vec![vec![1.0, 2.0]], 1, 2); + let b = CountSketch::from_legacy_matrix(vec![vec![3.0, 4.0]], 1, 2); + let c = CountSketch::from_legacy_matrix(vec![vec![5.0, 6.0]], 1, 2); + let merged = CountSketch::merge_refs(&[&a, &b, &c]).unwrap(); + assert_eq!(merged.sketch(), &vec![vec![9.0, 12.0]]); + } + + #[test] + fn test_msgpack_round_trip() { + let original = + CountSketch::from_legacy_matrix(vec![vec![1.5, -2.5], vec![3.5, -4.5]], 2, 2); + let bytes = original.serialize_msgpack(); + let decoded = CountSketch::deserialize_msgpack(&bytes).unwrap(); + assert_eq!(decoded.sketch(), original.sketch()); + assert_eq!(decoded.row_num, original.row_num); + assert_eq!(decoded.col_num, original.col_num); + } +} diff --git a/asap-common/sketch-core/src/dd_sketch.rs b/asap-common/sketch-core/src/dd_sketch.rs new file mode 100644 index 00000000..789367a1 --- /dev/null +++ b/asap-common/sketch-core/src/dd_sketch.rs @@ -0,0 +1,225 @@ +//! DDSketch — log-bucketed quantile sketch, mergeable by store-index alignment. +//! +//! Parallel to `count_sketch::CountSketch`: the minimum viable surface +//! needed for the modified-OTLP `Metric.data = DDSketch{…}` hot path +//! (PR C-CountSketch follow-up). Holds the bucket counts, their +//! absolute-index base offset, and the aggregate `{count, sum, min, max}`. +//! +//! Merge semantics: two sketches with the same relative-accuracy +//! parameter `alpha` are merged by aligning bucket arrays along their +//! `store_offset` and summing counts element-wise, with `min`/`max` +//! combined via min/max and `count`/`sum` added. +//! +//! The wire format is the protobuf-encoded +//! `asap_sketchlib::proto::sketchlib::DDSketchState` emitted by +//! DataCollector's `ddsketchprocessor`. Quantile estimation against +//! stored data is intentionally deferred — queries currently return +//! a placeholder error and fall through to the §5.2 fallback. + +use serde::{Deserialize, Serialize}; + +/// Minimal DDSketch state — bucket counts + alpha + aggregates. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DdSketch { + /// Relative accuracy parameter; must satisfy `0 < alpha < 1`. + pub alpha: f64, + /// Bucket counts in absolute-index order. The absolute index of + /// `store_counts[i]` is `i + store_offset`. + pub store_counts: Vec, + /// Absolute bucket index corresponding to `store_counts[0]`. May + /// be negative. + pub store_offset: i32, + pub count: u64, + pub sum: f64, + pub min: f64, + pub max: f64, +} + +impl DdSketch { + /// Construct an empty sketch. + pub fn new(alpha: f64) -> Self { + Self { + alpha, + store_counts: Vec::new(), + store_offset: 0, + count: 0, + sum: 0.0, + min: f64::INFINITY, + max: f64::NEG_INFINITY, + } + } + + /// Construct from the decoded wire fields. + #[allow(clippy::too_many_arguments)] + pub fn from_raw( + alpha: f64, + store_counts: Vec, + store_offset: i32, + count: u64, + sum: f64, + min: f64, + max: f64, + ) -> Self { + Self { + alpha, + store_counts, + store_offset, + count, + sum, + min, + max, + } + } + + /// Merge one other sketch into self by aligning bucket arrays on + /// absolute indices. Both operands must share the same `alpha`. + pub fn merge( + &mut self, + other: &DdSketch, + ) -> Result<(), Box> { + if (self.alpha - other.alpha).abs() > f64::EPSILON { + return Err(format!( + "DdSketch alpha mismatch: self={}, other={}", + self.alpha, other.alpha + ) + .into()); + } + + if other.store_counts.is_empty() { + self.count += other.count; + self.sum += other.sum; + if other.min < self.min { + self.min = other.min; + } + if other.max > self.max { + self.max = other.max; + } + return Ok(()); + } + if self.store_counts.is_empty() { + self.store_counts = other.store_counts.clone(); + self.store_offset = other.store_offset; + } else { + let self_start = self.store_offset as i64; + let self_end = self_start + self.store_counts.len() as i64; + let other_start = other.store_offset as i64; + let other_end = other_start + other.store_counts.len() as i64; + let new_start = self_start.min(other_start); + let new_end = self_end.max(other_end); + let new_len = (new_end - new_start) as usize; + let mut merged = vec![0u64; new_len]; + for (i, c) in self.store_counts.iter().enumerate() { + let idx = (self_start + i as i64 - new_start) as usize; + merged[idx] = merged[idx].saturating_add(*c); + } + for (i, c) in other.store_counts.iter().enumerate() { + let idx = (other_start + i as i64 - new_start) as usize; + merged[idx] = merged[idx].saturating_add(*c); + } + self.store_counts = merged; + self.store_offset = new_start as i32; + } + self.count += other.count; + self.sum += other.sum; + if other.min < self.min { + self.min = other.min; + } + if other.max > self.max { + self.max = other.max; + } + Ok(()) + } + + /// Merge a slice of references into a single new sketch. Returns + /// `Err` on alpha mismatch or an empty input. + pub fn merge_refs( + inputs: &[&DdSketch], + ) -> Result> { + let first = inputs + .first() + .ok_or("DdSketch::merge_refs called with empty input")?; + let mut merged = DdSketch::new(first.alpha); + for d in inputs { + merged.merge(d)?; + } + Ok(merged) + } + + /// Serialize to MessagePack bytes. + pub fn serialize_msgpack(&self) -> Vec { + rmp_serde::to_vec(self).unwrap_or_default() + } + + /// Deserialize from MessagePack bytes. + pub fn deserialize_msgpack( + buffer: &[u8], + ) -> Result> { + Ok(rmp_serde::from_slice(buffer)?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_empty() { + let d = DdSketch::new(0.01); + assert_eq!(d.count, 0); + assert!(d.store_counts.is_empty()); + assert_eq!(d.min, f64::INFINITY); + assert_eq!(d.max, f64::NEG_INFINITY); + } + + #[test] + fn test_merge_aligned_same_offset() { + let mut a = DdSketch::from_raw(0.01, vec![1, 2, 3], -1, 6, 30.0, 1.0, 5.0); + let b = DdSketch::from_raw(0.01, vec![10, 20, 30], -1, 60, 300.0, 0.5, 6.0); + a.merge(&b).unwrap(); + assert_eq!(a.store_counts, vec![11, 22, 33]); + assert_eq!(a.store_offset, -1); + assert_eq!(a.count, 66); + assert_eq!(a.sum, 330.0); + assert_eq!(a.min, 0.5); + assert_eq!(a.max, 6.0); + } + + #[test] + fn test_merge_overlapping_offsets() { + // a covers indices [-1, 0, 1]; b covers indices [0, 1, 2] + let mut a = DdSketch::from_raw(0.01, vec![1, 1, 1], -1, 3, 3.0, 1.0, 3.0); + let b = DdSketch::from_raw(0.01, vec![10, 10, 10], 0, 30, 30.0, 1.0, 3.0); + a.merge(&b).unwrap(); + // Merged window is [-1, 0, 1, 2] → [1, 11, 11, 10] + assert_eq!(a.store_counts, vec![1, 11, 11, 10]); + assert_eq!(a.store_offset, -1); + assert_eq!(a.count, 33); + } + + #[test] + fn test_merge_disjoint_offsets() { + let mut a = DdSketch::from_raw(0.01, vec![1, 2], 0, 3, 3.0, 1.0, 2.0); + let b = DdSketch::from_raw(0.01, vec![3, 4], 5, 7, 7.0, 5.0, 6.0); + a.merge(&b).unwrap(); + // Window [0..7) → [1,2,0,0,0,3,4] + assert_eq!(a.store_counts, vec![1, 2, 0, 0, 0, 3, 4]); + assert_eq!(a.store_offset, 0); + } + + #[test] + fn test_merge_alpha_mismatch() { + let mut a = DdSketch::new(0.01); + let b = DdSketch::new(0.02); + assert!(a.merge(&b).is_err()); + } + + #[test] + fn test_msgpack_round_trip() { + let original = DdSketch::from_raw(0.01, vec![1, 2, 3], -2, 6, 30.0, 1.0, 5.0); + let bytes = original.serialize_msgpack(); + let decoded = DdSketch::deserialize_msgpack(&bytes).unwrap(); + assert_eq!(decoded.store_counts, original.store_counts); + assert_eq!(decoded.store_offset, original.store_offset); + assert_eq!(decoded.count, original.count); + } +} diff --git a/asap-common/sketch-core/src/hll_sketch.rs b/asap-common/sketch-core/src/hll_sketch.rs new file mode 100644 index 00000000..5e687ba9 --- /dev/null +++ b/asap-common/sketch-core/src/hll_sketch.rs @@ -0,0 +1,212 @@ +//! HyperLogLog sketch — register-wise mergeable cardinality estimator. +//! +//! Parallel to `count_sketch::CountSketch`: the minimum viable surface +//! needed for the modified-OTLP `Metric.data = HLLSketch{…}` hot path +//! (PR C-CountSketch follow-up). Wraps a flat `Vec` of register +//! values (length = `2^precision`) and merges element-wise by taking +//! the maximum across aligned registers, which is the standard HLL +//! merge semantics. +//! +//! The wire format is the protobuf-encoded +//! `asap_sketchlib::proto::sketchlib::HyperLogLogState` emitted by +//! DataCollector's `hllprocessor`. This type carries the register +//! bytes and the variant/precision metadata losslessly, so the +//! merge + store round-trip works end-to-end. Cardinality estimation +//! against stored HLL data is intentionally deferred to a follow-up +//! — queries currently return a placeholder error and fall through +//! to the §5.2 fallback. + +use serde::{Deserialize, Serialize}; + +/// HLL estimator variant. Mirrors `asap_sketchlib::proto::sketchlib::HllVariant` +/// so the proto round-trip preserves the algorithm identity — the three +/// variants are not mutually compatible on register contents. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum HllVariant { + Unspecified, + Regular, + Datafusion, + Hip, +} + +/// Minimal HLL state — registers + variant + precision. Register-wise +/// mergeable (max over aligned cells). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HllSketch { + pub variant: HllVariant, + pub precision: u32, + /// Flat register array, length = `2^precision`. + pub registers: Vec, + /// HIP accumulator components — populated only when `variant == Hip`. + pub hip_kxq0: f64, + pub hip_kxq1: f64, + pub hip_est: f64, +} + +impl HllSketch { + /// Construct an empty sketch at the given precision. + pub fn new(variant: HllVariant, precision: u32) -> Self { + let n = 1usize << precision; + Self { + variant, + precision, + registers: vec![0u8; n], + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + } + } + + /// Construct from pre-built register bytes (used by the modified-OTLP + /// proto-decode path). + pub fn from_raw( + variant: HllVariant, + precision: u32, + registers: Vec, + hip_kxq0: f64, + hip_kxq1: f64, + hip_est: f64, + ) -> Self { + Self { + variant, + precision, + registers, + hip_kxq0, + hip_kxq1, + hip_est, + } + } + + /// Merge one other sketch into self via register-wise max. Both + /// operands must have identical variant and precision. + pub fn merge( + &mut self, + other: &HllSketch, + ) -> Result<(), Box> { + if self.variant != other.variant { + return Err(format!( + "HllSketch variant mismatch: self={:?}, other={:?}", + self.variant, other.variant + ) + .into()); + } + if self.precision != other.precision { + return Err(format!( + "HllSketch precision mismatch: self={}, other={}", + self.precision, other.precision + ) + .into()); + } + if self.registers.len() != other.registers.len() { + return Err(format!( + "HllSketch register-length mismatch: self={}, other={}", + self.registers.len(), + other.registers.len() + ) + .into()); + } + for (s, o) in self.registers.iter_mut().zip(other.registers.iter()) { + if *o > *s { + *s = *o; + } + } + // HIP accumulators add on merge (each source carried its own + // running estimate; merged state inherits the combined + // components). + if self.variant == HllVariant::Hip { + self.hip_kxq0 += other.hip_kxq0; + self.hip_kxq1 += other.hip_kxq1; + self.hip_est += other.hip_est; + } + Ok(()) + } + + /// Merge a slice of references into a single new sketch. All inputs + /// must share the same variant and precision; returns `Err` on + /// mismatch or an empty input. + pub fn merge_refs( + inputs: &[&HllSketch], + ) -> Result> { + let first = inputs + .first() + .ok_or("HllSketch::merge_refs called with empty input")?; + let mut merged = HllSketch::new(first.variant, first.precision); + for hll in inputs { + merged.merge(hll)?; + } + Ok(merged) + } + + /// Serialize to MessagePack bytes (used by the legacy Arroyo path + /// and by PR I's `_ENCODING_MSGPACK` variant when that lands). + pub fn serialize_msgpack(&self) -> Vec { + rmp_serde::to_vec(self).unwrap_or_default() + } + + /// Deserialize from MessagePack bytes. + pub fn deserialize_msgpack( + buffer: &[u8], + ) -> Result> { + Ok(rmp_serde::from_slice(buffer)?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_empty() { + let h = HllSketch::new(HllVariant::Regular, 4); + assert_eq!(h.registers.len(), 16); + assert!(h.registers.iter().all(|&r| r == 0)); + } + + #[test] + fn test_merge_register_wise_max() { + let mut a = HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0); + let b = HllSketch::from_raw(HllVariant::Regular, 2, vec![4, 2, 6, 0], 0.0, 0.0, 0.0); + a.merge(&b).unwrap(); + assert_eq!(a.registers, vec![4, 5, 6, 7]); + } + + #[test] + fn test_merge_variant_mismatch() { + let mut a = HllSketch::new(HllVariant::Regular, 4); + let b = HllSketch::new(HllVariant::Datafusion, 4); + assert!(a.merge(&b).is_err()); + } + + #[test] + fn test_merge_precision_mismatch() { + let mut a = HllSketch::new(HllVariant::Regular, 4); + let b = HllSketch::new(HllVariant::Regular, 5); + assert!(a.merge(&b).is_err()); + } + + #[test] + fn test_merge_refs() { + let a = HllSketch::from_raw(HllVariant::Regular, 1, vec![1, 0], 0.0, 0.0, 0.0); + let b = HllSketch::from_raw(HllVariant::Regular, 1, vec![0, 3], 0.0, 0.0, 0.0); + let c = HllSketch::from_raw(HllVariant::Regular, 1, vec![2, 2], 0.0, 0.0, 0.0); + let merged = HllSketch::merge_refs(&[&a, &b, &c]).unwrap(); + assert_eq!(merged.registers, vec![2, 3]); + } + + #[test] + fn test_msgpack_round_trip() { + let original = HllSketch::from_raw( + HllVariant::Hip, + 3, + vec![0, 1, 2, 3, 4, 5, 6, 7], + 1.0, + 2.0, + 3.0, + ); + let bytes = original.serialize_msgpack(); + let decoded = HllSketch::deserialize_msgpack(&bytes).unwrap(); + assert_eq!(decoded.registers, original.registers); + assert_eq!(decoded.precision, original.precision); + assert_eq!(decoded.hip_kxq0, 1.0); + } +} diff --git a/asap-common/sketch-core/src/lib.rs b/asap-common/sketch-core/src/lib.rs index 3ddd32b7..3cfda08d 100644 --- a/asap-common/sketch-core/src/lib.rs +++ b/asap-common/sketch-core/src/lib.rs @@ -9,7 +9,10 @@ pub mod count_min; pub mod count_min_sketchlib; pub mod count_min_with_heap; pub mod count_min_with_heap_sketchlib; +pub mod count_sketch; +pub mod dd_sketch; pub mod delta_set_aggregator; +pub mod hll_sketch; pub mod hydra_kll; pub mod kll; pub mod kll_sketchlib; diff --git a/asap-query-engine/src/drivers/ingest/otel.rs b/asap-query-engine/src/drivers/ingest/otel.rs index 072c9efe..cf8a42c7 100644 --- a/asap-query-engine/src/drivers/ingest/otel.rs +++ b/asap-query-engine/src/drivers/ingest/otel.rs @@ -696,7 +696,10 @@ fn decode_modified_otlp_sketch_bytes( encoding: i32, bytes: &[u8], ) -> Result, Box> { - use crate::precompute_operators::CountMinSketchAccumulator; + use crate::precompute_operators::{ + CountMinSketchAccumulator, CountSketchAccumulator, DDSketchAccumulator, + DatasketchesKLLAccumulator, HllSketchAccumulator, + }; // The encoding value is the raw i32 from the proto enum. We only // accept ENCODING_PROTO (= 1) for now; ENCODING_PROTO_DELTA (= 2) @@ -719,12 +722,21 @@ fn decode_modified_otlp_sketch_bytes( let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; Ok(Box::new(acc)) } - SketchKind::Kll | SketchKind::DdSketch | SketchKind::CountSketch | SketchKind::Hll => { - Err(format!( - "modified-OTLP sketch decoder for {kind:?} not yet implemented \ - (tracked in PR C, task #8)" - ) - .into()) + SketchKind::CountSketch => { + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; + Ok(Box::new(acc)) + } + SketchKind::Kll => { + let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(bytes)?; + Ok(Box::new(acc)) + } + SketchKind::DdSketch => { + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; + Ok(Box::new(acc)) + } + SketchKind::Hll => { + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; + Ok(Box::new(acc)) } } } diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index 2d693cbc..8d8cf9e1 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -168,7 +168,6 @@ struct Args { // --persistence-* flags as the config. Forces LockStrategy::PerKey // regardless of --lock-strategy; the Global variant is // intentionally left in-memory-only. - /// Enable the disk-backed persistence layer for SimpleMapStore #[arg(long)] persistence_enabled: bool, diff --git a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs new file mode 100644 index 00000000..f05936f2 --- /dev/null +++ b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs @@ -0,0 +1,304 @@ +//! Count Sketch accumulator — wraps `sketch_core::count_sketch::CountSketch`. +//! +//! This is the concrete accumulator reached from the modified-OTLP +//! `Metric.data = CountSketch{…}` hot path (PR C-CountSketch). Its +//! inner matrix is the same shape as `CountMinSketchAccumulator`'s +//! but with signed counts and no per-row heap tracking. +//! +//! Minimum viable surface: +//! - `AggregateCore` impl for precompute-engine worker merge +//! - `SerializableToSink` impl for store write-out +//! - `from_sketchlib_proto_bytes(buf)` — decoder for the modified +//! OTLP `CountSketchDataPoint.sketch` bytes (prost-encoded +//! `asap_sketchlib::proto::sketchlib::CountSketchState`). Mirrors +//! the CountMin decoder in PR B. +//! +//! Query semantics (median-of-estimators heavy-hitter tracking, +//! `TopKState` integration) are intentionally deferred — queries +//! against stored CountSketch data return a placeholder error today. +//! The wire format carries the matrix losslessly, so the merge + store +//! round-trip works end-to-end without that richer query surface. + +use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use serde_json::Value; +use sketch_core::count_sketch::CountSketch; +use std::collections::HashMap; + +/// Count Sketch accumulator — inner matrix of signed counts. +#[derive(Debug, Clone)] +pub struct CountSketchAccumulator { + pub inner: CountSketch, +} + +impl CountSketchAccumulator { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + inner: CountSketch::new(row_num, col_num), + } + } + + /// Decode from the modified OTLP wire format's + /// `CountSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::CountSketchState` message + /// that DataCollector's `countsketchprocessor` emits when + /// `encoding = COUNT_SKETCH_ENCODING_PROTO`. + /// + /// Mirrors `CountMinSketchAccumulator::from_sketchlib_proto_bytes` + /// but on the signed-counter `CountSketchState`. The resulting + /// accumulator is constructed via + /// `CountSketch::from_legacy_matrix` after reshaping the flat + /// `counts_int` / `counts_float` field into a `Vec>`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{CountSketchState, CounterType}; + use prost::Message; + + let state = CountSketchState::decode(buffer) + .map_err(|e| format!("decode CountSketchState: {e}"))?; + let rows = state.rows as usize; + let cols = state.cols as usize; + if rows == 0 || cols == 0 { + return Err( + format!("CountSketchState has zero dims (rows={rows}, cols={cols})").into(), + ); + } + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + format!( + "CountSketchState has unknown counter_type tag {}", + state.counter_type + ) + })?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountSketchState counts_int has {} entries, expected rows*cols = {}", + state.counts_int.len(), + expected_len + ) + .into()); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountSketchState counts_float has {} entries, expected rows*cols = {}", + state.counts_float.len(), + expected_len + ) + .into()); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountSketchState counter_type {other:?} not yet supported \ + (INT128 stores interleaved hi/lo pairs; will be added when needed)" + ) + .into()); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(Self { + inner: CountSketch::from_legacy_matrix(matrix, rows, cols), + }) + } +} + +impl SerializableToSink for CountSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "row_num": self.inner.row_num, + "col_num": self.inner.col_num, + "sketch": self.inner.sketch(), + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.serialize_msgpack() + } +} + +impl AggregateCore for CountSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountSketchAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_cs = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountSketchAccumulator")?; + + let merged_inner = CountSketch::merge_refs(&[&self.inner, &other_cs.inner])?; + Ok(Box::new(Self { + inner: merged_inner, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountSketch + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + _statistic: promql_utilities::query_logics::enums::Statistic, + _key: &Option, + _query_kwargs: &HashMap, + ) -> Result> { + // Query semantics (median-of-estimators heavy-hitter, TopKState) + // are deferred to a follow-up. The matrix round-trip through the + // modified-OTLP hot path already works end-to-end without this; + // queries against stored CountSketch data return a placeholder + // error and fall through to the §5.2 fallback. + Err( + "CountSketchAccumulator: query_statistic not yet implemented \ + (matrix round-trip works, but query semantics deferred; \ + tracked as a PR C-CountSketch follow-up)" + .into(), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode_state( + rows: u32, + cols: u32, + counter_type: i32, + counts_int: Vec, + counts_float: Vec, + ) -> Vec { + use asap_sketchlib::proto::sketchlib::CountSketchState; + use prost::Message; + let state = CountSketchState { + rows, + cols, + counter_type, + counts_int, + counts_float, + l2: Vec::new(), + topk: None, + }; + state.encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_int64() { + use asap_sketchlib::proto::sketchlib::CounterType; + // Signed 2x3 matrix: row 0 = [1,-2,3], row 1 = [-4,5,-6] + let bytes = encode_state( + 2, + 3, + CounterType::Int64 as i32, + vec![1, -2, 3, -4, 5, -6], + Vec::new(), + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]); + assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_float64() { + use asap_sketchlib::proto::sketchlib::CounterType; + let bytes = encode_state( + 2, + 2, + CounterType::Float64 as i32, + Vec::new(), + vec![1.5, -2.5, 3.5, -4.5], + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.5, -2.5]); + assert_eq!(matrix[1], vec![3.5, -4.5]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_dimension_mismatch() { + use asap_sketchlib::proto::sketchlib::CounterType; + // 2x3 declared but only 5 int entries + let bytes = encode_state( + 2, + 3, + CounterType::Int64 as i32, + vec![1, 2, 3, 4, 5], + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains("counts_int"), + "error should mention counts_int dim mismatch" + ); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_dims_rejected() { + use asap_sketchlib::proto::sketchlib::CountSketchState; + use prost::Message; + let state = CountSketchState::default(); + let bytes = state.encode_to_vec(); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("zero dims")); + } + + #[test] + fn test_aggregate_core_merge_matches_matrix_add() { + let a = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![1.0, -2.0], vec![3.0, -4.0]], 2, 2), + }; + let b = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![-1.0, 2.0], vec![-3.0, 4.0]], 2, 2), + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + let m = merged.inner.sketch(); + assert_eq!(m[0], vec![0.0, 0.0]); + assert_eq!(m[1], vec![0.0, 0.0]); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::precompute_operators::count_min_sketch_accumulator::CountMinSketchAccumulator; + let cs = CountSketchAccumulator::new(2, 3); + let cms = CountMinSketchAccumulator::new(2, 3); + let result = cs.merge_with(&cms); + assert!(result.is_err()); + } +} diff --git a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs index 528e3101..317403d2 100644 --- a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs @@ -70,6 +70,77 @@ impl DatasketchesKLLAccumulator { }) } + /// Decode from the modified OTLP wire format's + /// `KLLSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::KllState` message that + /// DataCollector's `kllprocessor` emits when + /// `encoding = KLL_SKETCH_ENCODING_PROTO`. + /// + /// ⚠ This is a **lossy statistical reconstruction**, not a + /// bit-identical round-trip: the `KllState` proto carries the + /// retained items in level order plus an explicit `levels[]` + /// boundary array, but sketch-core's `KllSketch` backend types + /// keep their level structure private. Rather than touch + /// upstream `asap_sketchlib` to add a typed-state constructor, + /// we build a fresh `DatasketchesKLLAccumulator` with the same + /// `k` and replay every retained item through `update()`. + /// Quantile estimates on the reconstructed sketch are + /// approximately equivalent to the source's — within KLL's + /// own rank-error bound, which is the same bound the source + /// already inherited — so Phase 1 hot-path queries that hit + /// the reconstructed sketch return answers the user would + /// already have accepted from the source. Bit-identical + /// reconstruction is tracked as a sketchlib upstream follow-up. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::KllState; + use prost::Message; + + let state = KllState::decode(buffer).map_err(|e| format!("decode KllState: {e}"))?; + if state.k < 8 { + return Err(format!("KllState.k must be >= 8 (got {})", state.k).into()); + } + if state.k > u16::MAX as u32 { + return Err(format!( + "KllState.k does not fit in u16 (got {}, max {})", + state.k, + u16::MAX + ) + .into()); + } + // Validate the levels[] boundary array if it is populated. The + // proto contract says `levels[0] == 0` and + // `levels[num_levels] == items.len()`. If the producer left + // levels empty (common when num_levels is zero), skip. + if !state.levels.is_empty() { + if state.levels.len() as u32 != state.num_levels + 1 { + return Err(format!( + "KllState levels length = {}, expected num_levels+1 = {}", + state.levels.len(), + state.num_levels + 1 + ) + .into()); + } + if state.levels[0] != 0 { + return Err(format!("KllState.levels[0] = {}, expected 0", state.levels[0]).into()); + } + if *state.levels.last().unwrap() as usize != state.items.len() { + return Err(format!( + "KllState.levels[{}] = {}, expected items.len() = {}", + state.num_levels, + state.levels.last().unwrap(), + state.items.len() + ) + .into()); + } + } + let k = state.k as u16; + let mut acc = Self::new(k); + for item in &state.items { + acc.update(*item); + } + Ok(acc) + } + /// Merge multiple accumulators efficiently without cloning all of them. pub fn merge_multiple( accumulators: &[Box], @@ -463,4 +534,83 @@ mod tests { let mixed_accs: Vec> = vec![Box::new(kll), Box::new(sum)]; assert!(DatasketchesKLLAccumulator::merge_multiple(&mixed_accs).is_err()); } + + #[test] + fn test_from_sketchlib_proto_bytes_reconstructs_quantiles() { + // Build a KllState with 64 items in level order; the decoder + // replays every item through `update()` so the reconstructed + // sketch is statistically equivalent — quantile estimates + // match the ground truth (sorted items) within KLL's own + // rank-error bound for k=200. + use asap_sketchlib::proto::sketchlib::KllState; + use prost::Message; + + let items: Vec = (0..64).map(|i| i as f64).collect(); + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 64], + items: items.clone(), + coin: None, + }; + let bytes = state.encode_to_vec(); + + let acc = + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.count(), 64); + // For 64 values 0..63, the true median is 31.5 and quantile + // error is ~1% × range = 0.63. KLL's own point query can + // legally be off by up to ε × N ~= 0.01 × 64 = 0.64. Allow a + // generous tolerance since the important invariant is "the + // decoded sketch is queryable and returns a sensible value". + let median = acc.get_quantile(0.5); + assert!( + (median - 31.5).abs() <= 10.0, + "reconstructed median {median} is outside tolerance of true median 31.5" + ); + let q01 = acc.get_quantile(0.01); + let q99 = acc.get_quantile(0.99); + assert!( + q01 <= q99, + "quantile monotonicity violated: q01={q01}, q99={q99}" + ); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_small_k() { + use asap_sketchlib::proto::sketchlib::KllState; + use prost::Message; + let state = KllState { + k: 4, // < minimum of 8 + m: 2, + num_levels: 0, + levels: Vec::new(), + items: Vec::new(), + coin: None, + }; + let bytes = state.encode_to_vec(); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("k must be >= 8")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_inconsistent_levels() { + use asap_sketchlib::proto::sketchlib::KllState; + use prost::Message; + // num_levels=1 but levels array has 3 entries instead of 2 + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 5, 10], + items: vec![1.0, 2.0, 3.0, 4.0, 5.0], + coin: None, + }; + let bytes = state.encode_to_vec(); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("levels length")); + } } diff --git a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs new file mode 100644 index 00000000..de8f6e24 --- /dev/null +++ b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs @@ -0,0 +1,209 @@ +//! DDSketch accumulator — wraps `sketch_core::dd_sketch::DdSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = DDSketch{…}` hot path (PR C-CountSketch follow-up). +//! Merge via bucket-index alignment on the inner sketch, serialize as +//! MessagePack for the sink, and decode from the sketchlib +//! `DDSketchState` proto. +//! +//! Query semantics (quantile estimation via log-bucket indices) are +//! intentionally deferred — the wire format carries the bucket counts, +//! offset, and aggregates losslessly, so the merge + store round-trip +//! works end-to-end without that richer query surface. + +use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use serde_json::Value; +use sketch_core::dd_sketch::DdSketch; +use std::collections::HashMap; + +/// DDSketch accumulator — inner log-bucketed sketch. +#[derive(Debug, Clone)] +pub struct DDSketchAccumulator { + pub inner: DdSketch, +} + +impl DDSketchAccumulator { + pub fn new(alpha: f64) -> Self { + Self { + inner: DdSketch::new(alpha), + } + } + + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::DDSketchState` message that + /// DataCollector's `ddsketchprocessor` emits when + /// `encoding = DD_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::DdSketchState; + use prost::Message; + + let state = + DdSketchState::decode(buffer).map_err(|e| format!("decode DDSketchState: {e}"))?; + if !(state.alpha > 0.0 && state.alpha < 1.0) { + return Err(format!( + "DDSketchState alpha {} out of range (expected 0 < alpha < 1)", + state.alpha + ) + .into()); + } + let inner = DdSketch::from_raw( + state.alpha, + state.store_counts.clone(), + state.store_offset, + state.count, + state.sum, + state.min, + state.max, + ); + Ok(Self { inner }) + } +} + +impl SerializableToSink for DDSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "alpha": self.inner.alpha, + "store_offset": self.inner.store_offset, + "bucket_count": self.inner.store_counts.len(), + "count": self.inner.count, + "sum": self.inner.sum, + "min": self.inner.min, + "max": self.inner.max, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.serialize_msgpack() + } +} + +impl AggregateCore for DDSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "DDSketchAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge DDSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_dd = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to DDSketchAccumulator")?; + let merged_inner = DdSketch::merge_refs(&[&self.inner, &other_dd.inner])?; + Ok(Box::new(Self { + inner: merged_inner, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::DDSketch + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + _statistic: promql_utilities::query_logics::enums::Statistic, + _key: &Option, + _query_kwargs: &HashMap, + ) -> Result> { + Err("DDSketchAccumulator: query_statistic not yet implemented \ + (bucket round-trip works, but quantile estimation deferred; \ + tracked as a PR C-CountSketch follow-up)" + .into()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode_state( + alpha: f64, + store_counts: Vec, + store_offset: i32, + count: u64, + sum: f64, + min: f64, + max: f64, + ) -> Vec { + use asap_sketchlib::proto::sketchlib::DdSketchState; + use prost::Message; + let state = DdSketchState { + alpha, + store_counts, + store_offset, + count, + sum, + min, + max, + }; + state.encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_round_trip() { + let bytes = encode_state(0.01, vec![1, 2, 3, 4], -2, 10, 50.0, 1.0, 4.0); + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.store_counts, vec![1, 2, 3, 4]); + assert_eq!(acc.inner.store_offset, -2); + assert_eq!(acc.inner.count, 10); + assert_eq!(acc.inner.sum, 50.0); + assert_eq!(acc.inner.min, 1.0); + assert_eq!(acc.inner.max, 4.0); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_invalid_alpha() { + let bytes = encode_state(0.0, vec![1], 0, 1, 1.0, 1.0, 1.0); + let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("alpha")); + } + + #[test] + fn test_aggregate_core_merge_aligns_buckets() { + let a = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 1, 1], -1, 3, 3.0, 1.0, 3.0), + }; + let b = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![10, 10, 10], 0, 30, 30.0, 1.0, 3.0), + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.inner.store_counts, vec![1, 11, 11, 10]); + assert_eq!(merged.inner.store_offset, -1); + assert_eq!(merged.inner.count, 33); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::precompute_operators::count_sketch_accumulator::CountSketchAccumulator; + let dd = DDSketchAccumulator::new(0.01); + let cs = CountSketchAccumulator::new(2, 3); + assert!(dd.merge_with(&cs).is_err()); + } +} diff --git a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs index 9efa72eb..859bf45f 100644 --- a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs @@ -247,8 +247,7 @@ impl AggregateCore for DeltaSetAggregatorAccumulator { fn approx_memory_bytes(&self) -> usize { // Two HashSets of KeyByLabelValues. const BYTES_PER_ENTRY: usize = 96; - std::mem::size_of::() - + (self.added.len() + self.removed.len()) * BYTES_PER_ENTRY + std::mem::size_of::() + (self.added.len() + self.removed.len()) * BYTES_PER_ENTRY } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs new file mode 100644 index 00000000..3647759b --- /dev/null +++ b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs @@ -0,0 +1,261 @@ +//! HLL accumulator — wraps `sketch_core::hll_sketch::HllSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = HLLSketch{…}` hot path (PR C-CountSketch follow-up). +//! Mirrors the CountSketch accumulator's shape: merge via register-wise +//! max on the inner sketch, serialize as MessagePack for the sink, and +//! decode from the sketchlib `HyperLogLogState` proto. +//! +//! Query semantics (cardinality estimation via the three HLL variants' +//! estimators) are intentionally deferred — the wire format carries the +//! registers + variant + HIP accumulators losslessly, so the merge + +//! store round-trip works end-to-end without that richer query surface. + +use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use serde_json::Value; +use sketch_core::hll_sketch::{HllSketch, HllVariant}; +use std::collections::HashMap; + +/// HLL accumulator — inner register array + variant metadata. +#[derive(Debug, Clone)] +pub struct HllSketchAccumulator { + pub inner: HllSketch, +} + +impl HllSketchAccumulator { + pub fn new(variant: HllVariant, precision: u32) -> Self { + Self { + inner: HllSketch::new(variant, precision), + } + } + + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::HyperLogLogState` message + /// that DataCollector's `hllprocessor` emits when + /// `encoding = HLL_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{HllVariant as ProtoVariant, HyperLogLogState}; + use prost::Message; + + let state = HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?; + if state.precision == 0 || state.precision > 20 { + return Err(format!( + "HyperLogLogState precision {} out of range (expected 1..=20)", + state.precision + ) + .into()); + } + let expected_len = 1usize << state.precision; + if state.registers.len() != expected_len { + return Err(format!( + "HyperLogLogState registers has {} bytes, expected 2^precision = {}", + state.registers.len(), + expected_len + ) + .into()); + } + let proto_variant = ProtoVariant::try_from(state.variant) + .map_err(|_| format!("HyperLogLogState has unknown variant tag {}", state.variant))?; + let variant = match proto_variant { + ProtoVariant::Unspecified => HllVariant::Unspecified, + ProtoVariant::Regular => HllVariant::Regular, + ProtoVariant::Datafusion => HllVariant::Datafusion, + ProtoVariant::Hip => HllVariant::Hip, + }; + let inner = HllSketch::from_raw( + variant, + state.precision, + state.registers.clone(), + state.hip_kxq0, + state.hip_kxq1, + state.hip_est, + ); + Ok(Self { inner }) + } +} + +impl SerializableToSink for HllSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "variant": format!("{:?}", self.inner.variant), + "precision": self.inner.precision, + "register_bytes": self.inner.registers.len(), + "hip_kxq0": self.inner.hip_kxq0, + "hip_kxq1": self.inner.hip_kxq1, + "hip_est": self.inner.hip_est, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.serialize_msgpack() + } +} + +impl AggregateCore for HllSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "HllSketchAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge HllSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_hll = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to HllSketchAccumulator")?; + let merged_inner = HllSketch::merge_refs(&[&self.inner, &other_hll.inner])?; + Ok(Box::new(Self { + inner: merged_inner, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::HLL + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + _statistic: promql_utilities::query_logics::enums::Statistic, + _key: &Option, + _query_kwargs: &HashMap, + ) -> Result> { + Err("HllSketchAccumulator: query_statistic not yet implemented \ + (register round-trip works, but cardinality estimation deferred; \ + tracked as a PR C-CountSketch follow-up)" + .into()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode_state( + variant: i32, + precision: u32, + registers: Vec, + hip_kxq0: f64, + hip_kxq1: f64, + hip_est: f64, + ) -> Vec { + use asap_sketchlib::proto::sketchlib::HyperLogLogState; + use prost::Message; + let state = HyperLogLogState { + variant, + precision, + registers, + hip_kxq0, + hip_kxq1, + hip_est, + }; + state.encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_regular() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + let bytes = encode_state( + ProtoVariant::Regular as i32, + 2, + vec![1, 2, 3, 4], + 0.0, + 0.0, + 0.0, + ); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Regular); + assert_eq!(acc.inner.precision, 2); + assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_hip_preserves_accumulators() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + let bytes = encode_state( + ProtoVariant::Hip as i32, + 2, + vec![0, 0, 0, 0], + 1.5, + 2.5, + 42.0, + ); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Hip); + assert_eq!(acc.inner.hip_kxq0, 1.5); + assert_eq!(acc.inner.hip_kxq1, 2.5); + assert_eq!(acc.inner.hip_est, 42.0); + } + + #[test] + fn test_from_sketchlib_proto_bytes_register_length_mismatch() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + // precision=2 → expected 4 registers; supply only 3 + let bytes = encode_state( + ProtoVariant::Regular as i32, + 2, + vec![1, 2, 3], + 0.0, + 0.0, + 0.0, + ); + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("registers")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_precision_rejected() { + use asap_sketchlib::proto::sketchlib::HyperLogLogState; + use prost::Message; + let state = HyperLogLogState::default(); + let bytes = state.encode_to_vec(); + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + } + + #[test] + fn test_aggregate_core_merge_matches_register_max() { + let a = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0), + }; + let b = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![4, 2, 6, 0], 0.0, 0.0, 0.0), + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.inner.registers, vec![4, 5, 6, 7]); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::precompute_operators::count_sketch_accumulator::CountSketchAccumulator; + let hll = HllSketchAccumulator::new(HllVariant::Regular, 2); + let cs = CountSketchAccumulator::new(2, 3); + assert!(hll.merge_with(&cs).is_err()); + } +} diff --git a/asap-query-engine/src/precompute_operators/mod.rs b/asap-query-engine/src/precompute_operators/mod.rs index 8a52e677..aafdb8ff 100644 --- a/asap-query-engine/src/precompute_operators/mod.rs +++ b/asap-query-engine/src/precompute_operators/mod.rs @@ -1,7 +1,10 @@ pub mod count_min_sketch_accumulator; pub mod count_min_sketch_with_heap_accumulator; +pub mod count_sketch_accumulator; pub mod datasketches_kll_accumulator; +pub mod dd_sketch_accumulator; pub mod delta_set_aggregator_accumulator; +pub mod hll_sketch_accumulator; pub mod hydra_kll_accumulator; pub mod increase_accumulator; pub mod min_max_accumulator; @@ -14,8 +17,11 @@ pub mod sum_accumulator; pub use count_min_sketch_accumulator::*; pub use count_min_sketch_with_heap_accumulator::*; +pub use count_sketch_accumulator::*; pub use datasketches_kll_accumulator::*; +pub use dd_sketch_accumulator::*; pub use delta_set_aggregator_accumulator::*; +pub use hll_sketch_accumulator::*; pub use hydra_kll_accumulator::*; pub use increase_accumulator::*; pub use min_max_accumulator::*; diff --git a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs index fe42ccc2..6fee9d49 100644 --- a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs @@ -242,9 +242,7 @@ impl AggregateCore for MultipleMinMaxAccumulator { fn approx_memory_bytes(&self) -> usize { const BYTES_PER_ENTRY: usize = 96; - std::mem::size_of::() - + self.values.len() * BYTES_PER_ENTRY - + self.sub_type.capacity() + std::mem::size_of::() + self.values.len() * BYTES_PER_ENTRY + self.sub_type.capacity() } fn get_keys(&self) -> Option> { diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index f2eb00bb..05ad9a0f 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -17,7 +17,11 @@ use std::time::{Duration, Instant}; use tracing::{debug, error, info, warn}; use super::persistence::{ - self, cache::PartCache, flusher::FlusherHandle, manifest::Manifest, recovery, + self, + cache::PartCache, + flusher::FlusherHandle, + manifest::Manifest, + recovery, source::{EpochSnapshot, EpochSnapshotEntry, EpochSource, SealedEpochRef}, PersistError, PersistResult, SimpleMapStorePersistenceConfig, }; @@ -329,11 +333,8 @@ impl SimpleMapStorePerKey { hard_cap_bytes, }); - let flusher = FlusherHandle::start( - persistence_cfg, - Arc::clone(&manifest), - Arc::clone(&inner), - )?; + let flusher = + FlusherHandle::start(persistence_cfg, Arc::clone(&manifest), Arc::clone(&inner))?; Ok(Self { inner, @@ -405,7 +406,12 @@ impl SimpleMapStorePerKey { // Skip destructive cleanup entirely — parts on disk are the // source of truth for cold data. if self.inner.persistence_enabled { - let _ = (num_aggregates_to_retain, metric, aggregation_id, read_count_threshold); + let _ = ( + num_aggregates_to_retain, + metric, + aggregation_id, + read_count_threshold, + ); return; } @@ -695,17 +701,19 @@ impl SimpleMapStorePerKey { ); continue; }; - let decoded = - match accumulator_serde::deserialize_accumulator(&disk_entry.sketch_bytes, &sketch_type) { - Ok(a) => a, - Err(e) => { - warn!( - "query_disk_parts: deserialize failed for {}: {}", - disk_entry.sketch_type_name, e - ); - continue; - } - }; + let decoded = match accumulator_serde::deserialize_accumulator( + &disk_entry.sketch_bytes, + &sketch_type, + ) { + Ok(a) => a, + Err(e) => { + warn!( + "query_disk_parts: deserialize failed for {}: {}", + disk_entry.sketch_type_name, e + ); + continue; + } + }; let arc_acc: Arc = Arc::from(decoded); results .entry(disk_entry.label.clone()) diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs index 8d108cf9..10f5fb11 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs @@ -267,7 +267,11 @@ fn run_tick(shared: &Arc, source: &S) -> PersistR if let Some(hot) = cfg.hot_window_ms { let cutoff = now.saturating_sub(hot); for r in &all { - if r.end_ts < cutoff && !selected.iter().any(|s| s.agg_id == r.agg_id && s.epoch_id == r.epoch_id) { + if r.end_ts < cutoff + && !selected + .iter() + .any(|s| s.agg_id == r.agg_id && s.epoch_id == r.epoch_id) + { selected.push(*r); } } @@ -305,8 +309,7 @@ fn run_tick(shared: &Arc, source: &S) -> PersistR let part_id = shared.next_part_id.fetch_add(1, Ordering::Relaxed); let part_dir = part_dir_path(&parts_root(&cfg.disk_path), part_id); let entries_total: usize = snapshots.iter().map(|s| s.len()).sum(); - let size_bytes_estimate: u64 = - snapshots.iter().map(|s| s.approx_bytes as u64).sum(); + let size_bytes_estimate: u64 = snapshots.iter().map(|s| s.approx_bytes as u64).sum(); let report = PartWriter::write_part(&part_dir, part_id, &snapshots)?; shared.manifest.append_add(PartEntry { @@ -398,9 +401,7 @@ fn now_ms() -> u64 { mod tests { use super::*; use crate::data_model::KeyByLabelValues; - use crate::stores::simple_map_store::persistence::source::{ - EpochSnapshot, EpochSnapshotEntry, - }; + use crate::stores::simple_map_store::persistence::source::{EpochSnapshot, EpochSnapshotEntry}; use std::sync::Mutex as StdMutex; use std::time::Duration; use tempfile::TempDir; diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs b/asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs index a29b78f8..7623e305 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs @@ -254,7 +254,10 @@ impl Manifest { write_snapshot_atomic(&self.snapshot_path(), &live)?; // Truncate log. let log_path = self.log_path(); - let f = OpenOptions::new().write(true).truncate(true).open(&log_path)?; + let f = OpenOptions::new() + .write(true) + .truncate(true) + .open(&log_path)?; f.sync_all()?; if let Ok(dir) = File::open(&self.disk_path) { let _ = dir.sync_all(); diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/part.rs b/asap-query-engine/src/stores/simple_map_store/persistence/part.rs index 3e4b41a7..768d78fe 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/part.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/part.rs @@ -188,7 +188,11 @@ impl PartWriter { for pe in &entries_plan { write_u32(&mut data_file, &mut data_crc, pe.label_bytes.len() as u32)?; write_u32(&mut data_file, &mut data_crc, pe.sketch_bytes.len() as u32)?; - write_u16(&mut data_file, &mut data_crc, pe.type_name_bytes.len() as u16)?; + write_u16( + &mut data_file, + &mut data_crc, + pe.type_name_bytes.len() as u16, + )?; write_u16(&mut data_file, &mut data_crc, 0)?; // _pad write_u32(&mut data_file, &mut data_crc, 0)?; // _pad write_padded(&mut data_file, &mut data_crc, &pe.label_bytes, 8)?; @@ -199,7 +203,9 @@ impl PartWriter { data_file.write_all(&data_crc_val.to_le_bytes())?; data_file.write_all(&[0u8; 4])?; // pad data_file.flush()?; - let data_file = data_file.into_inner().map_err(|e| PersistError::Io(e.into_error()))?; + let data_file = data_file + .into_inner() + .map_err(|e| PersistError::Io(e.into_error()))?; data_file.sync_all()?; drop(data_file); @@ -222,7 +228,9 @@ impl PartWriter { index_file.write_all(&index_crc_val.to_le_bytes())?; index_file.write_all(&[0u8; 4])?; index_file.flush()?; - let index_file = index_file.into_inner().map_err(|e| PersistError::Io(e.into_error()))?; + let index_file = index_file + .into_inner() + .map_err(|e| PersistError::Io(e.into_error()))?; index_file.sync_all()?; drop(index_file); @@ -481,8 +489,7 @@ impl PartReader { // 36..40 reserved let data_len = u64::from_le_bytes(buf[40..48].try_into().unwrap()); let index_len = u64::from_le_bytes(buf[48..56].try_into().unwrap()); - let created_unix_secs = - u32::from_le_bytes(buf[56..60].try_into().unwrap()) as u64; + let created_unix_secs = u32::from_le_bytes(buf[56..60].try_into().unwrap()) as u64; let created_unix_ms = created_unix_secs * 1000; let crc_expected = u32::from_le_bytes(buf[60..64].try_into().unwrap()); let crc_actual = crc32fast::hash(&buf[..60]); @@ -511,18 +518,13 @@ impl PartReader { let mut out = Vec::with_capacity(n); for i in 0..n { let off = i * INDEX_ENTRY_SIZE; - let agg_id = u64::from_le_bytes( - self.index_mmap[off..off + 8].try_into().unwrap(), - ); - let start_ts = u64::from_le_bytes( - self.index_mmap[off + 8..off + 16].try_into().unwrap(), - ); - let end_ts = u64::from_le_bytes( - self.index_mmap[off + 16..off + 24].try_into().unwrap(), - ); - let data_offset = u64::from_le_bytes( - self.index_mmap[off + 24..off + 32].try_into().unwrap(), - ); + let agg_id = u64::from_le_bytes(self.index_mmap[off..off + 8].try_into().unwrap()); + let start_ts = + u64::from_le_bytes(self.index_mmap[off + 8..off + 16].try_into().unwrap()); + let end_ts = + u64::from_le_bytes(self.index_mmap[off + 16..off + 24].try_into().unwrap()); + let data_offset = + u64::from_le_bytes(self.index_mmap[off + 24..off + 32].try_into().unwrap()); out.push(IndexRecord { agg_id, start_ts, @@ -543,8 +545,8 @@ impl PartReader { "data.bin offset out of range".to_string(), )); } - let label_len = u32::from_le_bytes(self.data_mmap[off..off + 4].try_into().unwrap()) - as usize; + let label_len = + u32::from_le_bytes(self.data_mmap[off..off + 4].try_into().unwrap()) as usize; let payload_len = u32::from_le_bytes(self.data_mmap[off + 4..off + 8].try_into().unwrap()) as usize; let type_name_len = @@ -629,8 +631,7 @@ mod tests { let tmp = TempDir::new().unwrap(); let part_dir = tmp.path().join("0000000000000001"); let snap = make_snapshot(); - let report = - PartWriter::write_part(&part_dir, 1, &[snap.clone()]).expect("write_part"); + let report = PartWriter::write_part(&part_dir, 1, &[snap.clone()]).expect("write_part"); assert_eq!(report.part_id, 1); assert_eq!(report.num_entries, 2); diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs b/asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs index 3b5f64f3..a0b47995 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs @@ -118,10 +118,8 @@ pub fn recover(disk_path: &Path) -> PersistResult<(Manifest, RecoveryReport)> { mod tests { use super::*; use crate::data_model::KeyByLabelValues; - use crate::stores::simple_map_store::persistence::part::{PartWriter, part_dir_path}; - use crate::stores::simple_map_store::persistence::source::{ - EpochSnapshot, EpochSnapshotEntry, - }; + use crate::stores::simple_map_store::persistence::part::{part_dir_path, PartWriter}; + use crate::stores::simple_map_store::persistence::source::{EpochSnapshot, EpochSnapshotEntry}; use tempfile::TempDir; fn dummy_snapshot() -> EpochSnapshot { @@ -178,15 +176,16 @@ mod tests { let (manifest, _) = recover(tmp.path()).unwrap(); let parts_root = parts_root(tmp.path()); let part_dir = part_dir_path(&parts_root, 42); - let report_write = - PartWriter::write_part(&part_dir, 42, &[dummy_snapshot()]).unwrap(); + let report_write = PartWriter::write_part(&part_dir, 42, &[dummy_snapshot()]).unwrap(); manifest - .append_add(crate::stores::simple_map_store::persistence::manifest::PartEntry { - part_id: 42, - min_ts: report_write.min_ts, - max_ts: report_write.max_ts, - size_bytes: report_write.data_len + report_write.index_len, - }) + .append_add( + crate::stores::simple_map_store::persistence::manifest::PartEntry { + part_id: 42, + min_ts: report_write.min_ts, + max_ts: report_write.max_ts, + size_bytes: report_write.data_len + report_write.index_len, + }, + ) .unwrap(); drop(manifest); diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index e8e2dd5a..010b50df 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -3,11 +3,11 @@ pub mod clickhouse_forwarding_tests; pub mod datafusion; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests; +pub mod persistence_integration_tests; +pub mod persistence_perf_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; pub mod sql_pattern_matching_tests; -pub mod persistence_integration_tests; -pub mod persistence_perf_tests; pub mod store_correctness_tests; pub mod trait_design_tests; diff --git a/asap-query-engine/src/tests/persistence_integration_tests.rs b/asap-query-engine/src/tests/persistence_integration_tests.rs index 1116f31d..f6e14e9d 100644 --- a/asap-query-engine/src/tests/persistence_integration_tests.rs +++ b/asap-query-engine/src/tests/persistence_integration_tests.rs @@ -92,9 +92,8 @@ fn with_persistence_flushes_sealed_epochs_to_disk() { // older than now-0 = now, so flush immediately on next tick." let persistence = persistence_cfg(&dir, Some(0)); - let store = - SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) - .expect("with_persistence"); + let store = SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); // Insert several windows so the rotator seals at least one epoch. // num_aggregates_to_retain = 2, so windows 3 will roll the epoch. @@ -134,7 +133,8 @@ fn with_persistence_flushes_sealed_epochs_to_disk() { .unwrap(); let total: usize = res.values().map(|v| v.len()).sum(); assert_eq!( - total, 4, + total, + 4, "expected 4 buckets across in-memory + disk; got {} (buckets: {:?})", total, res.values() @@ -148,9 +148,8 @@ fn query_read_through_merges_memory_and_disk_ranges() { let dir = TempDir::new().unwrap(); let cfg = make_streaming_config(42); let persistence = persistence_cfg(&dir, Some(0)); - let store = - SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) - .expect("with_persistence"); + let store = SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); // Insert 6 windows — more than enough to guarantee the rotator // seals multiple epochs. @@ -195,9 +194,8 @@ fn construct_and_drop_shuts_flusher_cleanly() { let dir = TempDir::new().unwrap(); let cfg = make_streaming_config(1); let persistence = persistence_cfg(&dir, None); - let store = - SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) - .expect("with_persistence"); + let store = SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); // Dropping the store should not deadlock or panic. drop(store); } @@ -232,9 +230,8 @@ fn hard_cap_back_pressure_blocks_inserts_until_flusher_drains() { disk_path: dir.path().to_path_buf(), part_cache_bytes: 0, }; - let store = - SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) - .expect("with_persistence"); + let store = SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); // Push well past the cap — 200 items × ~16 bytes each, vs. a // 640-byte cap — so the insert path is forced to block on the diff --git a/asap-query-engine/src/tests/persistence_perf_tests.rs b/asap-query-engine/src/tests/persistence_perf_tests.rs index 2267b458..4a9108f4 100644 --- a/asap-query-engine/src/tests/persistence_perf_tests.rs +++ b/asap-query-engine/src/tests/persistence_perf_tests.rs @@ -158,10 +158,7 @@ fn insert_throughput_in_memory_vs_persistent() { // -- baseline: in-memory, NoCleanup -- { - let store = SimpleMapStorePerKey::new( - streaming_config(1, None), - CleanupPolicy::NoCleanup, - ); + let store = SimpleMapStorePerKey::new(streaming_config(1, None), CleanupPolicy::NoCleanup); let items = gen_items(1, N); let d = insert_all(&store, items, BATCH); println!( @@ -244,10 +241,7 @@ fn query_latency_memory_only_vs_disk_through() { // -- in-memory baseline -- { - let store = SimpleMapStorePerKey::new( - streaming_config(1, None), - CleanupPolicy::NoCleanup, - ); + let store = SimpleMapStorePerKey::new(streaming_config(1, None), CleanupPolicy::NoCleanup); let items = gen_items(1, POPULATE); insert_all(&store, items, 1_000); @@ -258,12 +252,7 @@ fn query_latency_memory_only_vs_disk_through() { // -- disk read-through (everything flushed) -- { let tmp = TempDir::new().unwrap(); - let cfg = persistence_cfg( - &tmp, - 4 * 1024 * 1024, - Some(0), - Duration::from_millis(10), - ); + let cfg = persistence_cfg(&tmp, 4 * 1024 * 1024, Some(0), Duration::from_millis(10)); let store = SimpleMapStorePerKey::with_persistence( streaming_config(1, Some(256)), CleanupPolicy::NoCleanup, @@ -460,12 +449,7 @@ fn memory_bound_adherence_under_overload() { ); let tmp = TempDir::new().unwrap(); - let cfg = persistence_cfg( - &tmp, - LIMIT_BYTES, - Some(0), - Duration::from_millis(10), - ); + let cfg = persistence_cfg(&tmp, LIMIT_BYTES, Some(0), Duration::from_millis(10)); let store = SimpleMapStorePerKey::with_persistence( streaming_config(1, Some(128)), CleanupPolicy::NoCleanup, @@ -495,12 +479,7 @@ fn memory_bound_adherence_under_overload() { } let final_tracked = store.diagnostic_info().total_sketch_bytes; - println!( - " inserted {} items in {:?} ({})", - N, - d, - fmt_rate(N, d) - ); + println!(" inserted {} items in {:?} ({})", N, d, fmt_rate(N, d)); println!( " peak tracked: {} KiB (high-water = {} KiB, ratio = {:.2}x)", peak_tracked / 1024, diff --git a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs index 0938b862..09f76d58 100644 --- a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs +++ b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs @@ -27,10 +27,15 @@ use asap_otel_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; use asap_otel_proto::tonic::common::v1::{any_value, AnyValue, KeyValue}; use asap_otel_proto::tonic::metrics::v1::{ - metric::Data, CountMinSketch, CountMinSketchDataPoint, CountMinSketchEncoding, Metric, - ResourceMetrics, ScopeMetrics, + metric::Data, CountMinSketch, CountMinSketchDataPoint, CountMinSketchEncoding, CountSketch, + CountSketchDataPoint, CountSketchEncoding, DdSketch, DdSketchDataPoint, DdSketchEncoding, + HllSketch, HllSketchDataPoint, HllSketchEncoding, KllSketch, KllSketchDataPoint, + KllSketchEncoding, Metric, ResourceMetrics, ScopeMetrics, +}; +use asap_sketchlib::proto::sketchlib::{ + CountMinState, CountSketchState, CounterType, DdSketchState, HllVariant as ProtoHllVariant, + HyperLogLogState, KllState, }; -use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; use asap_types::aggregation_config::AggregationConfig; use asap_types::enums::{AggregationType, WindowType}; use prost::Message; @@ -42,7 +47,10 @@ use query_engine_rust::drivers::ingest::{OtlpReceiver, OtlpReceiverConfig}; use query_engine_rust::precompute_engine::config::{LateDataPolicy, PrecomputeEngineConfig}; use query_engine_rust::precompute_engine::output_sink::CapturingOutputSink; use query_engine_rust::precompute_engine::PrecomputeEngine; -use query_engine_rust::precompute_operators::CountMinSketchAccumulator; +use query_engine_rust::precompute_operators::{ + CountMinSketchAccumulator, CountSketchAccumulator, DDSketchAccumulator, + DatasketchesKLLAccumulator, HllSketchAccumulator, +}; /// Build a tumbling-window `CountMinSketch` aggregation for one metric, /// grouped by a single label. Mirrors the helper in @@ -308,3 +316,768 @@ async fn e2e_count_min_sketch_modified_otlp_path() { ); } } + +// ─── CountSketch path ──────────────────────────────────────────────────── + +/// Parallel to `make_count_min_agg_config` but for `CountSketch`. Uses +/// `AggregationType::CountSketch` added in PR C-CountSketch. +fn make_count_sketch_agg_config( + id: u64, + metric: &str, + window_secs: u64, + grouping: Vec<&str>, + rows: usize, + cols: usize, +) -> AggregationConfig { + let mut params = HashMap::new(); + params.insert("row_num".to_string(), serde_json::Value::from(rows as u64)); + params.insert("col_num".to_string(), serde_json::Value::from(cols as u64)); + AggregationConfig::new( + id, + AggregationType::CountSketch, + String::new(), + params, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new( + grouping.iter().map(|s| s.to_string()).collect(), + ), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + window_secs, + 0, + WindowType::Tumbling, + metric.to_string(), + metric.to_string(), + None, + None, + None, + None, + ) +} + +/// Build a `CountSketchState` proto from a signed matrix in row-major order. +fn build_count_sketch_state(rows: u32, cols: u32, counts_int: Vec) -> CountSketchState { + assert_eq!( + counts_int.len() as u32, + rows * cols, + "counts_int length must equal rows * cols" + ); + CountSketchState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int, + counts_float: Vec::new(), + l2: Vec::new(), + topk: None, + } +} + +/// Build an `ExportMetricsServiceRequest` carrying a single +/// `Metric.data = CountSketch{…}` payload with the given sketch bytes, +/// timestamped at `time_unix_nano` and labeled with `service`. +fn build_count_sketch_export_request( + metric_name: &str, + service_label: &str, + time_unix_nano: u64, + sketch_bytes: Vec, +) -> ExportMetricsServiceRequest { + let dp = CountSketchDataPoint { + attributes: vec![KeyValue { + key: "service".to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(service_label.to_string())), + }), + }], + start_time_unix_nano: 0, + time_unix_nano, + sketch: sketch_bytes, + encoding: CountSketchEncoding::Proto as i32, + dimension: String::new(), + epsilon: 0.0, + delta: 0.0, + flags: 0, + series_id: 0, + }; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: None, + scope_metrics: vec![ScopeMetrics { + scope: None, + metrics: vec![Metric { + name: metric_name.to_string(), + description: String::new(), + unit: String::new(), + metadata: Vec::new(), + data: Some(Data::Countsketch(CountSketch { + data_points: vec![dp], + aggregation_temporality: 0, + })), + }], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn e2e_count_sketch_modified_otlp_path() { + // Same topology as the CountMin test, different ports, different + // aggregation type, and signed counters. + let agg_id = 43u64; + let metric_name = "request_events_total"; + let service_label = "checkout"; + let window_secs = 1u64; + let rows = 2u32; + let cols = 4u32; + + let precompute_port = 19510u16; + let otlp_grpc_port = 19511u16; + let otlp_http_port = 19512u16; + + let cs_config = make_count_sketch_agg_config( + agg_id, + metric_name, + window_secs, + vec!["service"], + rows as usize, + cols as usize, + ); + let mut agg_map = HashMap::new(); + agg_map.insert(agg_id, cs_config); + let streaming_config = Arc::new(StreamingConfig::new(agg_map)); + + let sink = Arc::new(CapturingOutputSink::new()); + let engine = PrecomputeEngine::new( + engine_config(precompute_port), + streaming_config, + sink.clone(), + ); + let ingest_state = engine.ingest_state(); + + tokio::spawn(async move { + let _ = engine.run().await; + }); + + let otlp_receiver = OtlpReceiver::with_ingest_state( + OtlpReceiverConfig { + grpc_port: otlp_grpc_port, + http_port: otlp_http_port, + }, + ingest_state, + ); + tokio::spawn(async move { + let _ = otlp_receiver.run().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(400)).await; + + // Signed counts — the whole point of Count Sketch is ±1 increments, + // so the matrix contains negative values to distinguish it from + // CountMin. + // row 0: [ 1, -2, 3, -4] + // row 1: [-5, 6, -7, 8] + let counts_int: Vec = vec![1, -2, 3, -4, -5, 6, -7, 8]; + let cs_state = build_count_sketch_state(rows, cols, counts_int.clone()); + let sketch_bytes = cs_state.encode_to_vec(); + + let client = reqwest::Client::new(); + + // First send: window 0 payload with the known matrix. + let req = + build_count_sketch_export_request(metric_name, service_label, 100_000_000, sketch_bytes); + post_otlp_http(&client, otlp_http_port, req).await; + + // Second send: watermark advance past window end. + let zero_state = build_count_sketch_state(rows, cols, vec![0i64; (rows * cols) as usize]); + let watermark_advance_req = build_count_sketch_export_request( + metric_name, + service_label, + 2_000_000_000, + zero_state.encode_to_vec(), + ); + post_otlp_http(&client, otlp_http_port, watermark_advance_req).await; + + tokio::time::sleep(tokio::time::Duration::from_millis(800)).await; + + let captured = sink.drain(); + assert!( + !captured.is_empty(), + "expected at least one closed window output, got 0" + ); + + let (window0_output, window0_acc_box) = captured + .iter() + .find(|(out, _)| out.start_timestamp == 0) + .expect("no captured output for window 0"); + + assert_eq!(window0_output.aggregation_id, agg_id); + assert_eq!(window0_output.end_timestamp, window_secs * 1_000); + + let window0_acc = window0_acc_box + .as_any() + .downcast_ref::() + .expect("captured accumulator should be CountSketchAccumulator"); + + let stored_matrix = window0_acc.inner.sketch(); + assert_eq!(stored_matrix.len(), rows as usize, "row count mismatch"); + for r in 0..rows as usize { + let expected: Vec = counts_int[r * cols as usize..(r + 1) * cols as usize] + .iter() + .map(|&v| v as f64) + .collect(); + assert_eq!( + stored_matrix[r], expected, + "matrix row {r} mismatch (expected {expected:?}, got {:?})", + stored_matrix[r] + ); + } +} + +// ─── KLL sketch path ───────────────────────────────────────────────────── + +fn make_kll_agg_config( + id: u64, + metric: &str, + window_secs: u64, + grouping: Vec<&str>, + k: u32, +) -> AggregationConfig { + let mut params = HashMap::new(); + params.insert("k".to_string(), serde_json::Value::from(k)); + AggregationConfig::new( + id, + AggregationType::DatasketchesKLL, + "DatasketchesKLL".to_string(), + params, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new( + grouping.iter().map(|s| s.to_string()).collect(), + ), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + window_secs, + 0, + WindowType::Tumbling, + metric.to_string(), + metric.to_string(), + None, + None, + None, + None, + ) +} + +/// Build a `KllState` proto carrying the given retained items. Level +/// metadata is not populated — the decoder replays items via `update()` +/// regardless, per the lossy-reconstruction strategy documented on +/// `DatasketchesKLLAccumulator::from_sketchlib_proto_bytes`. +fn build_kll_state(k: u32, items: Vec) -> KllState { + KllState { + k, + m: 8, + num_levels: 0, + levels: Vec::new(), + items, + coin: None, + } +} + +fn build_kll_export_request( + metric_name: &str, + service_label: &str, + time_unix_nano: u64, + sketch_bytes: Vec, +) -> ExportMetricsServiceRequest { + let dp = KllSketchDataPoint { + attributes: vec![KeyValue { + key: "service".to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(service_label.to_string())), + }), + }], + start_time_unix_nano: 0, + time_unix_nano, + count: 0, + sum: 0.0, + min: 0.0, + max: 0.0, + sketch: sketch_bytes, + encoding: KllSketchEncoding::Proto as i32, + flags: 0, + series_id: 0, + }; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: None, + scope_metrics: vec![ScopeMetrics { + scope: None, + metrics: vec![Metric { + name: metric_name.to_string(), + description: String::new(), + unit: String::new(), + metadata: Vec::new(), + data: Some(Data::Kllsketch(KllSketch { + data_points: vec![dp], + aggregation_temporality: 0, + })), + }], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn e2e_kll_sketch_modified_otlp_path() { + let agg_id = 44u64; + let metric_name = "request_latency_ms"; + let service_label = "api"; + let window_secs = 1u64; + let k = 200u32; + + let precompute_port = 19520u16; + let otlp_grpc_port = 19521u16; + let otlp_http_port = 19522u16; + + let kll_config = make_kll_agg_config(agg_id, metric_name, window_secs, vec!["service"], k); + let mut agg_map = HashMap::new(); + agg_map.insert(agg_id, kll_config); + let streaming_config = Arc::new(StreamingConfig::new(agg_map)); + + let sink = Arc::new(CapturingOutputSink::new()); + let engine = PrecomputeEngine::new( + engine_config(precompute_port), + streaming_config, + sink.clone(), + ); + let ingest_state = engine.ingest_state(); + + tokio::spawn(async move { + let _ = engine.run().await; + }); + + let otlp_receiver = OtlpReceiver::with_ingest_state( + OtlpReceiverConfig { + grpc_port: otlp_grpc_port, + http_port: otlp_http_port, + }, + ingest_state, + ); + tokio::spawn(async move { + let _ = otlp_receiver.run().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(400)).await; + + let items: Vec = (1..=100).map(|v| v as f64).collect(); + let kll_state = build_kll_state(k, items); + let sketch_bytes = kll_state.encode_to_vec(); + + let client = reqwest::Client::new(); + let req = build_kll_export_request(metric_name, service_label, 100_000_000, sketch_bytes); + post_otlp_http(&client, otlp_http_port, req).await; + + let empty_state = build_kll_state(k, Vec::new()); + let watermark_req = build_kll_export_request( + metric_name, + service_label, + 2_000_000_000, + empty_state.encode_to_vec(), + ); + post_otlp_http(&client, otlp_http_port, watermark_req).await; + + tokio::time::sleep(tokio::time::Duration::from_millis(800)).await; + + let captured = sink.drain(); + assert!( + !captured.is_empty(), + "expected at least one closed window output, got 0" + ); + + let (window0_output, window0_acc_box) = captured + .iter() + .find(|(out, _)| out.start_timestamp == 0) + .expect("no captured output for window 0"); + + assert_eq!(window0_output.aggregation_id, agg_id); + assert_eq!(window0_output.end_timestamp, window_secs * 1_000); + + let kll_acc = window0_acc_box + .as_any() + .downcast_ref::() + .expect("captured accumulator should be DatasketchesKLLAccumulator"); + + let p50 = kll_acc.get_quantile(0.5); + assert!( + (30.0..=70.0).contains(&p50), + "p50 should be close to 50, got {p50}" + ); +} + +// ─── DDSketch path ─────────────────────────────────────────────────────── + +fn make_dd_sketch_agg_config( + id: u64, + metric: &str, + window_secs: u64, + grouping: Vec<&str>, + alpha: f64, +) -> AggregationConfig { + let mut params = HashMap::new(); + params.insert("alpha".to_string(), serde_json::Value::from(alpha)); + AggregationConfig::new( + id, + AggregationType::DDSketch, + String::new(), + params, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new( + grouping.iter().map(|s| s.to_string()).collect(), + ), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + window_secs, + 0, + WindowType::Tumbling, + metric.to_string(), + metric.to_string(), + None, + None, + None, + None, + ) +} + +fn build_dd_sketch_state( + alpha: f64, + store_counts: Vec, + store_offset: i32, + count: u64, + sum: f64, + min: f64, + max: f64, +) -> DdSketchState { + DdSketchState { + alpha, + store_counts, + store_offset, + count, + sum, + min, + max, + } +} + +fn build_dd_sketch_export_request( + metric_name: &str, + service_label: &str, + time_unix_nano: u64, + sketch_bytes: Vec, +) -> ExportMetricsServiceRequest { + let dp = DdSketchDataPoint { + attributes: vec![KeyValue { + key: "service".to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(service_label.to_string())), + }), + }], + start_time_unix_nano: 0, + time_unix_nano, + count: 0, + sketch: sketch_bytes, + encoding: DdSketchEncoding::DdsketchEncodingProto as i32, + exemplars: Vec::new(), + flags: 0, + series_id: 0, + sum: None, + min: None, + max: None, + }; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: None, + scope_metrics: vec![ScopeMetrics { + scope: None, + metrics: vec![Metric { + name: metric_name.to_string(), + description: String::new(), + unit: String::new(), + metadata: Vec::new(), + data: Some(Data::Ddsketch(DdSketch { + data_points: vec![dp], + aggregation_temporality: 0, + })), + }], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn e2e_dd_sketch_modified_otlp_path() { + let agg_id = 45u64; + let metric_name = "request_duration_seconds"; + let service_label = "payments"; + let window_secs = 1u64; + let alpha = 0.01; + + let precompute_port = 19530u16; + let otlp_grpc_port = 19531u16; + let otlp_http_port = 19532u16; + + let dd_config = + make_dd_sketch_agg_config(agg_id, metric_name, window_secs, vec!["service"], alpha); + let mut agg_map = HashMap::new(); + agg_map.insert(agg_id, dd_config); + let streaming_config = Arc::new(StreamingConfig::new(agg_map)); + + let sink = Arc::new(CapturingOutputSink::new()); + let engine = PrecomputeEngine::new( + engine_config(precompute_port), + streaming_config, + sink.clone(), + ); + let ingest_state = engine.ingest_state(); + + tokio::spawn(async move { + let _ = engine.run().await; + }); + + let otlp_receiver = OtlpReceiver::with_ingest_state( + OtlpReceiverConfig { + grpc_port: otlp_grpc_port, + http_port: otlp_http_port, + }, + ingest_state, + ); + tokio::spawn(async move { + let _ = otlp_receiver.run().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(400)).await; + + let store_counts = vec![5u64, 10, 15, 20]; + let dd_state = build_dd_sketch_state(alpha, store_counts.clone(), -1, 50, 150.0, 0.25, 8.0); + let sketch_bytes = dd_state.encode_to_vec(); + + let client = reqwest::Client::new(); + let req = build_dd_sketch_export_request(metric_name, service_label, 100_000_000, sketch_bytes); + post_otlp_http(&client, otlp_http_port, req).await; + + let watermark_state = build_dd_sketch_state(alpha, Vec::new(), 0, 0, 0.0, 0.0, 0.0); + let watermark_req = build_dd_sketch_export_request( + metric_name, + service_label, + 2_000_000_000, + watermark_state.encode_to_vec(), + ); + post_otlp_http(&client, otlp_http_port, watermark_req).await; + + tokio::time::sleep(tokio::time::Duration::from_millis(800)).await; + + let captured = sink.drain(); + assert!(!captured.is_empty(), "expected at least one output"); + + let (window0_output, window0_acc_box) = captured + .iter() + .find(|(out, _)| out.start_timestamp == 0) + .expect("no captured output for window 0"); + + assert_eq!(window0_output.aggregation_id, agg_id); + + let dd_acc = window0_acc_box + .as_any() + .downcast_ref::() + .expect("captured accumulator should be DDSketchAccumulator"); + + assert_eq!(dd_acc.inner.store_counts, store_counts); + assert_eq!(dd_acc.inner.store_offset, -1); + assert_eq!(dd_acc.inner.count, 50); + assert_eq!(dd_acc.inner.sum, 150.0); + assert!((dd_acc.inner.alpha - alpha).abs() < f64::EPSILON); +} + +// ─── HLL sketch path ───────────────────────────────────────────────────── + +fn make_hll_agg_config( + id: u64, + metric: &str, + window_secs: u64, + grouping: Vec<&str>, + precision: u32, +) -> AggregationConfig { + let mut params = HashMap::new(); + params.insert("precision".to_string(), serde_json::Value::from(precision)); + AggregationConfig::new( + id, + AggregationType::HLL, + String::new(), + params, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new( + grouping.iter().map(|s| s.to_string()).collect(), + ), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + window_secs, + 0, + WindowType::Tumbling, + metric.to_string(), + metric.to_string(), + None, + None, + None, + None, + ) +} + +fn build_hll_state(precision: u32, registers: Vec) -> HyperLogLogState { + HyperLogLogState { + variant: ProtoHllVariant::Regular as i32, + precision, + registers, + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + } +} + +fn build_hll_export_request( + metric_name: &str, + service_label: &str, + time_unix_nano: u64, + sketch_bytes: Vec, + precision: u32, +) -> ExportMetricsServiceRequest { + let dp = HllSketchDataPoint { + attributes: vec![KeyValue { + key: "service".to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(service_label.to_string())), + }), + }], + start_time_unix_nano: 0, + time_unix_nano, + count: 0, + cardinality: 0, + sketch: sketch_bytes, + encoding: HllSketchEncoding::Proto as i32, + precision, + flags: 0, + series_id: 0, + }; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: None, + scope_metrics: vec![ScopeMetrics { + scope: None, + metrics: vec![Metric { + name: metric_name.to_string(), + description: String::new(), + unit: String::new(), + metadata: Vec::new(), + data: Some(Data::Hllsketch(HllSketch { + data_points: vec![dp], + aggregation_temporality: 0, + })), + }], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn e2e_hll_sketch_modified_otlp_path() { + let agg_id = 46u64; + let metric_name = "unique_users"; + let service_label = "login"; + let window_secs = 1u64; + let precision = 4u32; + let num_registers = 1usize << precision; + + let precompute_port = 19540u16; + let otlp_grpc_port = 19541u16; + let otlp_http_port = 19542u16; + + let hll_config = + make_hll_agg_config(agg_id, metric_name, window_secs, vec!["service"], precision); + let mut agg_map = HashMap::new(); + agg_map.insert(agg_id, hll_config); + let streaming_config = Arc::new(StreamingConfig::new(agg_map)); + + let sink = Arc::new(CapturingOutputSink::new()); + let engine = PrecomputeEngine::new( + engine_config(precompute_port), + streaming_config, + sink.clone(), + ); + let ingest_state = engine.ingest_state(); + + tokio::spawn(async move { + let _ = engine.run().await; + }); + + let otlp_receiver = OtlpReceiver::with_ingest_state( + OtlpReceiverConfig { + grpc_port: otlp_grpc_port, + http_port: otlp_http_port, + }, + ingest_state, + ); + tokio::spawn(async move { + let _ = otlp_receiver.run().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(400)).await; + + let registers: Vec = (0..num_registers as u8).collect(); + let hll_state = build_hll_state(precision, registers.clone()); + let sketch_bytes = hll_state.encode_to_vec(); + + let client = reqwest::Client::new(); + let req = build_hll_export_request( + metric_name, + service_label, + 100_000_000, + sketch_bytes, + precision, + ); + post_otlp_http(&client, otlp_http_port, req).await; + + let watermark_state = build_hll_state(precision, vec![0u8; num_registers]); + let watermark_req = build_hll_export_request( + metric_name, + service_label, + 2_000_000_000, + watermark_state.encode_to_vec(), + precision, + ); + post_otlp_http(&client, otlp_http_port, watermark_req).await; + + tokio::time::sleep(tokio::time::Duration::from_millis(800)).await; + + let captured = sink.drain(); + assert!(!captured.is_empty(), "expected at least one output"); + + let (window0_output, window0_acc_box) = captured + .iter() + .find(|(out, _)| out.start_timestamp == 0) + .expect("no captured output for window 0"); + + assert_eq!(window0_output.aggregation_id, agg_id); + + let hll_acc = window0_acc_box + .as_any() + .downcast_ref::() + .expect("captured accumulator should be HllSketchAccumulator"); + + assert_eq!(hll_acc.inner.precision, precision); + assert_eq!(hll_acc.inner.registers, registers); +}