From cacefdc5d888d341a54378f1ca423999a4b13ee0 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 26 Aug 2026 15:23:00 -0400 Subject: [PATCH] refactor(query-engine): centralize range topk maintenance --- .../src/engines/simple_engine/mod.rs | 66 +++++++++++++------ ...stage_e_instant_range_equivalence_tests.rs | 13 +--- 2 files changed, 48 insertions(+), 31 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 99a1856..4e583ea 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -990,15 +990,8 @@ impl SimpleEngine { &context.metric, enable_topk_formatting, ); - if enable_topk_limiting { - if let Some(k) = context - .metadata - .query_kwargs - .get("k") - .and_then(|s| s.parse::().ok()) - { - results.truncate(k); - } + if enable_topk_limiting && context.metadata.statistic_to_compute == Statistic::Topk { + results.truncate(Self::parse_topk_limit(&context.metadata.query_kwargs)?); } Ok(results) @@ -1022,6 +1015,24 @@ impl SimpleEngine { .then_with(|| a_labels.cmp(b_labels)) } + /// Returns the required `k` parameter for a Topk query. + /// + /// PromQL context construction validates this before execution, but the + /// pipeline is also used directly in tests and by non-PromQL callers. Do + /// not silently turn malformed Topk metadata into an unbounded result. + fn parse_topk_limit(query_kwargs: &HashMap) -> Result { + query_kwargs + .get("k") + .ok_or_else(|| "Topk query is missing required `k` parameter".to_string())? + .parse::() + .map_err(|_| "Topk query has an invalid `k` parameter".to_string()) + } + + /// Adds Prometheus's metric-name label value to a Topk output key. + fn prepend_metric_name(metric: &str, key: &mut KeyByLabelValues) { + key.labels.insert(0, metric.to_string()); + } + /// Executes the complete query pipeline: plan, execute, collect, and format. /// /// The two top-k flags are deliberately separate because the two engines @@ -1120,9 +1131,7 @@ impl SimpleEngine { .into_iter() .map(|(key_opt, value)| { let updated_key = key_opt.map(|mut key| { - let mut new_labels = vec![metric.to_string()]; - new_labels.extend(key.labels); - key.labels = new_labels; + Self::prepend_metric_name(metric, &mut key); key }); (updated_key, value) @@ -2024,12 +2033,7 @@ impl SimpleEngine { let topk_k: Option = if enable_topk_limiting && context.base.metadata.statistic_to_compute == Statistic::Topk { - context - .base - .metadata - .query_kwargs - .get("k") - .and_then(|s| s.parse::().ok()) + Some(Self::parse_topk_limit(&context.base.metadata.query_kwargs)?) } else { None }; @@ -2216,9 +2220,7 @@ impl SimpleEngine { // groups/samples survive. if enable_topk_formatting && context.base.metadata.statistic_to_compute == Statistic::Topk { for elem in results.values_mut() { - let mut new_labels = vec![context.base.metric.clone()]; - new_labels.extend(elem.labels.labels.clone()); - elem.labels.labels = new_labels; + Self::prepend_metric_name(&context.base.metric, &mut elem.labels); } } @@ -2226,6 +2228,28 @@ impl SimpleEngine { } } +#[cfg(test)] +mod topk_metadata_tests { + use super::SimpleEngine; + use std::collections::HashMap; + + #[test] + fn topk_limit_requires_a_parseable_k() { + assert_eq!( + SimpleEngine::parse_topk_limit(&HashMap::new()), + Err("Topk query is missing required `k` parameter".to_string()) + ); + assert_eq!( + SimpleEngine::parse_topk_limit(&HashMap::from([("k".to_string(), "nope".to_string())])), + Err("Topk query has an invalid `k` parameter".to_string()) + ); + assert_eq!( + SimpleEngine::parse_topk_limit(&HashMap::from([("k".to_string(), "3".to_string())])), + Ok(3) + ); + } +} + #[cfg(test)] mod range_query_tests { use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; 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 5c12c3f..c9711f3 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 @@ -612,16 +612,9 @@ mod tests { ); } - // RED test for PR #629 review finding 2 (see - // .design_docs/pr-629-review-findings-handoff.md): `apply_range_topk`'s - // `candidates.sort_by(|a, b| b.1.partial_cmp(&a.1)...)` (mod.rs) sorts by - // value only, with no tiebreak. `candidates`'s order comes from - // iterating `results: HashMap<...>`, whose iteration order is - // per-process randomized (SipHash) -- two groups tied at the k-th value - // boundary can keep different groups run to run. This asserts the - // *desired* deterministic tiebreak (ties broken by ascending label - // values, so "host-b" beats "host-c"), which the current unfixed sort - // has no mechanism to guarantee. + // Regression test for the deterministic tie-break in `apply_range_topk`. + // Candidates are ordered by value descending, then label values ascending, + // so a tie at the k-th boundary is stable despite HashMap iteration order. #[tokio::test(flavor = "multi_thread")] async fn topk_range_tie_break_is_deterministic_by_label_values() { let query = "topk(2, cpu_load)";