diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 73f0094..897d352 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/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) -> Vec { + 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 = 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); } } } 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 e1bc962..1a91294 100644 --- a/asap-query-engine/src/tests/range_query_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/range_query_arithmetic_tests.rs @@ -409,25 +409,10 @@ mod tests { ) } - // 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. + // Regression test for issue #631: Topk preserves the original grouping + // labels for binary matching, while arithmetic drops the metric name. #[tokio::test(flavor = "multi_thread")] - async fn test_range_vector_vector_topk_lhs_plus_plain_rhs_returns_none() { + async fn test_range_vector_vector_topk_lhs_plus_plain_rhs_joins_by_original_labels() { let query = "topk(2, metric_a) + sum(metric_b) by (host)"; let engine = build_range_topk_plus_plain_engine( "topk(2, metric_a)", @@ -436,12 +421,23 @@ mod tests { &[("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:?}" - ); + let (_, qr) = engine + .handle_range_query_promql(query.to_string(), 1.0, 2.0, 1.0) + .expect("Expected result for range topk/plain query"); + let elements = matrix_values(qr); + + assert_eq!(elements.len(), 2); + let mut values: HashMap = elements + .into_iter() + .map(|element| { + assert_eq!(element.labels.labels.len(), 1); + assert_eq!(element.samples.len(), 1); + (element.labels.labels[0].clone(), element.samples[0].value) + }) + .collect(); + assert_eq!(values.remove("host-a"), Some(1100.0)); + assert_eq!(values.remove("host-b"), Some(2050.0)); + assert!(values.is_empty()); } /// Builds a SimpleEngine with two independent self-keyed topk-capable @@ -535,25 +531,63 @@ mod tests { ) } - // 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 instant_vector_vector_topk_lhs_topk_rhs_joins_by_original_labels() { + 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 (_, qr) = engine + .handle_query_promql(query.to_string(), 1.0) + .expect("Expected result for instant topk/topk query"); + let elements = match qr { + QueryResult::Vector(vector) => vector.values, + other => panic!("Expected instant vector result, got {other:?}"), + }; + + assert_eq!(elements.len(), 1); + assert_eq!(elements[0].labels.labels, vec!["host-b".to_string()]); + assert!((elements[0].value - 250.0).abs() < 1e-10); + } + + #[tokio::test(flavor = "multi_thread")] + async fn instant_vector_vector_topk_lhs_plus_plain_rhs_joins_by_original_labels() { + 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 (_, qr) = engine + .handle_query_promql(query.to_string(), 1.0) + .expect("Expected result for instant topk/plain query"); + let elements = match qr { + QueryResult::Vector(vector) => vector.values, + other => panic!("Expected instant vector result, got {other:?}"), + }; + + assert_eq!(elements.len(), 2); + let mut values: HashMap = elements + .into_iter() + .map(|element| { + assert_eq!(element.labels.labels.len(), 1); + (element.labels.labels[0].clone(), element.value) + }) + .collect(); + assert_eq!(values.remove("host-a"), Some(1100.0)); + assert_eq!(values.remove("host-b"), Some(2050.0)); + assert!(values.is_empty()); + } + + // Regression test for issue #631: different Topk metric names must not + // become part of the intermediate vector-matching identity. + #[tokio::test(flavor = "multi_thread")] 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. @@ -579,4 +613,263 @@ mod tests { assert_eq!(elements[0].samples.len(), 1); assert!((elements[0].samples[0].value - 250.0).abs() < 1e-10); } + + #[tokio::test(flavor = "multi_thread")] + async fn range_topk_binary_join_preserves_matching_for_all_arithmetic_ops() { + for (op, expected, candidates_a, candidates_b) in [ + ( + "+", + 250.0, + vec![("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + vec![("host-a", 5.0), ("host-b", 200.0), ("host-c", 300.0)], + ), + ( + "-", + -150.0, + vec![("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + vec![("host-a", 5.0), ("host-b", 200.0), ("host-c", 300.0)], + ), + ( + "*", + 10_000.0, + vec![("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + vec![("host-a", 5.0), ("host-b", 200.0), ("host-c", 300.0)], + ), + ( + "/", + 0.25, + vec![("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + vec![("host-a", 5.0), ("host-b", 200.0), ("host-c", 300.0)], + ), + ( + "%", + 50.0, + vec![("host-a", 100.0), ("host-b", 50.0), ("host-c", 10.0)], + vec![("host-a", 5.0), ("host-b", 200.0), ("host-c", 300.0)], + ), + ( + "^", + 8.0, + vec![("host-a", 1.0), ("host-b", 2.0), ("host-c", 0.0)], + vec![("host-a", 1.0), ("host-b", 3.0), ("host-c", 2.0)], + ), + ] { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &candidates_a, + &candidates_b, + ); + let query = format!("topk(2, metric_a) {op} topk(2, metric_b)"); + let (_, qr) = engine + .handle_range_query_promql(query, 1.0, 2.0, 1.0) + .unwrap_or_else(|| panic!("Expected result for operator {op}")); + let elements = matrix_values(qr); + + assert_eq!(elements.len(), 1, "operator {op}"); + assert_eq!(elements[0].labels.labels, vec!["host-b".to_string()]); + assert_eq!(elements[0].samples.len(), 1, "operator {op}"); + assert!( + (elements[0].samples[0].value - expected).abs() < 1e-10, + "operator {op}: expected {expected}, got {}", + elements[0].samples[0].value + ); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_topk_binary_join_preserves_matching_for_power() { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 2.0), ("host-b", 1.0)], + &[("host-a", 3.0), ("host-b", 1.0)], + ); + + let (_, qr) = engine + .handle_range_query_promql( + "topk(2, metric_a) ^ topk(2, metric_b)".to_string(), + 1.0, + 2.0, + 1.0, + ) + .expect("Expected result for power operator"); + let elements = matrix_values(qr); + let values: HashMap = elements + .into_iter() + .map(|element| (element.labels.labels[0].clone(), element.samples[0].value)) + .collect(); + + assert_eq!(values.get("host-a"), Some(&8.0)); + assert_eq!(values.get("host-b"), Some(&1.0)); + } + + #[tokio::test(flavor = "multi_thread")] + async fn instant_topk_binary_join_with_no_matches_returns_empty_vector() { + 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", 300.0), ("host-d", 200.0)], + ); + + let result = + engine.handle_query_promql("topk(2, metric_a) + topk(2, metric_b)".to_string(), 1.0); + let (_, qr) = result.expect("valid no-match join should return an empty vector"); + let elements = match qr { + QueryResult::Vector(vector) => vector.values, + other => panic!("Expected instant vector result, got {other:?}"), + }; + assert!(elements.is_empty()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_topk_binary_join_with_no_matches_returns_empty_matrix() { + 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", 300.0), ("host-d", 200.0)], + ); + + let result = engine.handle_range_query_promql( + "topk(2, metric_a) + topk(2, metric_b)".to_string(), + 1.0, + 2.0, + 1.0, + ); + let (_, qr) = result.expect("valid no-match join should return an empty matrix"); + assert!(matrix_values(qr).is_empty()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_topk_scalar_arithmetic_drops_metric_name() { + 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)], + &[], + ); + + let (_, qr) = engine + .handle_range_query_promql("topk(2, metric_a) + 2".to_string(), 1.0, 2.0, 1.0) + .expect("Expected result for range topk/scalar query"); + let elements = matrix_values(qr); + + assert_eq!(elements.len(), 2); + let values: HashMap = elements + .into_iter() + .map(|element| { + assert_eq!(element.labels.labels.len(), 1); + assert_eq!(element.samples.len(), 1); + (element.labels.labels[0].clone(), element.samples[0].value) + }) + .collect(); + assert_eq!(values.get("host-a"), Some(&102.0)); + assert_eq!(values.get("host-b"), Some(&52.0)); + } + + #[tokio::test(flavor = "multi_thread")] + async fn instant_topk_binary_comparison_falls_back_to_prometheus() { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 100.0), ("host-b", 50.0)], + &[("host-a", 5.0), ("host-b", 200.0)], + ); + + assert!(engine + .handle_query_promql("topk(2, metric_a) == topk(2, metric_b)".to_string(), 1.0,) + .is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_topk_binary_comparison_falls_back_to_prometheus() { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 100.0), ("host-b", 50.0)], + &[("host-a", 5.0), ("host-b", 200.0)], + ); + + assert!(engine + .handle_range_query_promql( + "topk(2, metric_a) == topk(2, metric_b)".to_string(), + 1.0, + 2.0, + 1.0, + ) + .is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn nested_topk_binary_comparison_falls_back_to_prometheus() { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 100.0), ("host-b", 50.0)], + &[("host-a", 5.0), ("host-b", 200.0)], + ); + + assert!(engine + .handle_query_promql( + "(topk(2, metric_a) == topk(2, metric_b)) + topk(2, metric_a)".to_string(), + 1.0, + ) + .is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn instant_topk_binary_matching_modifier_falls_back_to_prometheus() { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 100.0), ("host-b", 50.0)], + &[("host-a", 5.0), ("host-b", 200.0)], + ); + + assert!(engine + .handle_query_promql( + "topk(2, metric_a) + on(__name__) topk(2, metric_b)".to_string(), + 1.0, + ) + .is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_topk_binary_matching_modifier_falls_back_to_prometheus() { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 100.0), ("host-b", 50.0)], + &[("host-a", 5.0), ("host-b", 200.0)], + ); + + assert!(engine + .handle_range_query_promql( + "topk(2, metric_a) + on(__name__) topk(2, metric_b)".to_string(), + 1.0, + 2.0, + 1.0, + ) + .is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn nested_topk_binary_matching_modifier_falls_back_to_prometheus() { + let engine = build_range_two_topk_engine( + "topk(2, metric_a)", + "topk(2, metric_b)", + &[("host-a", 100.0), ("host-b", 50.0)], + &[("host-a", 5.0), ("host-b", 200.0)], + ); + + assert!(engine + .handle_query_promql( + "(topk(2, metric_a) + on(__name__) topk(2, metric_b)) + topk(2, metric_a)" + .to_string(), + 1.0, + ) + .is_none()); + } }