From 00c754c9edd343d4eb1391c8638cd5d332eab7e0 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 09:06:57 -0400 Subject: [PATCH 1/2] fix(query-engine): run top-k CMSHeap on FastPath --- .../src/engines/simple_engine/mod.rs | 2 +- .../src/engines/simple_engine/promql.rs | 10 +- .../src/engines/simple_engine/sql.rs | 8 +- .../precompute_engine/accumulator_factory.rs | 20 +- .../count_min_sketch_with_heap_accumulator.rs | 210 +++++++----------- .../src/tests/native_range_query_tests.rs | 16 +- .../src/tests/range_query_arithmetic_tests.rs | 4 +- ...stage_e_instant_range_equivalence_tests.rs | 2 +- 8 files changed, 121 insertions(+), 151 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index c0bb98bd..4d71067f 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -5318,7 +5318,7 @@ mod stage_e4_instant_wrapper_equivalence_tests { )); let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for (host, value) in [("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)] { - sketch.inner.update(host, value); + sketch.update(host, value as i32); } let output = PrecomputedOutput::new(2000, 3000, None, 1); store diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 0ecc764d..d42dcb3e 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -1776,7 +1776,7 @@ mod topk_pipeline_tests { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for i in 1..=15u64 { let srcip = format!("10.0.0.{i}"); - sketch.inner.update(&srcip, (i * 10) as f64); + sketch.update(&srcip, (i * 10) as i32); } let output = @@ -1894,7 +1894,7 @@ mod topk_pipeline_tests { let window = &context.store_plan.values_query; let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for i in 1..=15u64 { - sketch.inner.update(&format!("10.0.0.{i}"), (i * 10) as f64); + sketch.update(&format!("10.0.0.{i}"), (i * 10) as i32); } store .insert_precomputed_output( @@ -1981,7 +1981,7 @@ mod topk_pipeline_tests { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for i in 1..=15u64 { let srcip = format!("10.0.0.{i}"); - sketch.inner.update(&srcip, (i * 10) as f64); + sketch.update(&srcip, (i * 10) as i32); } let output = @@ -2155,7 +2155,7 @@ mod topk_pipeline_tests { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for i in 1..=15u64 { let srcip = format!("10.0.0.{i}"); - sketch.inner.update(&srcip, (i * 10) as f64); + sketch.update(&srcip, (i * 10) as i32); } let output = @@ -2213,7 +2213,7 @@ mod topk_pipeline_tests { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for i in 1..=15u64 { let srcip = format!("10.0.0.{i}"); - sketch.inner.update(&srcip, (i * 10) as f64); + sketch.update(&srcip, (i * 10) as i32); } let output = PrecomputedOutput::new(window.start_timestamp, window.end_timestamp, None, AGG_ID); diff --git a/asap-query-engine/src/engines/simple_engine/sql.rs b/asap-query-engine/src/engines/simple_engine/sql.rs index 11b71320..716ed3ff 100644 --- a/asap-query-engine/src/engines/simple_engine/sql.rs +++ b/asap-query-engine/src/engines/simple_engine/sql.rs @@ -1477,7 +1477,7 @@ mod topk_pipeline_tests { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for i in 1..=15u64 { let srcip = format!("10.0.0.{i}"); - sketch.inner.update(&srcip, (i * 10) as f64); + sketch.update(&srcip, (i * 10) as i32); } let output = @@ -1547,9 +1547,9 @@ mod topk_pipeline_tests { // heap keys join in that order regardless of query_output_labels' // own order. let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); - sketch.inner.update("10.0.0.1;shared-src", 100.0); - sketch.inner.update("10.0.0.2;shared-src", 200.0); - sketch.inner.update("10.0.0.3;other-src", 50.0); + sketch.update("10.0.0.1;shared-src", 100); + sketch.update("10.0.0.2;shared-src", 200); + sketch.update("10.0.0.3;other-src", 50); let output = PrecomputedOutput::new(window.start_timestamp, window.end_timestamp, None, AGG_ID); diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index e55149b5..581309d9 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -758,8 +758,12 @@ impl AccumulatorUpdater for CmsWithHeapAccumulatorUpdater { } fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { - let weight = if self.count_events { 1.0 } else { value }; - self.acc.inner.update(&key.to_semicolon_str(), weight); + let count = if self.count_events { + 1 + } else { + value.round() as i32 + }; + self.acc.update(&key.to_semicolon_str(), count); } impl_accumulator_methods!(acc); @@ -775,7 +779,7 @@ impl AccumulatorUpdater for CmsWithHeapAccumulatorUpdater { fn memory_usage_bytes(&self) -> usize { std::mem::size_of::() - + self.row_num * self.col_num * std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() } } @@ -1938,6 +1942,16 @@ mod tests { assert_eq!(cms.query_key(&key), 50.0, "sum of 5×10 == 50"); } + #[test] + fn test_cms_with_heap_memory_uses_i32_counters() { + let updater = CmsWithHeapAccumulatorUpdater::new(3, 7, 32, true); + + assert_eq!( + updater.memory_usage_bytes(), + std::mem::size_of::() + 3 * 7 * 4 + ); + } + #[test] fn test_cms_heap_params_reads_depth_width_heapsize() { let mut params = cms_heap_params_required(); diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs index 486ed23b..55919000 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs @@ -2,113 +2,92 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; -use asap_sketchlib::{message_pack_format::MessagePackCodec, CmsHeapItem, CountMinSketchWithHeap}; +use asap_sketchlib::{CMSHeap, DataInput, FastPath, HeapItem, Vector2D}; use asap_types::traits::SerializationError; use serde_json::Value; use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -/// Count-Min Sketch with Heap accumulator — wraps `asap_sketchlib::sketches::CountMinSketchWithHeap`. -/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches::countminsketch_topk`. -/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. -#[derive(Debug, Clone)] +/// Count-Min Sketch with Heap accumulator backed by the benchmarked FastPath implementation. +#[derive(Clone)] pub struct CountMinSketchWithHeapAccumulator { - pub inner: CountMinSketchWithHeap, + inner: CMSHeap, FastPath>, } -// Re-export HeapItem so existing code using CountMinSketchWithHeapAccumulator::HeapItem still works. -pub use asap_sketchlib::CmsHeapItem as HeapItemReexport; +impl std::fmt::Debug for CountMinSketchWithHeapAccumulator { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("CountMinSketchWithHeapAccumulator") + .field("rows", &self.inner.rows()) + .field("cols", &self.inner.cols()) + .field("heap_size", &self.inner.heap().capacity()) + .finish() + } +} impl CountMinSketchWithHeapAccumulator { pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { Self { - inner: CountMinSketchWithHeap::new(row_num, col_num, heap_size), + inner: CMSHeap::new(row_num, col_num, heap_size), } } - pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { - let key_string = key.labels.join(";"); - self.inner.estimate(&key_string) - } - - /// This function seems will never be used anymore. Keep it for possible future use. - pub fn deserialize_from_json(data: &Value) -> Result> { - let row_num = data["row_num"] - .as_f64() - .ok_or("Missing or invalid 'row_num' field")? as usize; - let col_num = data["col_num"] - .as_f64() - .ok_or("Missing or invalid 'col_num' field")? as usize; - let heap_size = data["heap_size"] - .as_f64() - .ok_or("Missing or invalid 'heap_size' field")? as usize; - - let sketch_data = data["sketch"] - .as_array() - .ok_or("Missing or invalid 'sketch' field")?; - - let mut sketch = Vec::new(); - for row in sketch_data { - let row_array = row.as_array().ok_or("Invalid row in sketch data")?; - let mut sketch_row = Vec::new(); - for cell in row_array { - let value = cell.as_f64().ok_or("Invalid cell value in sketch data")?; - sketch_row.push(value); - } - sketch.push(sketch_row); - } - - let topk_heap_data = data["topk_heap"] - .as_array() - .ok_or("Missing or invalid 'topk_heap' field")?; - - let mut topk_heap = Vec::new(); - for item in topk_heap_data { - let key = item["key"] - .as_str() - .ok_or("Missing or invalid 'key' in heap item")? - .to_string(); - let value = item["value"] - .as_f64() - .ok_or("Missing or invalid 'value' in heap item")?; - topk_heap.push(CmsHeapItem { key, value }); + pub fn update(&mut self, key: &str, count: i32) { + if count > 0 { + self.inner.insert_many(&DataInput::Str(key), count); } - - Ok(Self { - inner: CountMinSketchWithHeap::from_legacy_matrix( - sketch, topk_heap, row_num, col_num, heap_size, - ), - }) } - pub fn deserialize_from_bytes(_buffer: &[u8]) -> Result> { - Err("deserialize_from_bytes for CountMinSketchWithHeapAccumulator not implemented".into()) + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + let key_string = key.labels.join(";"); + self.inner.estimate(&DataInput::Str(&key_string)) as f64 } /// Get all keys from the top-k heap. pub fn get_topk_keys(&self) -> Vec { self.inner - .topk_heap_items() + .heap() + .heap() .iter() .map(|item| { - let labels: Vec = item.key.split(';').map(|s| s.to_string()).collect(); + let key = match &item.key { + HeapItem::String(key) => key, + other => panic!("expected string heap key, got {other:?}"), + }; + let labels: Vec = key.split(';').map(str::to_owned).collect(); KeyByLabelValues { labels } }) .collect() } + + fn sketch_matrix(&self) -> Vec> { + let storage = self.inner.cms().as_storage(); + (0..storage.rows()) + .map(|row| { + (0..storage.cols()) + .map(|col| *storage.get(row, col).expect("matrix index is in bounds")) + .collect() + }) + .collect() + } } impl SerializableToSink for CountMinSketchWithHeapAccumulator { fn serialize_to_json(&self) -> Result { let heap_items: Vec = self .inner - .topk_heap_items() + .heap() + .heap() .iter() .map(|item| { + let key = match &item.key { + HeapItem::String(key) => key, + other => panic!("expected string heap key, got {other:?}"), + }; serde_json::json!({ - "key": item.key, - "value": item.value + "key": key, + "value": item.count }) }) .collect(); @@ -116,19 +95,17 @@ impl SerializableToSink for CountMinSketchWithHeapAccumulator { Ok(serde_json::json!({ "row_num": self.inner.rows(), "col_num": self.inner.cols(), - "heap_size": self.inner.heap_size, - "sketch": self.inner.sketch_matrix(), + "heap_size": self.inner.heap().capacity(), + "sketch": self.sketch_matrix(), "topk_heap": heap_items })) } fn serialize_to_bytes(&self) -> Result, SerializationError> { - self.inner - .to_msgpack() - .map_err(|e| SerializationError::Bytes { - type_name: "CountMinSketchWithHeapAccumulator", - source: e.to_string().into(), - }) + Err(SerializationError::Bytes { + type_name: "CountMinSketchWithHeapAccumulator", + source: "byte serialization is not supported for FastPath CMSHeap".into(), + }) } } @@ -213,7 +190,10 @@ impl MergeableAccumulator for CountMinSketchW let mut iter = accumulators.into_iter(); let mut merged = iter.next().unwrap(); for acc in iter { - merged.inner.merge(&acc.inner)?; + if merged.inner.rows() != acc.inner.rows() || merged.inner.cols() != acc.inner.cols() { + return Err("CountMinSketchWithHeap dimension mismatch".into()); + } + merged.inner.merge(&acc.inner); } Ok(merged) } @@ -228,8 +208,8 @@ mod tests { let cms = CountMinSketchWithHeapAccumulator::new(4, 1000, 20); assert_eq!(cms.inner.rows(), 4); assert_eq!(cms.inner.cols(), 1000); - assert_eq!(cms.inner.heap_size, 20); - assert_eq!(cms.inner.topk_heap_items().len(), 0); + assert_eq!(cms.inner.heap().capacity(), 20); + assert_eq!(cms.inner.heap().len(), 0); } #[test] @@ -244,50 +224,17 @@ mod tests { #[test] fn test_count_min_sketch_with_heap_merge() { - // Build controlled state via from_legacy_matrix (works regardless of backend config). - let sketch1 = vec![ - vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - ]; - let heap1 = vec![ - CmsHeapItem { - key: "key1".to_string(), - value: 100.0, - }, - CmsHeapItem { - key: "key2".to_string(), - value: 50.0, - }, - ]; - let sketch2 = vec![ - vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - ]; - let heap2 = vec![ - CmsHeapItem { - key: "key3".to_string(), - value: 75.0, - }, - CmsHeapItem { - key: "key1".to_string(), - value: 80.0, - }, - ]; - - let cms1 = CountMinSketchWithHeapAccumulator { - inner: CountMinSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), - }; - let cms2 = CountMinSketchWithHeapAccumulator { - inner: CountMinSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), - }; - - let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); - assert!(result.is_ok()); - let merged = result.unwrap(); - assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); - assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); - assert_eq!(merged.inner.heap_size, 3); - assert!(merged.inner.topk_heap_items().len() <= 3); + let mut left = CountMinSketchWithHeapAccumulator::new(3, 1024, 5); + let mut right = CountMinSketchWithHeapAccumulator::new(3, 1024, 3); + left.update("key1", 100); + right.update("key1", 80); + right.update("key2", 75); + + let merged = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![left, right]) + .expect("matching sketches merge"); + let key = KeyByLabelValues::new_with_labels(vec!["key1".to_owned()]); + assert_eq!(merged.query_key(&key), 180.0); + assert_eq!(merged.inner.heap().capacity(), 5); } #[test] @@ -298,7 +245,7 @@ mod tests { let merged = result.unwrap(); assert_eq!(merged.inner.rows(), cms.inner.rows()); assert_eq!(merged.inner.cols(), cms.inner.cols()); - assert_eq!(merged.inner.heap_size, cms.inner.heap_size); + assert_eq!(merged.inner.heap().capacity(), cms.inner.heap().capacity()); } #[test] @@ -319,12 +266,12 @@ mod tests { #[test] fn test_get_topk_keys() { let mut cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); - cms.inner.update("label1;label2", 100.0); - cms.inner.update("label3;label4", 50.0); + cms.update("label1;label2", 100); + cms.update("label3;label4", 50); let keys = cms.get_topk_keys(); assert_eq!(keys.len(), 2); - // Top-k order can differ between Legacy and Sketchlib backends (heap ordering / estimates). + // Heap ordering is not part of the top-k interface. let label_sets: std::collections::HashSet<_> = keys.iter().map(|k| k.labels.clone()).collect(); assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); @@ -344,4 +291,13 @@ mod tests { assert!(keys.is_some()); assert_eq!(keys.unwrap().len(), 0); } + + #[test] + fn byte_serialization_is_unsupported() { + let cms = CountMinSketchWithHeapAccumulator::new(3, 50, 10); + + let error = cms.serialize_to_bytes().unwrap_err(); + + assert!(error.to_string().contains("not supported")); + } } diff --git a/asap-query-engine/src/tests/native_range_query_tests.rs b/asap-query-engine/src/tests/native_range_query_tests.rs index c900a7e9..f3283f47 100644 --- a/asap-query-engine/src/tests/native_range_query_tests.rs +++ b/asap-query-engine/src/tests/native_range_query_tests.rs @@ -936,8 +936,8 @@ mod tests { // keys, at every step, so a wrong implementation is caught by // key-set mismatch, not just a wrong value. for heap in [&mut heap_1, &mut heap_2] { - heap.inner.update("host-x;evt-x", 30.0); - heap.inner.update("host-y;evt-y", 20.0); + heap.update("host-x;evt-x", 30); + heap.update("host-y;evt-y", 20); } let mut keys_1000 = DeltaSetAggregatorAccumulator::new(); @@ -2282,9 +2282,9 @@ mod tests { // call, this None-keyed group is dropped and the range query returns // empty instead of the expanded top-k keys. let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); - sketch.inner.update("host-a", 30.0); - sketch.inner.update("host-b", 20.0); - sketch.inner.update("host-c", 10.0); + sketch.update("host-a", 30); + sketch.update("host-b", 20); + sketch.update("host-c", 10); let query = "topk(5, transfer_events)"; let engine = @@ -2491,9 +2491,9 @@ mod tests { )); let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); - sketch.inner.update("host-a", 30.0); - sketch.inner.update("host-b", 20.0); - sketch.inner.update("host-c", 10.0); + sketch.update("host-a", 30); + sketch.update("host-b", 20); + sketch.update("host-c", 10); // Non-None outer key -- e.g. Some(KeyByLabelValues{labels: [""]}), // exactly what real worker.rs ingestion produces for "" grouping. diff --git a/asap-query-engine/src/tests/range_query_arithmetic_tests.rs b/asap-query-engine/src/tests/range_query_arithmetic_tests.rs index a18c6497..82d1484a 100644 --- a/asap-query-engine/src/tests/range_query_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/range_query_arithmetic_tests.rs @@ -370,7 +370,7 @@ mod tests { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for (host, value) in topk_candidates { - sketch.inner.update(host, *value); + sketch.update(host, value.round() as i32); } let topk_output = PrecomputedOutput::new(0, WINDOW_MS, None, 1); store @@ -505,7 +505,7 @@ mod tests { for (agg_id, candidates) in [(1u64, candidates_a), (2u64, candidates_b)] { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for (host, value) in candidates { - sketch.inner.update(host, *value); + sketch.update(host, value.round() as i32); } let output = PrecomputedOutput::new(0, WINDOW_MS, None, agg_id); store diff --git a/asap-query-engine/src/tests/stage_e_instant_range_equivalence_tests.rs b/asap-query-engine/src/tests/stage_e_instant_range_equivalence_tests.rs index 510ba17c..2e8341e2 100644 --- a/asap-query-engine/src/tests/stage_e_instant_range_equivalence_tests.rs +++ b/asap-query-engine/src/tests/stage_e_instant_range_equivalence_tests.rs @@ -449,7 +449,7 @@ mod tests { for (start, end, candidates) in buckets { let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); for (host, value) in *candidates { - sketch.inner.update(host, *value); + sketch.update(host, value.round() as i32); } let output = PrecomputedOutput::new(*start, *end, None, 1); store From deacf1fa9d0ea1dfcd2dc3bafb578627a2351e09 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 09:25:10 -0400 Subject: [PATCH 2/2] test(query-engine): cover FastPath top-k behavior --- .../precompute_engine/accumulator_factory.rs | 4 ++-- .../count_min_sketch_with_heap_accumulator.rs | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index 581309d9..e6aefba2 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -1932,14 +1932,14 @@ mod tests { let key = KeyByLabelValues::new_with_labels(vec!["10.0.0.1".to_string()]); for _ in 0..5 { - updater.update_keyed(&key, 10.0, 0); + updater.update_keyed(&key, 1.6, 0); } let acc = updater.take_accumulator(); let cms = acc .as_any() .downcast_ref::() .unwrap(); - assert_eq!(cms.query_key(&key), 50.0, "sum of 5×10 == 50"); + assert_eq!(cms.query_key(&key), 10.0, "sum of 5×round(1.6) == 10"); } #[test] diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs index 55919000..343364fd 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs @@ -222,6 +222,30 @@ mod tests { assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); } + #[test] + fn fastpath_fixed_stream_matches_expected_estimates_and_topk() { + let mut cms = CountMinSketchWithHeapAccumulator::new(3, 1024, 2); + cms.update("api", 7); + cms.update("worker", 3); + cms.update("api", 5); + cms.update("batch", 1); + + let api = KeyByLabelValues::new_with_labels(vec!["api".to_owned()]); + let worker = KeyByLabelValues::new_with_labels(vec!["worker".to_owned()]); + assert_eq!(cms.query_key(&api), 12.0); + assert_eq!(cms.query_key(&worker), 3.0); + + let keys: std::collections::HashSet<_> = cms + .get_topk_keys() + .into_iter() + .map(|key| key.labels) + .collect(); + assert_eq!( + keys, + std::collections::HashSet::from([vec!["api".to_owned()], vec!["worker".to_owned()]]) + ); + } + #[test] fn test_count_min_sketch_with_heap_merge() { let mut left = CountMinSketchWithHeapAccumulator::new(3, 1024, 5);