feat(query-engine): computed top-k in range queries + stage-E equivalence tests (#581 stage E prep)

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<Vec<crate::engines::query_result::RangeVectorElement>, String> {
         use crate::engines::query_result::RangeVectorElement;
         use crate::engines::window_merger::create_window_merger;
@@ -1829,8 +1839,144 @@ 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.
+    /// `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<KeyByLabelValues, crate::engines::query_result::RangeVectorElement>,
+        statistic: &Statistic,
+        query_kwargs: &HashMap<String, String>,
+        metric: &str,
+        enable_topk_formatting: bool,
+        enable_topk_limiting: bool,
+    ) -> Vec<crate::engines::query_result::RangeVectorElement> {
+        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::<usize>().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<KeyByLabelValues> = results.keys().cloned().collect();
+                let key_to_idx: HashMap<KeyByLabelValues, usize> = 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<u64, Vec<(usize, f64)>> = 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((idx, sample.value));
+                    }
+                }
+
+                let mut kept_timestamps_by_idx: HashMap<usize, HashSet<u64>> = HashMap::new();
+                for (timestamp, mut candidates) in by_timestamp {
+                    // 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 (idx, _) in candidates {
+                        kept_timestamps_by_idx
+                            .entry(idx)
+                            .or_default()
+                            .insert(timestamp);
+                    }
+                }
+
+                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_idx[&idx];
+                    elem.samples.retain(|s| keep.contains(&s.timestamp));
+                }
+            }
+        }
+
+        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()
     }
 }
 

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<RangeVectorElement> = 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<KeyByLabelValues, HashMap<u64, f64>> = 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<RangeVectorElement> = self
-            .execute_range_query_pipeline(&context)
+            .execute_range_query_pipeline(&context, true, true)
             .map_err(|e| {
                 warn!("Range query execution failed: {}", e);
                 e

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;