From cbeae1bc1c82288484f4d318a2e537b44d8925c3 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 23:53:20 -0400 Subject: [PATCH 1/4] feat(query-engine): computed top-k in range queries + stage-E equivalence tests (#581 stage E prep) Adds step-major topk ranking/truncation to execute_range_query_pipeline (previously range had no top-k support at all), wires it through PromQL's range call sites, and adds an instant/range equivalence test matrix across Tumbling/Sliding window shapes and SetAgg/DeltaSetAgg keys configs, ahead of stage E's full pipeline collapse. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/mod.rs | 118 +++- .../src/engines/simple_engine/promql.rs | 20 +- asap-query-engine/src/tests/mod.rs | 1 + ...stage_e_instant_range_equivalence_tests.rs | 614 ++++++++++++++++++ 4 files changed, 745 insertions(+), 8 deletions(-) create mode 100644 asap-query-engine/src/tests/stage_e_instant_range_equivalence_tests.rs diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 6dee86a..1faca13 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1500,10 +1500,20 @@ impl SimpleEngine { } } - /// Execute the range query pipeline + /// Execute the range query pipeline. + /// + /// `enable_topk_limiting`/`enable_topk_formatting` mirror + /// `execute_query_pipeline`'s flags of the same name (see that method's + /// doc comment) -- both no-ops unless + /// `context.base.metadata.statistic_to_compute == Statistic::Topk`. The + /// actual ranking/truncation is delegated to `apply_range_topk` below; + /// see its doc comment for why range's version can't just reuse + /// instant's `format_final_results` truncate-once shape. fn execute_range_query_pipeline( &self, context: &RangeQueryExecutionContext, + enable_topk_limiting: bool, + enable_topk_formatting: bool, ) -> Result, String> { use crate::engines::query_result::RangeVectorElement; use crate::engines::window_merger::create_window_merger; @@ -1829,8 +1839,110 @@ impl SimpleEngine { } } - // Convert to Vec - Ok(results.into_values().collect()) + Ok(self.apply_range_topk( + results, + &context.base.metadata.statistic_to_compute, + &context.base.metadata.query_kwargs, + &context.base.metric, + enable_topk_formatting, + enable_topk_limiting, + )) + } + + /// Applies PromQL top-k semantics to a range query's raw per-group + /// results. No-op unless `statistic == Statistic::Topk` (mirrors + /// `format_final_results`). + /// + /// This is deliberately NOT a straight port of `format_final_results` + /// (sort all groups once by value, then truncate to k): that shape only + /// works because instant queries have exactly one value per group. A + /// range query's `RangeVectorElement` carries many per-timestamp + /// samples, and real PromQL `topk(k, range_vector)` semantics rank + /// independently AT EACH timestamp -- the surviving key set can differ + /// from step to step. So this ranks/truncates per-timestamp + /// ("step-major"), across all groups, as its own pass over the + /// already-assembled results -- rather than restructuring the group-major + /// fetch/merge loop above into a step-major shape. Issue #581's own + /// scoping decided the fetch/merge loop itself becomes step-major only + /// as part of stage E, the full instant/range pipeline collapse (not + /// done here, deliberately -- this is stage-E prep). Doing the ranking + /// as a separate pass gets the same correctness (each timestamp's kept + /// set is decided across all groups, never one group at a time) without + /// front-running that larger, separately-staged restructure. + fn apply_range_topk( + &self, + mut results: HashMap, + statistic: &Statistic, + query_kwargs: &HashMap, + metric: &str, + enable_topk_formatting: bool, + enable_topk_limiting: bool, + ) -> Vec { + if *statistic != Statistic::Topk { + return results.into_values().collect(); + } + + // Limiting MUST run before formatting: it matches + // `kept_timestamps_by_key`'s keys (read from each element's + // `labels` field) against `results`' own HashMap keys via + // `retain`. Formatting rewrites `elem.labels` (the field) without + // touching the HashMap's outer key, so if formatting ran first the + // two would no longer agree and `retain` would drop every group. + if enable_topk_limiting { + if let Some(k) = query_kwargs.get("k").and_then(|s| s.parse::().ok()) { + use std::collections::HashSet; + + // Step-major ranking: group every group's samples by + // timestamp first, so each timestamp's top-k decision sees + // every group's value at that timestamp. + let mut by_timestamp: HashMap> = HashMap::new(); + for elem in results.values() { + for sample in &elem.samples { + by_timestamp + .entry(sample.timestamp) + .or_default() + .push((elem.labels.clone(), sample.value)); + } + } + + let mut kept_timestamps_by_key: HashMap> = + HashMap::new(); + for (timestamp, mut candidates) in by_timestamp { + candidates + .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + candidates.truncate(k); + for (key, _) in candidates { + kept_timestamps_by_key + .entry(key) + .or_default() + .insert(timestamp); + } + } + + results.retain(|key, _| kept_timestamps_by_key.contains_key(key)); + for elem in results.values_mut() { + let keep = &kept_timestamps_by_key[&elem.labels]; + elem.samples.retain(|s| keep.contains(&s.timestamp)); + } + // A group that made no timestamp's top-k has no samples left + // -- drop it entirely rather than emitting an empty series. + results.retain(|_, elem| !elem.samples.is_empty()); + } + } + + if enable_topk_formatting { + // Prepend metric name to each key's label values (PromQL shape), + // same rewrite as format_final_results does for instant. Safe to + // mutate `elem.labels` now -- nothing below matches it back + // against the HashMap's outer key. + for elem in results.values_mut() { + let mut new_labels = vec![metric.to_string()]; + new_labels.extend(elem.labels.labels.clone()); + elem.labels.labels = new_labels; + } + } + + results.into_values().collect() } } diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 8465b46..c0bd531 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -717,7 +717,10 @@ impl SimpleEngine { if let Some((scalar, vector_arm, scalar_on_left)) = detect_scalar_arm(lhs, rhs) { let (ctx, labels) = self.build_arm_range_context(vector_arm, start, end, step)?; - let results = self.execute_range_query_pipeline(&ctx).ok()?; + // (true, true): self-gated, same as instant's binary-arm call + // (evaluate_binary_arm) -- both flags are no-ops unless the arm's + // statistic is Topk. + let results = self.execute_range_query_pipeline(&ctx, true, true).ok()?; let combined: Vec = results .into_iter() .map(|mut elem| { @@ -745,8 +748,13 @@ impl SimpleEngine { if lhs_labels != rhs_labels { return None; } - let lhs_results = self.execute_range_query_pipeline(&lhs_ctx).ok()?; - let rhs_results = self.execute_range_query_pipeline(&rhs_ctx).ok()?; + // (true, true): self-gated, same rationale as the scalar-arm call above. + let lhs_results = self + .execute_range_query_pipeline(&lhs_ctx, true, true) + .ok()?; + let rhs_results = self + .execute_range_query_pipeline(&rhs_ctx, true, true) + .ok()?; // Build lookup: label_key -> {timestamp -> value} for rhs let mut rhs_map: HashMap> = HashMap::new(); @@ -1307,9 +1315,11 @@ impl SimpleEngine { let context = self.build_range_query_execution_context_from_parsed(&ast, &query, start, end, step)?; - // Execute range query pipeline + // Execute range query pipeline. (true, true): self-gated, same as + // instant's handle_query_promql -- both flags are no-ops unless this + // query's statistic is Topk. let results: Vec = self - .execute_range_query_pipeline(&context) + .execute_range_query_pipeline(&context, true, true) .map_err(|e| { warn!("Range query execution failed: {}", e); e diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index b4a5bbe..262119e 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -12,6 +12,7 @@ pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; pub mod range_query_arithmetic_tests; pub mod sql_pattern_matching_tests; +pub mod stage_e_instant_range_equivalence_tests; pub mod store_correctness_tests; pub mod structural_matching_tests; pub mod trait_design_tests; 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 new file mode 100644 index 0000000..fe25e52 --- /dev/null +++ b/asap-query-engine/src/tests/stage_e_instant_range_equivalence_tests.rs @@ -0,0 +1,614 @@ +//! Stage-E prep for #581 (unify PromQL instant/range fetch+merge paths): +//! equivalence tests between the instant pipeline (`handle_query_promql`) +//! and a single-step range pipeline (`handle_range_query_promql`), across +//! window-shape, dual-population, and keys-aggregator-kind axes. All of +//! stage E's blockers (B/#612, C/#596/#614, #608/#621) are closed on main as +//! of this writing, so every case here is expected to PASS -- there's no +//! known bug being pinned. A single-step range query and the equivalent +//! instant query at the same timestamp read and merge the exact same +//! underlying buckets, so they're compared for exact equality, not +//! approximate. +//! +//! Two corrections to the originally-scoped matrix, discovered while writing +//! these (both from reading `execute_range_query_pipeline` and its callers, +//! not from any doc): +//! +//! 1. A Sliding-window aggregation's query lookback is REQUIRED to equal its +//! `window_size_ms` -- an active `assert!` in `execute_range_query_pipeline` +//! (mod.rs, #608 review), not just a convention. So "Sliding, +//! window_size == lookback/2, /3, ..." isn't a constructible scenario -- +//! it would panic, not merge incorrectly. The window-shape axis below +//! instead varies Tumbling's lookback-to-bucket-width ratio (1/2/3, via +//! `sum_over_time(metric[Ns])` selector width) and adds one genuinely +//! overlapping Sliding case (`window_size_ms > slide_interval_ms`, +//! lookback == window_size_ms as required). +//! +//! 2. `SetAggregator`/`DeltaSetAggregator` are keys-side (dual-population) +//! aggregation types in this codebase -- never a general value-aggregation +//! choice (see every existing dual-population fixture in +//! `native_range_query_tests.rs`). And `DeltaSetAggregator`'s +//! Tumbling-only restriction (#588/#606) is about ITS OWN window_type, +//! independent of whatever window_type the paired VALUE aggregation uses +//! -- `RangeQueryExecutionContext::keys_window_type`'s doc comment says as +//! much ("can legitimately differ from window_type"). So there's no +//! Sliding-value-excludes-DeltaSetAgg-keys interaction to exclude: keys +//! config is fixed at Tumbling for DeltaSetAgg regardless of the value +//! shape under test, and every {value shape} x {no keys / SetAgg keys / +//! DeltaSetAgg keys} x {instant / range} cell is constructible. 4 value +//! shapes x 3 accumulator configs x 2 paths = 24 cases, no exclusions. + +#[cfg(test)] +mod tests { + use crate::data_model::{ + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + KeyByLabelValues, PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, + SchemaConfig, StreamingConfig, WindowType, + }; + use crate::engines::query_result::QueryResult; + use crate::engines::simple_engine::SimpleEngine; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::precompute_operators::{ + CountMinSketchWithHeapAccumulator, DeltaSetAggregatorAccumulator, SetAggregatorAccumulator, + }; + use crate::stores::simple_map_store::SimpleMapStore; + use crate::stores::Store; + use crate::AggregateCore; + use promql_utilities::data_model::KeyByLabelNames; + use std::collections::HashMap; + use std::sync::Arc; + + /// One value-side window shape under test. `query` is a complete + /// `sum_over_time(cpu_load[Ns])` selector sized to exercise this shape's + /// lookback-to-bucket-width ratio (Tumbling) or to satisfy the + /// Sliding lookback==window_size_ms invariant. + struct ValueShape { + window_type: WindowType, + window_size_ms: u64, + slide_interval_ms: u64, + query: &'static str, + /// Output timestamp (seconds) both the instant query and the + /// single-step range query are evaluated at. + query_time_s: f64, + } + + const VALUE_SHAPES: [ValueShape; 4] = [ + ValueShape { + window_type: WindowType::Tumbling, + window_size_ms: 1000, + slide_interval_ms: 1000, + query: "sum_over_time(cpu_load[1s])", + query_time_s: 3.0, + }, + ValueShape { + window_type: WindowType::Tumbling, + window_size_ms: 1000, + slide_interval_ms: 1000, + query: "sum_over_time(cpu_load[2s])", + query_time_s: 3.0, + }, + ValueShape { + window_type: WindowType::Tumbling, + window_size_ms: 1000, + slide_interval_ms: 1000, + query: "sum_over_time(cpu_load[3s])", + query_time_s: 3.0, + }, + ValueShape { + window_type: WindowType::Sliding, + window_size_ms: 2000, + slide_interval_ms: 1000, + query: "sum_over_time(cpu_load[2s])", + query_time_s: 3.0, + }, + ]; + + /// Which keys/labels accumulator (if any) is paired with the value + /// aggregation for a given matrix cell. + enum KeysConfig { + None, + SetAgg, + DeltaSetAgg, + } + + /// Builds an engine with one value aggregation (id=1, Sum, shaped by + /// `shape`) covering timestamps 1000/2000/3000 for group "host-a", and + /// optionally a second keys aggregation (id=2) that resolves the same + /// group's label key. `keys` is fixed at Tumbling, window_size_ms == + /// slide_interval_ms == 1000, with one bucket at [2000,3000) -- SetAgg's + /// instant window is [end-window_size, end] so this always resolves at + /// query_time_s=3.0 regardless of the value shape's own ratio; DeltaSetAgg's + /// instant window is [0, end] so it finds this (and would find any + /// earlier bucket) regardless too. + fn build_engine(shape: &ValueShape, keys: KeysConfig) -> SimpleEngine { + let grouping_labels = vec!["host".to_string()]; + let host_a = Some(KeyByLabelValues { + labels: vec!["host-a".to_string()], + }); + + let mut aggregation_configs = HashMap::new(); + aggregation_configs.insert( + 1u64, + AggregationConfig { + aggregation_id: 1, + aggregation_type: AggregationType::Sum, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::new(grouping_labels.clone()), + aggregated_labels: KeyByLabelNames::empty(), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: shape.window_size_ms, + slide_interval_ms: shape.slide_interval_ms, + window_type: shape.window_type, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: "cpu_load".to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + + let mut query_config = QueryConfig::new(shape.query.to_string()) + .add_aggregation(AggregationReference::new(1, None)); + + if !matches!(keys, KeysConfig::None) { + let key_agg_type = match keys { + KeysConfig::SetAgg => AggregationType::SetAggregator, + KeysConfig::DeltaSetAgg => AggregationType::DeltaSetAggregator, + KeysConfig::None => unreachable!(), + }; + aggregation_configs.insert( + 2u64, + AggregationConfig { + aggregation_id: 2, + aggregation_type: key_agg_type, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::new(grouping_labels.clone()), + aggregated_labels: KeyByLabelNames::new(vec!["host".to_string()]), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: 1000, + slide_interval_ms: 1000, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: "cpu_load".to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + query_config = query_config.add_aggregation(AggregationReference::new(2, None)); + } + + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs, + }); + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + + // Value data: three consecutive 1000ms-wide panes. For Tumbling + // shapes these ARE the stored buckets directly. For the Sliding + // shape, worker.rs-style pre-merging isn't done here -- instead we + // insert the two window_size_ms=2000-wide merged buckets a real + // Sliding worker would have produced from these same panes + // ([0,2000) from panes at 1000/2000, [1000,3000) from panes at + // 2000/3000), matching `create_engine_multi_timestamp_with_window`'s + // own merge logic but explicit here since this factory also needs a + // second (keys) aggregation those single-aggregation factories don't + // support. + let value_buckets: Vec<(u64, u64, f64)> = match shape.window_type { + WindowType::Tumbling => vec![(0, 1000, 1.0), (1000, 2000, 10.0), (2000, 3000, 100.0)], + WindowType::Sliding => vec![(0, 2000, 11.0), (1000, 3000, 110.0)], + }; + for (start, end, value) in value_buckets { + let output = PrecomputedOutput::new(start, end, host_a.clone(), 1); + store + .insert_precomputed_output(output, Box::new(SumAccumulator::with_sum(value))) + .unwrap(); + } + + if !matches!(keys, KeysConfig::None) { + let acc: Box = match keys { + KeysConfig::SetAgg => { + let mut a = SetAggregatorAccumulator::new(); + a.add_key(KeyByLabelValues { + labels: vec!["host-a".to_string()], + }); + Box::new(a) + } + KeysConfig::DeltaSetAgg => { + let mut a = DeltaSetAggregatorAccumulator::new(); + a.add_key(KeyByLabelValues { + labels: vec!["host-a".to_string()], + }); + Box::new(a) + } + KeysConfig::None => unreachable!(), + }; + let output = PrecomputedOutput::new(2000, 3000, host_a.clone(), 2); + store.insert_precomputed_output(output, acc).unwrap(); + } + + let promql_schema = PromQLSchema::new().add_metric( + "cpu_load".to_string(), + KeyByLabelNames::new(grouping_labels), + ); + let inference_config = InferenceConfig { + schema: SchemaConfig::PromQL(promql_schema), + query_configs: vec![query_config], + cleanup_policy: CleanupPolicy::NoCleanup, + }; + + SimpleEngine::new( + store, + inference_config, + streaming_config, + 1000, + QueryLanguage::promql, + ) + } + + /// Runs `shape`'s query as both an instant query (at `shape.query_time_s`) + /// and a single-step range query (start == end == `shape.query_time_s`), + /// and asserts they produce the exact same (label, value) set. Panics + /// with a diagnostic including `case_name` on mismatch -- this is a + /// correctness check, not a fixture to adjust if it fails. + fn assert_instant_range_equivalent(shape: &ValueShape, keys: KeysConfig, case_name: &str) { + let engine = build_engine(shape, keys); + + let (_, instant_result) = engine + .handle_query_promql(shape.query.to_string(), shape.query_time_s) + .unwrap_or_else(|| panic!("{case_name}: instant query returned None")); + let mut instant_pairs: Vec<(Vec, f64)> = match instant_result { + QueryResult::Vector(v) => v + .values + .into_iter() + .map(|e| (e.labels.labels, e.value)) + .collect(), + QueryResult::Matrix(_) => panic!("{case_name}: instant query returned a Matrix"), + }; + + // start must be strictly < end (validate_range_query_params); end is + // set half a step past start so the per-step loop + // (`while current_time <= end_ms`) still fires exactly once, at + // current_time == start == shape.query_time_s -- the same instant + // the instant query above ran at. + let (_, range_result) = engine + .handle_range_query_promql( + shape.query.to_string(), + shape.query_time_s, + shape.query_time_s + 0.5, + 1.0, + ) + .unwrap_or_else(|| panic!("{case_name}: range query returned None")); + let mut range_pairs: Vec<(Vec, f64)> = match range_result { + QueryResult::Matrix(m) => m + .values + .into_iter() + .map(|e| { + assert_eq!( + e.samples.len(), + 1, + "{case_name}: single-step range query produced {} samples for {:?}, expected exactly 1", + e.samples.len(), + e.labels.labels + ); + (e.labels.labels, e.samples[0].value) + }) + .collect(), + QueryResult::Vector(_) => panic!("{case_name}: range query returned a Vector"), + }; + + instant_pairs.sort_by(|a, b| a.0.cmp(&b.0)); + range_pairs.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!( + instant_pairs, range_pairs, + "{case_name}: instant and single-step range results diverge" + ); + } + + macro_rules! equivalence_test { + ($test_name:ident, $shape_index:expr, $keys:expr) => { + #[tokio::test(flavor = "multi_thread")] + async fn $test_name() { + assert_instant_range_equivalent( + &VALUE_SHAPES[$shape_index], + $keys, + stringify!($test_name), + ); + } + }; + } + + // ── single accumulator (no keys_query) ────────────────────────────── + equivalence_test!(equiv_single_tumbling_ratio1, 0, KeysConfig::None); + equivalence_test!(equiv_single_tumbling_ratio2, 1, KeysConfig::None); + equivalence_test!(equiv_single_tumbling_ratio3, 2, KeysConfig::None); + equivalence_test!(equiv_single_sliding_overlapping, 3, KeysConfig::None); + + // ── dual accumulator, SetAggregator keys ──────────────────────────── + equivalence_test!(equiv_dual_setagg_tumbling_ratio1, 0, KeysConfig::SetAgg); + equivalence_test!(equiv_dual_setagg_tumbling_ratio2, 1, KeysConfig::SetAgg); + equivalence_test!(equiv_dual_setagg_tumbling_ratio3, 2, KeysConfig::SetAgg); + equivalence_test!(equiv_dual_setagg_sliding_overlapping, 3, KeysConfig::SetAgg); + + // ── dual accumulator, DeltaSetAggregator keys ─────────────────────── + equivalence_test!( + equiv_dual_deltasetagg_tumbling_ratio1, + 0, + KeysConfig::DeltaSetAgg + ); + equivalence_test!( + equiv_dual_deltasetagg_tumbling_ratio2, + 1, + KeysConfig::DeltaSetAgg + ); + equivalence_test!( + equiv_dual_deltasetagg_tumbling_ratio3, + 2, + KeysConfig::DeltaSetAgg + ); + equivalence_test!( + equiv_dual_deltasetagg_sliding_overlapping, + 3, + KeysConfig::DeltaSetAgg + ); + + // ════════════════════════════════════════════════════════════════════ + // ── Computed top-k in range queries (new plumbing, this PR) ───────── + // ════════════════════════════════════════════════════════════════════ + // + // IMPORTANT, discovered while writing these: `Statistic::Topk` as a + // per-group value query (`AggregateCore::query_statistic`) is only + // answerable by accumulator types that special-case or ignore it -- + // `CountMinSketchWithHeapAccumulator::query` ignores the `_statistic` + // argument entirely and always returns `query_key`. + // `SumAccumulator::query` strictly matches on `Statistic::Sum | Count` + // and returns `Err("Unsupported statistic")` for anything else, + // including Topk. An earlier version of these tests tried a per-group + // plain `SumAccumulator` (`get_keys() == None`, to force resolution + // through the fallback_key path rather than self-keyed expansion) to + // isolate "ranking/truncation across independently-resolved groups" + // from "self-keyed expansion within one group" -- both instant AND + // range returned zero results, identically, because the per-group value + // query fails before ranking ever runs. This is pre-existing on `main`, + // unrelated to this PR's range changes, and it means "topk over an + // arbitrary non-self-keyed expression" (e.g. `topk(5, rate(foo[5m]))`) + // is not a supported query shape ANYWHERE in this codebase today, not + // just missing in range -- the scope of "topk in range" that's + // actually constructible is narrower than originally framed. + // + // So these tests use the one accumulator shape that both (a) actually + // answers Topk queries and (b) is how this codebase's own working topk + // fixture (`build_topk_engine` in promql.rs, `topk_pipeline_tests`) is + // built: one ungrouped `CountMinSketchWithHeapAccumulator` per output + // timestamp. `range_query_self_keyed_topk_expands_*` + // (native_range_query_tests.rs, #595) already covers "expansion + // resolves correctly within one step" -- these add the genuinely new + // case #595 didn't: MULTIPLE steps, where the top-k set changes from + // step to step (step-major ranking, this PR's new plumbing) and, + // separately, that a range query now truncates to k at all (before + // this PR, `execute_range_query_pipeline` had no truncation step, so + // every self-keyed candidate below `heap_size` would have come back + // untruncated). + + /// (bucket_start_ms, bucket_end_ms, [(host, value), ...]). + type TopkBucket<'a> = (u64, u64, &'a [(&'a str, f64)]); + + fn build_range_topk_engine(query: &str, buckets: &[TopkBucket]) -> SimpleEngine { + let mut aggregation_configs = HashMap::new(); + aggregation_configs.insert( + 1u64, + AggregationConfig { + aggregation_id: 1, + aggregation_type: AggregationType::CountMinSketchWithHeap, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::empty(), + aggregated_labels: KeyByLabelNames::new(vec!["host".to_string()]), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: 1000, + slide_interval_ms: 1000, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: "cpu_load".to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs, + }); + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + for (start, end, candidates) in buckets { + let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); + for (host, value) in *candidates { + sketch.inner.update(host, *value); + } + let output = PrecomputedOutput::new(*start, *end, None, 1); + store + .insert_precomputed_output(output, Box::new(sketch)) + .unwrap(); + } + let promql_schema = PromQLSchema::new().add_metric( + "cpu_load".to_string(), + KeyByLabelNames::new(vec!["host".to_string()]), + ); + let query_config = + QueryConfig::new(query.to_string()).add_aggregation(AggregationReference::new(1, None)); + let inference_config = InferenceConfig { + schema: SchemaConfig::PromQL(promql_schema), + query_configs: vec![query_config], + cleanup_policy: CleanupPolicy::NoCleanup, + }; + SimpleEngine::new( + store, + inference_config, + streaming_config, + 1000, + QueryLanguage::promql, + ) + } + + fn topk_samples_for( + elements: &[crate::engines::query_result::RangeVectorElement], + host: &str, + ) -> Vec<(u64, f64)> { + elements + .iter() + .find(|e| e.labels.labels.contains(&host.to_string())) + .map(|e| e.samples.iter().map(|s| (s.timestamp, s.value)).collect()) + .unwrap_or_default() + } + + /// The ranked-and-truncated set of surviving hosts must differ from + /// step to step, ranked independently at each timestamp (step-major) -- + /// NOT a single global ranking applied to every step (group-major would + /// wrongly keep/drop the same hosts at every timestamp regardless of + /// that timestamp's actual values). Before this PR, range had no + /// truncation at all, so this would have returned all 3 hosts at both + /// steps. + /// + /// t=1000: host-a=100, host-b=50, host-c=10 -> top2 = {a, b}, c dropped. + /// t=2000: host-a=5, host-b=50, host-c=100 -> top2 = {b, c}, a dropped. + #[tokio::test(flavor = "multi_thread")] + async fn topk_range_step_major_ranking_differs_per_step() { + let query = "topk(2, cpu_load)"; + let engine = build_range_topk_engine( + query, + &[ + ( + 0, + 1000, + &[("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + ), + ( + 1000, + 2000, + &[("host-a", 5.0), ("host-b", 50.0), ("host-c", 100.0)], + ), + ], + ); + + let (_, result) = engine + .handle_range_query_promql(query.to_string(), 1.0, 2.5, 1.0) + .expect("range topk query failed"); + let elements = match result { + QueryResult::Matrix(m) => m.values, + QueryResult::Vector(_) => panic!("expected a Matrix"), + }; + + assert_eq!( + topk_samples_for(&elements, "host-a"), + vec![(1000, 100.0)], + "host-a should survive only t=1000's top-2 (100 > 50,10), not t=2000's (5 is last of 3)" + ); + assert_eq!( + topk_samples_for(&elements, "host-b"), + vec![(1000, 50.0), (2000, 50.0)], + "host-b (50) is top-2 at both steps (t=1000: 100,50 beat 10; t=2000: 100,50 beat 5)" + ); + assert_eq!( + topk_samples_for(&elements, "host-c"), + vec![(2000, 100.0)], + "host-c should survive only t=2000's top-2 (100 > 50,5), not t=1000's (10 is last of 3)" + ); + } + + /// A single-step range topk query must match the equivalent instant + /// topk query exactly -- both gated identically on + /// `Statistic::Topk` + `query_kwargs["k"]`, and both now truncate. + #[tokio::test(flavor = "multi_thread")] + async fn topk_range_single_step_matches_instant() { + let query = "topk(2, cpu_load)"; + let engine = build_range_topk_engine( + query, + &[( + 0, + 1000, + &[("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + )], + ); + + let (_, instant_result) = engine + .handle_query_promql(query.to_string(), 1.0) + .expect("instant topk query failed"); + let mut instant_pairs: Vec<(Vec, f64)> = match instant_result { + QueryResult::Vector(v) => v + .values + .into_iter() + .map(|e| (e.labels.labels, e.value)) + .collect(), + QueryResult::Matrix(_) => panic!("expected a Vector"), + }; + + let (_, range_result) = engine + .handle_range_query_promql(query.to_string(), 1.0, 1.5, 1.0) + .expect("range topk query failed"); + let mut range_pairs: Vec<(Vec, f64)> = match range_result { + QueryResult::Matrix(m) => m + .values + .into_iter() + .map(|e| { + assert_eq!(e.samples.len(), 1); + (e.labels.labels, e.samples[0].value) + }) + .collect(), + QueryResult::Vector(_) => panic!("expected a Matrix"), + }; + + assert_eq!( + instant_pairs.len(), + 2, + "topk(2, ...) must truncate to 2 results" + ); + instant_pairs.sort_by(|a, b| a.0.cmp(&b.0)); + range_pairs.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!(instant_pairs, range_pairs); + } + + /// The new range plumbing must not interfere with the existing + /// self-keyed range topk path (#595): a single-step range query over a + /// dual-population self-keyed group, exercised the same way + /// `native_range_query_tests.rs`'s `range_query_self_keyed_topk_expands_*` + /// tests do, must still resolve and expand correctly. + #[tokio::test(flavor = "multi_thread")] + async fn topk_range_does_not_break_existing_self_keyed_expansion() { + let query = "topk(5, cpu_load)"; + let engine = + build_range_topk_engine(query, &[(0, 1000, &[("host-a", 30.0), ("host-b", 20.0)])]); + + let (_, result) = engine + .handle_range_query_promql(query.to_string(), 1.0, 1.5, 1.0) + .expect("range topk query failed"); + let elements = match result { + QueryResult::Matrix(m) => m.values, + QueryResult::Vector(_) => panic!("expected a Matrix"), + }; + + assert_eq!( + topk_samples_for(&elements, "host-a"), + vec![(1000, 30.0)], + "self-keyed expansion (#595) must still resolve host-a" + ); + assert_eq!( + topk_samples_for(&elements, "host-b"), + vec![(1000, 20.0)], + "self-keyed expansion (#595) must still resolve host-b" + ); + } +} From e76a147cc2a2ee320714d690017a751a919f0451 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 26 Aug 2026 00:01:57 -0400 Subject: [PATCH 2/4] fix(query-engine): drop dead-code retain in apply_range_topk (roborev #23) kept_timestamps_by_key's entries are only ever inserted alongside a timestamp drawn from that same element's own samples, so the per-sample filter can never leave a surviving key's samples empty -- the trailing retain was unreachable dead code. Co-Authored-By: Claude Sonnet 5 --- asap-query-engine/src/engines/simple_engine/mod.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 1faca13..e6ca7e9 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1921,12 +1921,14 @@ impl SimpleEngine { results.retain(|key, _| kept_timestamps_by_key.contains_key(key)); for elem in results.values_mut() { + // `keep` is built only from timestamps that already + // appear in this same element's `samples` (see the + // `by_timestamp` loop above), and is non-empty for every + // key that survives the `retain` just above -- so this + // filter can never leave `elem.samples` empty. let keep = &kept_timestamps_by_key[&elem.labels]; elem.samples.retain(|s| keep.contains(&s.timestamp)); } - // A group that made no timestamp's top-k has no samples left - // -- drop it entirely rather than emitting an empty series. - results.retain(|_, elem| !elem.samples.is_empty()); } } From c73d52ee5427c61f408f82641c493a58922915af Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 26 Aug 2026 08:55:42 -0400 Subject: [PATCH 3/4] test(query-engine): pin topk binary-expr label bugs found via PR #629 Finding 1 topk(...) as one arm of a binary expression is broken two different ways, neither of which is what Finding 1's review comment described nor fixable as part of #629 -- both are pre-existing/orthogonal and tracked in #631 instead. One test pins the current (surprising) None-return behavior; the other reproduces the join-corruption bug Finding 1 actually describes, and is #[ignore]d since it isn't fixed here. Co-Authored-By: Claude Sonnet 5 --- .../src/tests/range_query_arithmetic_tests.rs | 293 ++++++++++++++++++ 1 file changed, 293 insertions(+) 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 cdf6720..e1bc962 100644 --- a/asap-query-engine/src/tests/range_query_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/range_query_arithmetic_tests.rs @@ -19,6 +19,7 @@ mod tests { use crate::engines::query_result::{QueryResult, RangeVectorElement}; use crate::engines::simple_engine::SimpleEngine; use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::precompute_operators::CountMinSketchWithHeapAccumulator; use crate::stores::simple_map_store::SimpleMapStore; use crate::stores::Store; use crate::AggregateCore; @@ -286,4 +287,296 @@ mod tests { values coincide, got {result:?}" ); } + + /// Builds a SimpleEngine with one self-keyed topk-capable metric + /// (`metric_a`, `CountMinSketchWithHeap`, one ungrouped sketch per + /// bucket) and one plain grouped-sum metric (`metric_b`, `SumAccumulator` + /// per host), so a mixed `topk(k, metric_a) OP sum(metric_b) by (host)` + /// range query is constructible. Mirrors `build_range_topk_engine` in + /// `stage_e_instant_range_equivalence_tests.rs` for the topk side, and + /// `create_range_engine_two_metrics`'s per-host Sum buckets for the + /// plain side. + fn build_range_topk_plus_plain_engine( + topk_query: &str, + plain_query: &str, + topk_candidates: &[(&str, f64)], + plain_values: &[(&str, f64)], + ) -> SimpleEngine { + let mut aggregation_configs = HashMap::new(); + aggregation_configs.insert( + 1u64, + AggregationConfig { + aggregation_id: 1, + aggregation_type: AggregationType::CountMinSketchWithHeap, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::empty(), + aggregated_labels: KeyByLabelNames::new(vec!["host".to_string()]), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: WINDOW_MS, + slide_interval_ms: WINDOW_MS, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: "metric_a".to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + aggregation_configs.insert( + 2u64, + AggregationConfig { + aggregation_id: 2, + aggregation_type: AggregationType::Sum, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::new(vec!["host".to_string()]), + aggregated_labels: KeyByLabelNames::empty(), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: WINDOW_MS, + slide_interval_ms: WINDOW_MS, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: "metric_b".to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs, + }); + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + + let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); + for (host, value) in topk_candidates { + sketch.inner.update(host, *value); + } + let topk_output = PrecomputedOutput::new(0, WINDOW_MS, None, 1); + store + .insert_precomputed_output(topk_output, Box::new(sketch)) + .unwrap(); + + for (host, value) in plain_values { + let key = KeyByLabelValues { + labels: vec![host.to_string()], + }; + let plain_output = PrecomputedOutput::new(0, WINDOW_MS, Some(key), 2); + store + .insert_precomputed_output( + plain_output, + Box::new(SumAccumulator::with_sum(*value)) as Box, + ) + .unwrap(); + } + + let promql_schema = PromQLSchema::new() + .add_metric( + "metric_a".to_string(), + KeyByLabelNames::new(vec!["host".to_string()]), + ) + .add_metric( + "metric_b".to_string(), + KeyByLabelNames::new(vec!["host".to_string()]), + ); + let inference_config = InferenceConfig { + schema: SchemaConfig::PromQL(promql_schema), + query_configs: vec![ + QueryConfig::new(topk_query.to_string()) + .add_aggregation(AggregationReference::new(1, None)), + QueryConfig::new(plain_query.to_string()) + .add_aggregation(AggregationReference::new(2, None)), + ], + cleanup_policy: CleanupPolicy::NoCleanup, + }; + + SimpleEngine::new( + store, + inference_config, + streaming_config, + WINDOW_MS, + QueryLanguage::promql, + ) + } + + // Documents a PRE-EXISTING bug, unrelated to PR #629, NOT one of the 4 + // review findings being addressed here (see + // .design_docs/pr-629-review-findings-handoff.md, Finding 1). Confirmed + // via `git log -L` that `build_promql_execution_context_tail` has + // unconditionally prepended `"__name__"` to a Topk arm's *label names* + // (promql.rs, the `if statistic_to_compute == Statistic::Topk` block) + // since commit 9ac794c ("simple engine split by language #284"), long + // before #629. Because a plain (non-topk) arm never gets that prepend, + // `handle_binary_expr_range_promql`'s `lhs_labels != rhs_labels` guard + // (label *names*, checked before any join) rejects EVERY + // `topk(...) OP plain_metric` binary expression outright, on both the + // range and instant paths identically -- the query never reaches the + // join code at all, let alone the `apply_range_topk` formatting-mutates + // `elem.labels` *values* bug Finding 1 actually describes. This is + // asserted here only so the (surprising) current behavior is pinned; + // fixing it is out of scope for PR #629's review comments -- tracked as + // its own issue, #631, alongside the topk+topk repro below. + #[tokio::test(flavor = "multi_thread")] + async fn test_range_vector_vector_topk_lhs_plus_plain_rhs_returns_none() { + let query = "topk(2, metric_a) + sum(metric_b) by (host)"; + let engine = build_range_topk_plus_plain_engine( + "topk(2, metric_a)", + "sum(metric_b) by (host)", + &[("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + &[("host-a", 1000.0), ("host-b", 2000.0), ("host-c", 3000.0)], + ); + + let result = engine.handle_range_query_promql(query.to_string(), 1.0, 2.0, 1.0); + assert!( + result.is_none(), + "pre-existing __name__ label-name mismatch (predates #629) should reject this \ + query outright, got {result:?}" + ); + } + + /// Builds a SimpleEngine with two independent self-keyed topk-capable + /// metrics (`metric_a`, `metric_b`, both `CountMinSketchWithHeap`, one + /// ungrouped sketch per bucket each), so a + /// `topk(k1, metric_a) OP topk(k2, metric_b)` range query is + /// constructible. Unlike the topk+plain mix + /// (`build_range_topk_plus_plain_engine`), both arms get `"__name__"` + /// prepended to their label *names* identically, so this shape actually + /// reaches `handle_binary_expr_range_promql`'s vector-vector join -- + /// this is the shape PR #629 review Finding 1's `apply_range_topk` + /// formatting-before-join bug is reachable through. + fn build_range_two_topk_engine( + query_a: &str, + query_b: &str, + candidates_a: &[(&str, f64)], + candidates_b: &[(&str, f64)], + ) -> SimpleEngine { + let mut aggregation_configs = HashMap::new(); + for (id, metric) in [(1u64, "metric_a"), (2u64, "metric_b")] { + aggregation_configs.insert( + id, + AggregationConfig { + aggregation_id: id, + aggregation_type: AggregationType::CountMinSketchWithHeap, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::empty(), + aggregated_labels: KeyByLabelNames::new(vec!["host".to_string()]), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: WINDOW_MS, + slide_interval_ms: WINDOW_MS, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: metric.to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + } + + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs, + }); + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + + 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); + } + let output = PrecomputedOutput::new(0, WINDOW_MS, None, agg_id); + store + .insert_precomputed_output(output, Box::new(sketch)) + .unwrap(); + } + + let promql_schema = PromQLSchema::new() + .add_metric( + "metric_a".to_string(), + KeyByLabelNames::new(vec!["host".to_string()]), + ) + .add_metric( + "metric_b".to_string(), + KeyByLabelNames::new(vec!["host".to_string()]), + ); + let inference_config = InferenceConfig { + schema: SchemaConfig::PromQL(promql_schema), + query_configs: vec![ + QueryConfig::new(query_a.to_string()) + .add_aggregation(AggregationReference::new(1, None)), + QueryConfig::new(query_b.to_string()) + .add_aggregation(AggregationReference::new(2, None)), + ], + cleanup_policy: CleanupPolicy::NoCleanup, + }; + + SimpleEngine::new( + store, + inference_config, + streaming_config, + WINDOW_MS, + QueryLanguage::promql, + ) + } + + // RED test for PR #629 review finding 1 (see + // .design_docs/pr-629-review-findings-handoff.md): `apply_range_topk`'s + // formatting step (`enable_topk_formatting=true`) prepends EACH arm's + // own metric name to `elem.labels` before the vector-vector join. Two + // *different* topk metrics joined by a shared grouping label ("host") + // pass the earlier label-*names* guard (both get `"__name__"` + // prepended identically), so this reaches the join -- but the join + // then compares `["metric_a", host]`-shaped values against + // `["metric_b", host]`-shaped values, which never match regardless of + // whether the host itself is common to both topk's surviving sets. + // Real PromQL vector matching ignores `__name__`/joins by the shared + // label ("host") alone, so this should succeed wherever both topks kept + // that host -- the premature per-arm metric-name prepend breaks that. + // + // Tracked in #631, not fixed as part of PR #629 -- ignored so this RED + // repro doesn't fail this PR's test suite. Real Prometheus semantics for + // this query shape haven't been confirmed yet either. + #[tokio::test(flavor = "multi_thread")] + #[ignore = "tracked in #631, not part of PR #629's scope"] + async fn test_range_vector_vector_topk_lhs_topk_rhs() { + // topk(2, metric_a): host-a=100, host-b=50 survive; host-c=10 dropped. + // topk(2, metric_b): host-b=200, host-c=300 survive; host-a=5 dropped. + // Only host-b survives both topks -> expected combined: host-b = 250. + let query = "topk(2, metric_a) + topk(2, metric_b)"; + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + &[("host-a", 5.0), ("host-b", 200.0), ("host-c", 300.0)], + ); + + let result = engine.handle_range_query_promql(query.to_string(), 1.0, 2.0, 1.0); + let (_, qr) = result.expect("Expected result for topk/topk range query"); + let elements = matrix_values(qr); + + assert_eq!( + elements.len(), + 1, + "Expected only host-b (present in both topks' surviving sets), got {elements:?}" + ); + assert!(elements[0].labels.labels.contains(&"host-b".to_string())); + assert_eq!(elements[0].samples.len(), 1); + assert!((elements[0].samples[0].value - 250.0).abs() < 1e-10); + } } From 8551d2fd5a52a70869feee8071edbe42fe186e61 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 26 Aug 2026 08:57:45 -0400 Subject: [PATCH 4/4] fix(query-engine): deterministic range topk tie-break, fewer label clones Addresses PR #629 review findings 2-4 (mod.rs::apply_range_topk): - Finding 2: candidates.sort_by was value-only with no tiebreak, so groups tied at the k-th value boundary kept a different survivor run to run (HashMap iteration order is randomized per-process). Confirmed via a flaky RED test (4/5 pass rate) before adding a label-values tiebreak; stable across 20+ runs after. - Finding 3: folded into the same restructure -- index each group once instead of cloning its label vector per (group, timestamp) sample (G clones instead of G*T). - Finding 4: documented why range has no observable (false, true) case for enable_topk_limiting/enable_topk_formatting, unlike instant's always-sort-when-Topk behavior. Finding 1 is not addressed here -- tracked separately in #631. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/mod.rs | 54 +++++++++++++++---- ...stage_e_instant_range_equivalence_tests.rs | 50 +++++++++++++++++ 2 files changed, 93 insertions(+), 11 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index e6ca7e9..9c63649 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1869,6 +1869,18 @@ impl SimpleEngine { /// as a separate pass gets the same correctness (each timestamp's kept /// set is decided across all groups, never one group at a time) without /// front-running that larger, separately-staged restructure. + /// `enable_topk_limiting` and `enable_topk_formatting` are independent + /// flags, mirroring instant's `execute_query_pipeline` contract -- but + /// unlike instant, range has no `(false, true)`-observable case. Instant + /// always sorts Topk results when formatting regardless of limiting, + /// because it returns a flat `Vec` where sort order is part of the + /// output. Range returns a `HashMap` (this function's `results`) whose + /// iteration order was never meaningful, and each surviving group carries + /// many per-timestamp samples rather than one value to sort the outer + /// collection by -- so skipping the ranking block when + /// `enable_topk_limiting` is false has no observable effect here beyond + /// formatting, even though every current call site passes both flags + /// together and never actually exercises `(false, true)`. fn apply_range_topk( &self, mut results: HashMap, @@ -1892,41 +1904,61 @@ impl SimpleEngine { if let Some(k) = query_kwargs.get("k").and_then(|s| s.parse::().ok()) { use std::collections::HashSet; + // Index each group once (G clones total) instead of cloning + // its label vector per (group, timestamp) sample -- G*T + // clones otherwise, for G groups over T steps. + let index_keys: Vec = results.keys().cloned().collect(); + let key_to_idx: HashMap = index_keys + .iter() + .cloned() + .enumerate() + .map(|(i, key)| (key, i)) + .collect(); + // Step-major ranking: group every group's samples by // timestamp first, so each timestamp's top-k decision sees // every group's value at that timestamp. - let mut by_timestamp: HashMap> = HashMap::new(); + let mut by_timestamp: HashMap> = HashMap::new(); for elem in results.values() { + let idx = key_to_idx[&elem.labels]; for sample in &elem.samples { by_timestamp .entry(sample.timestamp) .or_default() - .push((elem.labels.clone(), sample.value)); + .push((idx, sample.value)); } } - let mut kept_timestamps_by_key: HashMap> = - HashMap::new(); + let mut kept_timestamps_by_idx: HashMap> = HashMap::new(); for (timestamp, mut candidates) in by_timestamp { - candidates - .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + // Tiebreak on label values: `candidates`'s order comes + // from iterating `results`, a HashMap, whose iteration + // order is randomized per-process -- without this, + // groups tied at the k-th value boundary would keep + // different survivors run to run. + candidates.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| index_keys[a.0].labels.cmp(&index_keys[b.0].labels)) + }); candidates.truncate(k); - for (key, _) in candidates { - kept_timestamps_by_key - .entry(key) + for (idx, _) in candidates { + kept_timestamps_by_idx + .entry(idx) .or_default() .insert(timestamp); } } - results.retain(|key, _| kept_timestamps_by_key.contains_key(key)); + results.retain(|key, _| kept_timestamps_by_idx.contains_key(&key_to_idx[key])); for elem in results.values_mut() { + let idx = key_to_idx[&elem.labels]; // `keep` is built only from timestamps that already // appear in this same element's `samples` (see the // `by_timestamp` loop above), and is non-empty for every // key that survives the `retain` just above -- so this // filter can never leave `elem.samples` empty. - let keep = &kept_timestamps_by_key[&elem.labels]; + let keep = &kept_timestamps_by_idx[&idx]; elem.samples.retain(|s| keep.contains(&s.timestamp)); } } 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 fe25e52..5c12c3f 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 @@ -611,4 +611,54 @@ mod tests { "self-keyed expansion (#595) must still resolve host-b" ); } + + // 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. + #[tokio::test(flavor = "multi_thread")] + async fn topk_range_tie_break_is_deterministic_by_label_values() { + let query = "topk(2, cpu_load)"; + // host-a is the clear #1. host-b and host-c are tied at 50.0 for the + // single remaining top-2 slot -- "host-b" < "host-c" lexicographically, + // so host-b should always win the tie once the sort is deterministic. + let engine = build_range_topk_engine( + query, + &[( + 0, + 1000, + &[("host-a", 100.0), ("host-b", 50.0), ("host-c", 50.0)], + )], + ); + + let (_, result) = engine + .handle_range_query_promql(query.to_string(), 1.0, 1.5, 1.0) + .expect("range topk query failed"); + let elements = match result { + QueryResult::Matrix(m) => m.values, + QueryResult::Vector(_) => panic!("expected a Matrix"), + }; + + assert_eq!( + topk_samples_for(&elements, "host-a"), + vec![(1000, 100.0)], + "host-a is the clear #1, always kept" + ); + assert_eq!( + topk_samples_for(&elements, "host-b"), + vec![(1000, 50.0)], + "host-b should deterministically win the tie against host-c (\"host-b\" < \"host-c\")" + ); + assert_eq!( + topk_samples_for(&elements, "host-c"), + vec![], + "host-c should deterministically lose the tie against host-b" + ); + } }