fix(query-engine): normalize topk labels in binary expressions (#644)

* fix(query-engine): normalize topk labels for instant binary arms

* fix(query-engine): normalize topk labels for range binary arms

* test(query-engine): cover range topk binary arithmetic operators

* test(query-engine): cover empty topk binary joins

* test(query-engine): cover range topk scalar arithmetic

* fix(query-engine): tighten topk binary operator coverage

* fix(query-engine): address roborev topk findings

* test(query-engine): cover nested topk comparison fallback

* fix(query-engine): fall back for topk matching modifiers

* test(query-engine): cover nested topk matching fallback

Author: milindsrivastava1997

asap-query-engine/src/engines/simple_engine/promql.rs

@@ -17,6 +17,8 @@ use std::collections::HashMap;
 use std::time::Instant;
 use tracing::{debug, warn};
 
+const METRIC_NAME_LABEL: &str = "__name__";
+
 /// Detects whether either side of a PromQL binary expression is a scalar
 /// (numeric literal), returning the scalar value, the other (vector) arm,
 /// and whether the scalar was on the left. Shared by instant and range
@@ -94,6 +96,18 @@ fn combine_scalar(
         .collect()
 }
 
+fn binary_matching_label_names(label_names: Vec<String>) -> Vec<String> {
+    label_names
+        .into_iter()
+        .filter(|label| label != METRIC_NAME_LABEL)
+        .collect()
+}
+
+fn is_supported_binary_arithmetic_op(op: &promql_parser::parser::token::TokenType) -> bool {
+    use promql_parser::parser::token::{T_ADD, T_DIV, T_MOD, T_MUL, T_POW, T_SUB};
+    matches!(op.id(), T_ADD | T_SUB | T_MUL | T_DIV | T_MOD | T_POW)
+}
+
 impl SimpleEngine {
     /// Aligns `end_timestamp` down to the nearest data-ingestion-interval
     /// boundary, unconditionally — mirroring SQL's `align_end_timestamp_sql`.
@@ -338,7 +352,7 @@ impl SimpleEngine {
         let statistic_to_compute = requirements.statistics[0];
 
         if statistic_to_compute == Statistic::Topk {
-            let mut new_labels = vec!["__name__".to_string()];
+            let mut new_labels = vec![METRIC_NAME_LABEL.to_string()];
             new_labels.extend(query_output_labels.labels);
             query_output_labels = KeyByLabelNames::new(new_labels);
         }
@@ -410,7 +424,8 @@ impl SimpleEngine {
             other => {
                 let config = self.find_query_config_promql_structural(other)?;
                 let ctx = self.build_query_execution_context_from_ast(other, &config, time)?;
-                let label_names = ctx.metadata.query_output_labels.labels.clone();
+                let label_names =
+                    binary_matching_label_names(ctx.metadata.query_output_labels.labels.clone());
                 Some((ctx, label_names))
             }
         }
@@ -437,6 +452,12 @@ impl SimpleEngine {
             Expr::NumberLiteral(_) => None, // caller handles scalars
             Expr::Paren(paren) => self.evaluate_binary_arm(&paren.expr, time),
             Expr::Binary(binary) => {
+                if binary.modifier.is_some() {
+                    return None;
+                }
+                if !is_supported_binary_arithmetic_op(&binary.op) {
+                    return None;
+                }
                 // Nested binary expression — recurse on both sides
                 let (lhs_results, lhs_labels) = self.evaluate_binary_arm(&binary.lhs, time)?;
                 let (rhs_results, rhs_labels) = self.evaluate_binary_arm(&binary.rhs, time)?;
@@ -460,11 +481,10 @@ impl SimpleEngine {
                 // for just this arm. Accepted behavior change (#567); warn loudly
                 // so it's visible rather than silent.
                 let results = self
-                    // (true, true): safe unconditionally — both flags are
-                    // self-gated on statistic == Topk / a "k" kwarg being
-                    // present, same as the main instant-query path (see
-                    // execute_query_pipeline's doc comment).
-                    .execute_query_pipeline(&ctx, true, true)
+                    // Binary arms need Topk limiting, but must remain in the
+                    // unformatted intermediate label representation until
+                    // after the binary join.
+                    .execute_query_pipeline(&ctx, true, false)
                     .map_err(|e| {
                         warn!(
                             "Binary-expr arm for metric '{}' failed ({}) — \
@@ -498,6 +518,13 @@ impl SimpleEngine {
             _ => return None,
         };
 
+        if !is_supported_binary_arithmetic_op(&binary.op) {
+            return None;
+        }
+        if binary.modifier.is_some() {
+            return None;
+        }
+
         let lhs = binary.lhs.as_ref();
         let rhs = binary.rhs.as_ref();
         let op = &binary.op;
@@ -691,16 +718,23 @@ impl SimpleEngine {
             _ => return None,
         };
 
+        if !is_supported_binary_arithmetic_op(&binary.op) {
+            return None;
+        }
+        if binary.modifier.is_some() {
+            return None;
+        }
+
         let lhs = binary.lhs.as_ref();
         let rhs = binary.rhs.as_ref();
         let op = &binary.op;
 
         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)?;
-            // (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()?;
+            // Binary arms need Topk limiting, but must remain in the
+            // unformatted intermediate label representation until after the
+            // arithmetic operation.
+            let results = self.execute_range_query_pipeline(&ctx, true, false).ok()?;
             let combined: Vec<RangeVectorElement> = results
                 .into_iter()
                 .map(|mut elem| {
@@ -728,12 +762,12 @@ impl SimpleEngine {
         if lhs_labels != rhs_labels {
             return None;
         }
-        // (true, true): self-gated, same rationale as the scalar-arm call above.
+        // Binary arms need Topk limiting, but not final presentation formatting.
         let lhs_results = self
-            .execute_range_query_pipeline(&lhs_ctx, true, true)
+            .execute_range_query_pipeline(&lhs_ctx, true, false)
             .ok()?;
         let rhs_results = self
-            .execute_range_query_pipeline(&rhs_ctx, true, true)
+            .execute_range_query_pipeline(&rhs_ctx, true, false)
             .ok()?;
 
         // Build lookup: label_key -> {timestamp -> value} for rhs
@@ -1577,14 +1611,12 @@ mod topk_pipeline_tests {
         }
     }
 
-    /// A topk leaf wrapped in a binary expr (`topk(10, ...) + 0`) must still
-    /// get the same top-10 truncation and metric-name-prefixed formatting as
-    /// the bare `topk(10, ...)` query — evaluate_arm_native's leaf branch
-    /// used to hardcode (false, false) for enable_topk_limiting/formatting,
-    /// which would have returned all 15 unformatted (single-label) rows here
-    /// instead of the top 10 with the metric-name prefix.
+    /// A topk leaf wrapped in an arithmetic binary expr (`topk(10, ...) + 0`)
+    /// must still truncate to the top 10, while arithmetic output drops the
+    /// metric name. Binary-arm evaluation must not apply standalone Topk
+    /// presentation formatting before the arithmetic operation.
     #[test]
-    fn topk_wrapped_in_binary_expr_still_truncates_and_formats() {
+    fn topk_wrapped_in_binary_expr_truncates_without_metric_name() {
         let (engine, store) = build_topk_engine();
 
         let context = engine
@@ -1624,16 +1656,10 @@ mod topk_pipeline_tests {
                 "results must stay sorted by count descending"
             );
         }
-        assert_eq!(
-            results[0].labels.labels,
-            vec![METRIC.to_string(), "10.0.0.15".to_string()],
-        );
+        assert_eq!(results[0].labels.labels, vec!["10.0.0.15".to_string()],);
         assert_eq!(results[0].value, 150.0);
         for element in &results {
-            assert_eq!(
-                element.labels.labels[0], METRIC,
-                "binary-expr path must still prepend the metric name (PromQL top-k formatting)",
-            );
+            assert_eq!(element.labels.labels.len(), 1);
         }
     }
 }