From b5d7e3aac87d874ecfdfe27aa36ec9d068bb6f1c Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Fri, 21 Aug 2026 11:31:31 -0400 Subject: [PATCH 01/10] fix(query-engine): merge all sliding-window buckets per key instead of taking first execute_and_merge_store_queries kept only the first precomputed bucket per key for Sliding-window queries, discarding the rest when the store returned more than expected. DataFusion's SummaryMergeMultipleExec already merges all of them correctly, so binary-expr queries (still on DataFusion) don't hit this. #567 will move binary-expr onto this native path, so this native/DataFusion behavior gap needed closing first. Part of ProjectASAP/ASAPQuery#567 Stage 1. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/mod.rs | 48 +++-- asap-query-engine/src/tests/mod.rs | 1 + .../src/tests/native_pipeline_merge_tests.rs | 174 ++++++++++++++++++ 3 files changed, 208 insertions(+), 15 deletions(-) create mode 100644 asap-query-engine/src/tests/native_pipeline_merge_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 501f7218..bd45027f 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -549,22 +549,40 @@ impl SimpleEngine { }; let merged_values = if plan.values_query.is_exact_query { - // Sliding window: no merge needed, extract buckets from timestamped data - debug!("Sliding window mode: Skipping merge (expecting 1 precompute per key)"); - values_map - .into_iter() - .map(|(key, timestamped_buckets)| { - if timestamped_buckets.len() != 1 { - warn!( - "Sliding window expected 1 precompute per key, found {}. Using first.", - timestamped_buckets.len() - ); + // Sliding window: expected exactly 1 precompute per key today + // (ponytail: hardcoded, #554 will make >1 legitimate — don't + // block on it). The store can legitimately return more than + // expected for one exact window; merge whatever came back + // instead of arbitrarily keeping the first and dropping the + // rest (see #567). + const EXPECTED_BUCKETS_PER_KEY: usize = 1; + debug!("Sliding window mode: merging {} keys", values_map.len()); + let mut merged = HashMap::with_capacity(values_map.len()); + for (key, timestamped_buckets) in values_map.into_iter() { + if timestamped_buckets.is_empty() { + continue; + } + if timestamped_buckets.len() != EXPECTED_BUCKETS_PER_KEY { + warn!( + "Sliding window expected {} precompute(s) per key, found {}. Merging all.", + EXPECTED_BUCKETS_PER_KEY, + timestamped_buckets.len() + ); + } + let precomputes: Vec> = timestamped_buckets + .into_iter() + .map(|(_, bucket)| bucket.as_ref().clone_boxed_core()) + .collect(); + match self.merge_accumulators(&precomputes) { + Ok(merged_accumulator) => { + merged.insert(key, merged_accumulator); + } + Err(e) => { + warn!("Failed to merge accumulators for key {:?}: {}", key, e); } - // Extract bucket from timestamped tuple - let (_, bucket) = timestamped_buckets.into_iter().next().unwrap(); - (key, bucket.as_ref().clone_boxed_core()) - }) - .collect() + } + } + merged } else { // Tumbling window: merge needed debug!("Tumbling window mode: Merging {} outputs", values_map.len()); diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 9fbe6fc1..563f8540 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -3,6 +3,7 @@ pub mod clickhouse_forwarding_tests; pub mod datafusion; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests; +pub mod native_pipeline_merge_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; pub mod sql_pattern_matching_tests; diff --git a/asap-query-engine/src/tests/native_pipeline_merge_tests.rs b/asap-query-engine/src/tests/native_pipeline_merge_tests.rs new file mode 100644 index 00000000..044b11a2 --- /dev/null +++ b/asap-query-engine/src/tests/native_pipeline_merge_tests.rs @@ -0,0 +1,174 @@ +//! Native pipeline merge tests (issue #567, Stage 1). +//! +//! `execute_and_merge_store_queries`'s Sliding-window branch +//! (`simple_engine/mod.rs`) must merge every precomputed bucket returned for +//! a key, not just the first one. The store's `query_precomputed_output_exact` +//! can legitimately return more than one bucket for the same key under one +//! exact window (see `per_key.rs::query_precomputed_output_exact`), and +//! DataFusion's `SummaryMergeMultipleExec` already merges all of them +//! correctly — native must match. + +use crate::data_model::{AggregationType, WindowType}; +use crate::engines::query_result::InstantVectorElement; +use crate::engines::simple_engine::SimpleEngine; +use crate::precompute_operators::sum_accumulator::SumAccumulator; +use crate::tests::test_utilities::engine_factories::create_engine_multi_timestamp_with_window; + +const QUERY_TIME: f64 = 1000.0; // -> data time 1_000_000ms, see convert_query_time_to_data_time +const DATA_TIME: u64 = 1_000_000; +const SLIDING_WINDOW_MS: u64 = 1_000; // matches create_engine_multi_timestamp_with_window's fixed bucket width + +/// Runs a query through the native pipeline (`execute_query_pipeline`), the +/// same path `execute_and_merge_store_queries` is reached from. +fn execute_native( + engine: &SimpleEngine, + query: &str, + query_time_sec: f64, +) -> Vec { + let context = engine + .build_query_execution_context_promql(query.to_string(), query_time_sec) + .expect("Failed to build context"); + engine + .execute_query_pipeline(&context, false, false) + .expect("execute_query_pipeline failed") +} + +#[tokio::test] +async fn sliding_single_bucket_returns_its_value() { + let data = vec![( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(42.0)) as Box, + )]; + let query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + SLIDING_WINDOW_MS, + WindowType::Sliding, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1); + assert!((results[0].value - 42.0).abs() < 1e-10); +} + +#[tokio::test] +async fn sliding_two_buckets_for_same_key_are_merged_not_dropped() { + // Two precomputed buckets land under the same key and the same exact + // window (both at DATA_TIME). Today's code takes the first and warns; + // it must merge both. + let data = vec![ + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ), + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(5.0)) as Box, + ), + ]; + let query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + SLIDING_WINDOW_MS, + WindowType::Sliding, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1, "expected one merged result for host-a"); + assert!( + (results[0].value - 15.0).abs() < 1e-10, + "expected both buckets merged into 15.0, got {}", + results[0].value + ); +} + +#[tokio::test] +async fn sliding_bucket_count_mismatch_still_returns_merged_result() { + // 3 buckets (not just 2) for one key: generalizes #2 beyond the + // exactly-one-extra case, and confirms a mismatch never errors/drops — + // it merges everything and only warns. + let data = vec![ + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ), + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(5.0)) as Box, + ), + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(3.0)) as Box, + ), + ]; + let query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + SLIDING_WINDOW_MS, + WindowType::Sliding, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1); + assert!( + (results[0].value - 18.0).abs() < 1e-10, + "expected all 3 buckets merged into 18.0, got {}", + results[0].value + ); +} + +#[tokio::test] +async fn tumbling_multi_bucket_merge_unaffected_by_sliding_fix() { + // Regression guard: the Sliding-branch edit lives in the same `if` as + // the Tumbling branch below it — prove Tumbling's (already-correct) + // multi-timestamp merge is untouched, through the native pipeline. + let timestamps = [996_000u64, 997_000, 998_000, 999_000, 1_000_000]; + let data = timestamps + .iter() + .map(|&ts| { + ( + ts, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ) + }) + .collect(); + let query = "sum_over_time(http_requests[5s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + // window_size_ms < query range so do_merge=true. Equal (5s window, + // 5s range) hits a separate, pre-existing panic — see #569, not this stage. + 1_000, + WindowType::Tumbling, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1); + assert!( + (results[0].value - 50.0).abs() < 1e-10, + "expected 5 timestamps merged into 50.0, got {}", + results[0].value + ); +} From 1df67d4f47d7c9108885fd40f4086360ed4861ad Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Fri, 21 Aug 2026 16:30:21 -0400 Subject: [PATCH 02/10] fix(query-engine): dedupe sliding-window merge path, fix double-clone and stale latency label - Sliding branch now delegates to merge_precomputed_outputs (do_merge=true) instead of hand-rolling its own extract/merge/insert loop, removing the duplication with the Tumbling branch's merge path. - merge_accumulators now takes ownership of the accumulator Vec so its single-element shortcut can move the value out instead of re-cloning it on top of the clone already done to build the Vec. - The [LATENCY] log's merge/no-merge label was hardcoded off window_type and said "no merge" even when merge_accumulators was in fact called; it now reflects whether merging actually occurs. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/mod.rs | 38 +++++++------------ 1 file changed, 13 insertions(+), 25 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index bd45027f..4c618ebf 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -557,8 +557,7 @@ impl SimpleEngine { // rest (see #567). const EXPECTED_BUCKETS_PER_KEY: usize = 1; debug!("Sliding window mode: merging {} keys", values_map.len()); - let mut merged = HashMap::with_capacity(values_map.len()); - for (key, timestamped_buckets) in values_map.into_iter() { + for timestamped_buckets in values_map.values() { if timestamped_buckets.is_empty() { continue; } @@ -569,20 +568,10 @@ impl SimpleEngine { timestamped_buckets.len() ); } - let precomputes: Vec> = timestamped_buckets - .into_iter() - .map(|(_, bucket)| bucket.as_ref().clone_boxed_core()) - .collect(); - match self.merge_accumulators(&precomputes) { - Ok(merged_accumulator) => { - merged.insert(key, merged_accumulator); - } - Err(e) => { - warn!("Failed to merge accumulators for key {:?}: {}", key, e); - } - } } - merged + // Sliding windows always merge (all buckets belong to one + // logical window) — reuse the same merge path as Tumbling. + self.merge_precomputed_outputs(&values_map, true, agg_info.aggregation_type_for_value) } else { // Tumbling window: merge needed debug!("Tumbling window mode: Merging {} outputs", values_map.len()); @@ -594,13 +583,12 @@ impl SimpleEngine { }; let merge_duration = merge_start_time.elapsed(); + let did_merge = window_type == WindowType::Sliding + || do_merge + || agg_info.aggregation_type_for_value == AggregationType::DeltaSetAggregator; debug!( "[LATENCY] Precomputed output processing ({}): {:.2}ms, resulted in {} merged outputs", - if window_type == WindowType::Sliding { - "no merge" - } else { - "merge" - }, + if did_merge { "merge" } else { "no merge" }, merge_duration.as_secs_f64() * 1000.0, merged_values.len() ); @@ -1098,7 +1086,7 @@ impl SimpleEngine { debug!(" Merging accumulators (should_merge=true)"); #[cfg(feature = "extra_debugging")] let merge_start = Instant::now(); - match self.merge_accumulators(&precomputes) { + match self.merge_accumulators(precomputes) { Ok(merged_accumulator) => { #[cfg(feature = "extra_debugging")] let merge_duration = merge_start.elapsed(); @@ -1141,21 +1129,21 @@ impl SimpleEngine { /// This follows the Python merge_accumulators approach fn merge_accumulators( &self, - accumulators: &[Box], + accumulators: Vec>, ) -> Result, AccumulatorError> { if accumulators.is_empty() { return Err(AccumulatorError::EmptySlice); } if accumulators.len() == 1 { - return Ok(accumulators[0].clone_boxed_core()); + return Ok(accumulators.into_iter().next().unwrap()); } // Try to use optimized batch merge for KLL accumulators if accumulators[0].get_accumulator_type() == AggregationType::DatasketchesKLL { use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; - match DatasketchesKLLAccumulator::merge_multiple(accumulators) { + match DatasketchesKLLAccumulator::merge_multiple(&accumulators) { Ok(merged) => return Ok(Box::new(merged)), Err(e) => { warn!( @@ -1171,7 +1159,7 @@ impl SimpleEngine { if accumulators[0].get_accumulator_type() == AggregationType::CountMinSketch { use crate::precompute_operators::count_min_sketch_accumulator::CountMinSketchAccumulator; - match CountMinSketchAccumulator::merge_multiple(accumulators) { + match CountMinSketchAccumulator::merge_multiple(&accumulators) { Ok(merged) => return Ok(Box::new(merged)), Err(e) => { warn!( From 8e768fa72b7a5bc08e49ece1cfd3ae4c6a2a73f9 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 22 Aug 2026 21:53:08 -0400 Subject: [PATCH 03/10] fix(query-engine): warn instead of silently dropping keys with no precompute buckets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit merge_precomputed_outputs silently skipped any key whose timestamped_buckets list was empty, with no log signal — unlike the "found N, expected 1" mismatch case a few lines up in the Sliding caller, which does warn. Since this function is shared by Sliding, Tumbling, and the keys-merge path, the warn now covers all three instead of being Sliding-only. Also files #575 to compute EXPECTED_BUCKETS_PER_KEY instead of hardcoding it to 1, since #554 will make >1 legitimate whenever a sliding-window query's range exceeds the window size. Co-Authored-By: Claude Sonnet 5 --- asap-query-engine/src/engines/simple_engine/mod.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 4c618ebf..4cf0fbd2 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1074,7 +1074,12 @@ impl SimpleEngine { let mut merged = HashMap::with_capacity(precomputed_outputs_map.len()); for (key, timestamped_buckets) in precomputed_outputs_map.iter() { - if !timestamped_buckets.is_empty() { + if timestamped_buckets.is_empty() { + warn!( + "Store returned key {:?} with no precompute buckets; skipping", + key + ); + } else { // Extract just the buckets (without timestamps) for merging let precomputes: Vec> = timestamped_buckets .iter() From 8311f609217d34f39cdec168886bdf4f17097b26 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Fri, 21 Aug 2026 13:28:36 -0400 Subject: [PATCH 04/10] feat(query-engine): add native instant binary-expr evaluator, staged ahead of cutover MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builds a native (non-DataFusion) implementation of PromQL binary- arithmetic instant queries: a recursive arm evaluator over Vec, plus vector-vector and scalar combiners lifted from the existing range-binary path. Not yet wired into production dispatch (handle_query_promql still calls the DataFusion path) — exposed via handle_query_promql_native for equivalence testing against the DataFusion path ahead of the Stage 3 cutover. One accepted, deliberately loud behavior change: an arm with zero current precomputed data now falls back to Prometheus (matching "not acceleratable" semantics) instead of DataFusion's silent empty-result behavior, with a warn! so it's observable. Part of ProjectASAP/ASAPQuery#567 Stage 2. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/promql.rs | 162 +++++++++- asap-query-engine/src/tests/mod.rs | 1 + .../src/tests/native_binary_instant_tests.rs | 298 ++++++++++++++++++ 3 files changed, 460 insertions(+), 1 deletion(-) create mode 100644 asap-query-engine/src/tests/native_binary_instant_tests.rs diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index a9a2d546..87dae33f 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -9,7 +9,7 @@ use super::{ RangeQueryParams, }; use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig}; -use crate::engines::query_result::{QueryResult, RangeVectorElement}; +use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVectorElement}; use asap_types::query_requirements::build_query_requirements_promql; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; @@ -36,6 +36,56 @@ fn detect_scalar_arm<'a>( } } +/// Native vector-vector combiner for one binary-expr level (instant query): +/// joins two arms' results by label key and applies `op` per matching key, +/// dropping non-matches (inner-join semantics, matching DataFusion's +/// `build_binary_vector_plan`). Positional `KeyByLabelValues` equality is +/// safe here: label *names* are always canonically sorted by +/// `KeyByLabelNames::new()`, and every source of PromQL label-name ordering +/// routes through it — see #567's label-order investigation. +fn combine_vector_vector_native( + lhs_results: Vec, + rhs_results: Vec, + op: &promql_parser::parser::token::TokenType, +) -> Vec { + let rhs_map: HashMap = rhs_results + .into_iter() + .map(|elem| (elem.labels, elem.value)) + .collect(); + + lhs_results + .into_iter() + .filter_map(|lhs_elem| { + rhs_map.get(&lhs_elem.labels).map(|&rhs_val| { + let value = SimpleEngine::apply_range_binary_op(op, lhs_elem.value, rhs_val); + InstantVectorElement::new(lhs_elem.labels.clone(), value) + }) + }) + .collect() +} + +/// Native scalar combiner for one binary-expr level (instant query): +/// applies `op(scalar, value)` or `op(value, scalar)` per `scalar_on_left` +/// to every element of the vector arm's results. +fn combine_scalar_native( + vector_results: Vec, + scalar: f64, + op: &promql_parser::parser::token::TokenType, + scalar_on_left: bool, +) -> Vec { + vector_results + .into_iter() + .map(|elem| { + let value = if scalar_on_left { + SimpleEngine::apply_range_binary_op(op, scalar, elem.value) + } else { + SimpleEngine::apply_range_binary_op(op, elem.value, scalar) + }; + InstantVectorElement::new(elem.labels, value) + }) + .collect() +} + impl SimpleEngine { /// Aligns `end_timestamp` down to the nearest data-ingestion-interval /// boundary, unconditionally — mirroring SQL's `align_end_timestamp_sql`. @@ -453,6 +503,116 @@ impl SimpleEngine { Some((output_labels, QueryResult::vector(results, query_time))) } + /// Recursively evaluates one arm of a binary arithmetic expression via + /// the native pipeline (`execute_query_pipeline`), instead of building a + /// DataFusion plan. Mirrors `build_arm_logical_plan`'s shape exactly, + /// including its limitation: nested `Binary` arms are combined as + /// vector-vector only — a scalar inside a nested arm (e.g. `(a+5)*b`) + /// is not supported, matching today's DataFusion path. + /// + /// - Leaf arm: resolved via `resolve_arm_leaf_context`, executed through + /// `execute_query_pipeline`. + /// - Binary arm: recursively evaluate both sides then combine with + /// `combine_vector_vector_native`. + /// - Scalar literal: returns `None` (handled by the caller separately). + fn evaluate_arm_native( + &self, + arm_ast: &promql_parser::parser::Expr, + time: f64, + ) -> Option<(Vec, Vec)> { + use promql_parser::parser::Expr; + + match arm_ast { + Expr::NumberLiteral(_) => None, // caller handles scalars + Expr::Paren(paren) => self.evaluate_arm_native(&paren.expr, time), + Expr::Binary(binary) => { + // Nested binary expression — recurse on both sides + let (lhs_results, lhs_labels) = self.evaluate_arm_native(&binary.lhs, time)?; + let (rhs_results, _) = self.evaluate_arm_native(&binary.rhs, time)?; + let combined = combine_vector_vector_native(lhs_results, rhs_results, &binary.op); + Some((combined, lhs_labels)) + } + _ => { + let (ctx, label_names) = self.resolve_arm_leaf_context(arm_ast, time)?; + // Unlike DataFusion's PrecomputedSummaryReadExec (which streams + // whatever rows exist, including zero, so a currently-empty arm + // still returns Some(empty vector)), execute_query_pipeline errors + // when the store has no precomputed outputs at all for this arm — + // that propagates to None here, triggering a full Prometheus + // fallback for the whole expression instead of an empty result + // for just this arm. Accepted behavior change (#567); warn loudly + // so it's visible rather than silent. + let results = self + .execute_query_pipeline(&ctx, false, false) + .map_err(|e| { + warn!( + "Native binary-expr arm for metric '{}' produced no results \ + ({}) — unlike DataFusion, this falls back to Prometheus for \ + the whole expression rather than returning an empty result \ + for just this arm", + ctx.metric, e + ); + e + }) + .ok()?; + Some((results, label_names)) + } + } + } + + /// Native (non-DataFusion) implementation of binary-expr handling, + /// staged ahead of #567's cutover so it can be equivalence-tested + /// against the DataFusion path (`handle_binary_expr_promql`) before + /// that path is rewired to call this one. Not yet reachable from + /// `handle_query_promql` — see `handle_query_promql_native` for a + /// test-facing entrypoint. + pub fn handle_binary_expr_promql_native( + &self, + ast: &promql_parser::parser::Expr, + time: f64, + ) -> Option<(KeyByLabelNames, QueryResult)> { + use promql_parser::parser::Expr; + + let query_time = Self::convert_query_time_to_data_time(time); + + let binary = match ast { + Expr::Binary(b) => b, + _ => 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 (vector_results, label_names) = self.evaluate_arm_native(vector_arm, time)?; + let combined = combine_scalar_native(vector_results, scalar, op, scalar_on_left); + return Some(( + KeyByLabelNames::new(label_names), + QueryResult::vector(combined, query_time), + )); + } + + // Vector–vector + let (lhs_results, lhs_labels) = self.evaluate_arm_native(lhs, time)?; + let (rhs_results, _) = self.evaluate_arm_native(rhs, time)?; + let combined = combine_vector_vector_native(lhs_results, rhs_results, op); + let output_labels = KeyByLabelNames::new(lhs_labels); + Some((output_labels, QueryResult::vector(combined, query_time))) + } + + /// Parses `query` and dispatches to `handle_binary_expr_promql_native`. + /// Test-facing convenience wrapper — mirrors the binary-expr branch of + /// `handle_query_promql`, but for the native path. + pub fn handle_query_promql_native( + &self, + query: String, + time: f64, + ) -> Option<(KeyByLabelNames, QueryResult)> { + let ast = promql_parser::parser::parse(&query).ok()?; + self.handle_binary_expr_promql_native(&ast, time) + } + /// Applies a PromQL binary arithmetic operator to two f64 values. fn apply_range_binary_op( op: &promql_parser::parser::token::TokenType, diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 563f8540..d5b31eec 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -3,6 +3,7 @@ pub mod clickhouse_forwarding_tests; pub mod datafusion; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests; +pub mod native_binary_instant_tests; pub mod native_pipeline_merge_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; diff --git a/asap-query-engine/src/tests/native_binary_instant_tests.rs b/asap-query-engine/src/tests/native_binary_instant_tests.rs new file mode 100644 index 00000000..45d58c10 --- /dev/null +++ b/asap-query-engine/src/tests/native_binary_instant_tests.rs @@ -0,0 +1,298 @@ +//! Native instant binary-expr evaluator tests (issue #567, Stage 2). +//! +//! `handle_binary_expr_promql_native`/`evaluate_arm_native` are a new, +//! parallel implementation of PromQL binary-arithmetic instant queries, +//! built to eventually replace the DataFusion-backed +//! `handle_binary_expr_promql` (#567's Stage 3 cutover). These tests compare +//! the two paths directly — both are reachable today (`handle_query_promql` +//! for DataFusion, `handle_query_promql_native` for the new path) so +//! equivalence can be proven before anything is rewired. + +#[cfg(test)] +mod tests { + use crate::data_model::{AggregationType, KeyByLabelValues}; + use crate::engines::query_result::QueryResult; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::precompute_operators::{CountMinSketchAccumulator, DeltaSetAggregatorAccumulator}; + use crate::tests::test_utilities::engine_factories::{ + create_engine_dual_input, create_engine_single_pop, create_engine_three_metrics, + create_engine_two_metrics, + }; + use crate::AggregateCore; + + const QUERY_TIME: f64 = 1000.0; + + fn vector_values(qr: QueryResult) -> Vec<(Vec, f64)> { + match qr { + QueryResult::Vector(iv) => iv + .values + .into_iter() + .map(|e| (e.labels.labels, e.value)) + .collect(), + _ => panic!("Expected vector result"), + } + } + + fn sorted(mut v: Vec<(Vec, f64)>) -> Vec<(Vec, f64)> { + v.sort_by(|a, b| a.0.cmp(&b.0)); + v + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_vector_vector_all_ops_match_datafusion() { + // `^` excluded: DataFusion's build_binary_vector_plan maps T_POW to + // Operator::BitwiseXor "as a proxy" (plan_builder.rs comment) and + // fails to even produce a result for it today — nothing to compare + // native against for that operator. See the dedicated `^` test below. + for (op, expected) in [ + ("+", 30.0), + ("-", -10.0), + ("*", 200.0), + ("/", 0.5), + ("%", 10.0), + ] { + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = format!("sum(metric_a) by (host) {op} sum(metric_b) by (host)"); + let (_, old_qr) = engine + .handle_query_promql(query.clone(), QUERY_TIME) + .unwrap_or_else(|| panic!("old path failed for op {op}")); + let (_, new_qr) = engine + .handle_query_promql_native(query.clone(), QUERY_TIME) + .unwrap_or_else(|| panic!("new path failed for op {op}")); + + let old_values = sorted(vector_values(old_qr)); + let new_values = sorted(vector_values(new_qr)); + assert_eq!(old_values.len(), 1, "op {op}"); + assert_eq!( + new_values, old_values, + "op {op}: native must match DataFusion" + ); + assert!( + (new_values[0].1 - expected).abs() < 1e-10, + "op {op}: expected {expected}, got {}", + new_values[0].1 + ); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_power_operator_computes_correctly_datafusion_unsupported() { + // DataFusion's `^` support is broken today (see comment above); this + // documents that native computes it correctly via f64::powf, + // independent of a DataFusion comparison. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(2.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = "sum(metric_a) by (host) ^ sum(metric_b) by (host)"; + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + let new_values = vector_values(new_qr); + assert!((new_values[0].1 - 1024.0).abs() < 1e-6, "2^10 = 1024"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_vector_scalar_both_orderings_match_datafusion() { + let engine = create_engine_single_pop( + "errors_total", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(7.0)) as Box, + )], + "sum(errors_total) by (host)", + ); + + for query in [ + "sum(errors_total) by (host) * 100", + "100 * sum(errors_total) by (host)", + ] { + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .unwrap_or_else(|| panic!("old path failed for {query}")); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .unwrap_or_else(|| panic!("new path failed for {query}")); + + let old_values = vector_values(old_qr); + let new_values = vector_values(new_qr); + assert_eq!(new_values, old_values, "{query}"); + assert!((new_values[0].1 - 700.0).abs() < 1e-10, "{query}"); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_nested_binary_matches_datafusion() { + // (metric_a + metric_b) * metric_c, all Sum, host-a: (10+20)*3 = 90 + let engine = create_engine_three_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + "metric_c", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(3.0)) as Box, + )], + "sum(metric_c) by (host)", + ); + + let query = "(sum(metric_a) by (host) + sum(metric_b) by (host)) * sum(metric_c) by (host)"; + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("old path failed"); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + + let old_values = vector_values(old_qr); + let new_values = vector_values(new_qr); + assert_eq!(new_values, old_values); + assert!((new_values[0].1 - 90.0).abs() < 1e-10, "(10+20)*3 = 90"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_binary_expr_no_data_falls_back_to_none() { + // metric_a is configured (schema + pattern match) but has zero + // precomputed data. Accepted behavior change (#567): unlike + // DataFusion (which would return an empty-but-present result), + // native falls back to Prometheus (returns None) for the whole + // expression — see evaluate_arm_native's leaf branch. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![], // no data at all for metric_a + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (host)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + assert!( + old_result.is_some(), + "DataFusion returns an empty-but-present result, not None, for a currently-empty arm" + ); + + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + assert!( + new_result.is_none(), + "native falls back to Prometheus (None) for a currently-empty arm" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_binary_expr_unsupported_arm_returns_none() { + // foo() is not a supported PromQL function -> arm lookup fails -> None, + // same as today's DataFusion path (dispatch_arithmetic_tests.rs). + let engine = create_engine_single_pop( + "requests_total", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(200.0)) as Box, + )], + "sum(requests_total) by (host)", + ); + + let query = "foo(errors_total[5m]) / sum(requests_total) by (host)"; + assert!(engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .is_none()); + assert!(engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_vector_vector_dual_population_matches_datafusion() { + // event_frequency is dual-population (CountMinSketch values + + // DeltaSetAggregator keys) -- confirms the leaf swap + // (ctx.to_logical_plan() -> execute_query_pipeline) still resolves + // keys_query correctly. Wrapped in `+ 0` to route through the + // binary-expr handler at all. + let cms = CountMinSketchAccumulator::new(2, 3); + let mut keys = DeltaSetAggregatorAccumulator::new(); + keys.add_key(KeyByLabelValues { + labels: vec!["host-a".to_string(), "evt-1".to_string()], + }); + + let engine = create_engine_dual_input( + "event_frequency", + AggregationType::CountMinSketch, + AggregationType::DeltaSetAggregator, + vec![], + vec!["host", "event"], + vec![(None, Box::new(cms))], + vec![(None, Box::new(keys))], + "count(event_frequency) by (host, event)", + ); + + let query = "count(event_frequency) by (host, event) + 0"; + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("old path failed"); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + + assert_eq!(sorted(vector_values(new_qr)), sorted(vector_values(old_qr))); + } +} From 90128c82d707735c5dc02e2be688058e6989f94f Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Fri, 21 Aug 2026 17:36:01 -0400 Subject: [PATCH 05/10] fix(query-engine): reject cross-label-set joins in native binary-expr combiner combine_vector_vector_native joined two arms purely by positional KeyByLabelValues equality (a bare Vec of values, no label names attached). DataFusion's build_binary_vector_plan joins on named columns instead, so it fails to resolve (-> None) whenever the two arms don't share the same label set. Native had no equivalent check, so it could silently fabricate a joined result whenever two differently-grouped arms' values happened to coincide (e.g. `sum(a) by (host) + sum(b) by (region)` with a host value equal to a region value). combine_vector_vector_native now takes both arms' label-name lists and returns None on a mismatch, matching DataFusion's failed-join behavior. Adds four regression tests proving the divergence (three were red before this fix) plus a same-label-set control case. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/promql.rs | 50 +++-- .../src/tests/native_binary_instant_tests.rs | 174 ++++++++++++++++++ 2 files changed, 208 insertions(+), 16 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 87dae33f..1c738621 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -40,28 +40,39 @@ fn detect_scalar_arm<'a>( /// joins two arms' results by label key and applies `op` per matching key, /// dropping non-matches (inner-join semantics, matching DataFusion's /// `build_binary_vector_plan`). Positional `KeyByLabelValues` equality is -/// safe here: label *names* are always canonically sorted by -/// `KeyByLabelNames::new()`, and every source of PromQL label-name ordering -/// routes through it — see #567's label-order investigation. +/// safe *once `lhs_labels == rhs_labels` is confirmed*: label names are +/// canonically sorted by `KeyByLabelNames::new()`, so two arms with the same +/// label set always order their values the same way. Returns `None` if the +/// two arms don't share the same label set — mirroring DataFusion's +/// `build_binary_vector_plan`, which fails to resolve a join column that +/// only exists on one side. fn combine_vector_vector_native( lhs_results: Vec, + lhs_labels: &[String], rhs_results: Vec, + rhs_labels: &[String], op: &promql_parser::parser::token::TokenType, -) -> Vec { +) -> Option> { + if lhs_labels != rhs_labels { + return None; + } + let rhs_map: HashMap = rhs_results .into_iter() .map(|elem| (elem.labels, elem.value)) .collect(); - lhs_results - .into_iter() - .filter_map(|lhs_elem| { - rhs_map.get(&lhs_elem.labels).map(|&rhs_val| { - let value = SimpleEngine::apply_range_binary_op(op, lhs_elem.value, rhs_val); - InstantVectorElement::new(lhs_elem.labels.clone(), value) + Some( + lhs_results + .into_iter() + .filter_map(|lhs_elem| { + rhs_map.get(&lhs_elem.labels).map(|&rhs_val| { + let value = SimpleEngine::apply_range_binary_op(op, lhs_elem.value, rhs_val); + InstantVectorElement::new(lhs_elem.labels.clone(), value) + }) }) - }) - .collect() + .collect(), + ) } /// Native scalar combiner for one binary-expr level (instant query): @@ -528,8 +539,14 @@ impl SimpleEngine { Expr::Binary(binary) => { // Nested binary expression — recurse on both sides let (lhs_results, lhs_labels) = self.evaluate_arm_native(&binary.lhs, time)?; - let (rhs_results, _) = self.evaluate_arm_native(&binary.rhs, time)?; - let combined = combine_vector_vector_native(lhs_results, rhs_results, &binary.op); + let (rhs_results, rhs_labels) = self.evaluate_arm_native(&binary.rhs, time)?; + let combined = combine_vector_vector_native( + lhs_results, + &lhs_labels, + rhs_results, + &rhs_labels, + &binary.op, + )?; Some((combined, lhs_labels)) } _ => { @@ -595,8 +612,9 @@ impl SimpleEngine { // Vector–vector let (lhs_results, lhs_labels) = self.evaluate_arm_native(lhs, time)?; - let (rhs_results, _) = self.evaluate_arm_native(rhs, time)?; - let combined = combine_vector_vector_native(lhs_results, rhs_results, op); + let (rhs_results, rhs_labels) = self.evaluate_arm_native(rhs, time)?; + let combined = + combine_vector_vector_native(lhs_results, &lhs_labels, rhs_results, &rhs_labels, op)?; let output_labels = KeyByLabelNames::new(lhs_labels); Some((output_labels, QueryResult::vector(combined, query_time))) } diff --git a/asap-query-engine/src/tests/native_binary_instant_tests.rs b/asap-query-engine/src/tests/native_binary_instant_tests.rs index 45d58c10..1599d8c4 100644 --- a/asap-query-engine/src/tests/native_binary_instant_tests.rs +++ b/asap-query-engine/src/tests/native_binary_instant_tests.rs @@ -295,4 +295,178 @@ mod tests { assert_eq!(sorted(vector_values(new_qr)), sorted(vector_values(old_qr))); } + + // --- Regression tests: combine_vector_vector_native joins purely on + // positional KeyByLabelValues equality (raw Vec of values, no + // label names attached at all -- see key_by_label_values.rs). DataFusion's + // build_binary_vector_plan joins on named columns (lhs.{label} = + // rhs.{label}) instead, so it fails to resolve (-> None) whenever the two + // arms don't share the same label set. Native has no equivalent check, so + // it can silently join two arms grouped by *different* labels whenever + // their values happen to coincide. These tests are expected to be RED + // against the current implementation -- they exist to prove the + // divergence before combine_vector_vector_native is fixed. + + #[tokio::test(flavor = "multi_thread")] + async fn native_mismatched_label_names_with_colliding_values_diverges_from_datafusion() { + // metric_a grouped by (host), metric_b grouped by (region) -- disjoint + // label sets -- but both happen to produce the value "us-east". + // DataFusion can't join lhs.host against a plan with no "host" column + // -> None. Native joins on the bare value vector ["us-east"] == + // ["us-east"] -> a spurious combined result. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["us-east".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["region"], + vec![( + Some(vec!["us-east".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (region)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + + match old_result { + None => assert!( + new_result.is_none(), + "BUG: DataFusion refuses to join arms with different label sets (None), \ + but native silently produced {new_result:?} by matching on value alone" + ), + Some((_, old_qr)) => assert_eq!( + new_result.map(|(_, qr)| sorted(vector_values(qr))), + Some(sorted(vector_values(old_qr))), + "native and DataFusion must agree when both produce a result" + ), + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_mismatched_label_names_multi_label_full_collision_diverges_from_datafusion() { + // Same bug with a 2-label grouping: (host, dc) vs (region, zone), but + // the *entire* ordered value vector coincides ("us-east", "az1" on + // both sides) -- confirms the collision isn't a single-label fluke. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host", "dc"], + vec![( + Some(vec!["us-east".to_string(), "az1".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host, dc)", + "metric_b", + AggregationType::Sum, + vec!["region", "zone"], + vec![( + Some(vec!["us-east".to_string(), "az1".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (region, zone)", + ); + + let query = "sum(metric_a) by (host, dc) + sum(metric_b) by (region, zone)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + + match old_result { + None => assert!( + new_result.is_none(), + "BUG: DataFusion refuses to join arms with different label sets (None), \ + but native silently produced {new_result:?} by matching on value alone" + ), + Some((_, old_qr)) => assert_eq!( + new_result.map(|(_, qr)| sorted(vector_values(qr))), + Some(sorted(vector_values(old_qr))), + "native and DataFusion must agree when both produce a result" + ), + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_mismatched_label_names_non_colliding_values_still_diverges_from_datafusion() { + // Same disjoint label sets (host vs region), but this time the values + // don't collide ("us-east" vs "eu-west"). DataFusion still can't + // resolve the join (column doesn't exist, independent of the data) + // -> None. Native's join correctly finds no match on the *value*, but + // that's incidental -- it still returns Some(empty) instead of None, + // because it never checked whether the label sets matched at all. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["us-east".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["region"], + vec![( + Some(vec!["eu-west".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (region)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + + assert_eq!( + old_result.is_none(), + new_result.is_none(), + "BUG: DataFusion returns {old_result:?} (no such column to join on), native \ + returns {new_result:?} -- native never checks that the two arms' label sets match" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_same_label_set_no_matching_values_both_agree_empty() { + // Control case: identical label sets (host on both sides) but + // disjoint values -- both paths should agree there's no match. This + // isolates the bug to *differing* label sets, not "no match" in + // general. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-b".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (host)"; + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("old path failed"); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + + let old_values = vector_values(old_qr); + assert_eq!(old_values, Vec::new()); + assert_eq!(sorted(vector_values(new_qr)), sorted(old_values)); + } } From d86abd654efad489ce58f5855b106015ac3c28d0 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 22 Aug 2026 22:23:39 -0400 Subject: [PATCH 06/10] fix(query-engine): enable topk limiting/formatting for binary-expr leaf arms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit evaluate_arm_native's leaf branch hardcoded execute_query_pipeline(&ctx, false, false), so a topk arm inside a binary expression (e.g. `topk(10, metric) + 0`) never got truncated to k or metric-name-prefixed — both flags are self-gated on statistic == Topk / a "k" kwarg being present (see execute_query_pipeline's doc comment), so passing (true, true) unconditionally is a no-op for non-topk arms, matching what the main non-binary instant-query path already does. Adds a regression test proving the divergence: topk(10, ...) + 0 returned all 15 unformatted rows before this fix, now correctly truncates to 10 with the metric-name prefix. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/promql.rs | 66 ++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 1c738621..bdf615e4 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -560,7 +560,11 @@ impl SimpleEngine { // for just this arm. Accepted behavior change (#567); warn loudly // so it's visible rather than silent. let results = self - .execute_query_pipeline(&ctx, false, false) + // (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) .map_err(|e| { warn!( "Native binary-expr arm for metric '{}' produced no results \ @@ -1652,4 +1656,64 @@ mod topk_pipeline_tests { assert!(pair[0] >= pair[1]); } } + + /// 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. + #[test] + fn topk_wrapped_in_binary_expr_still_truncates_and_formats() { + let (engine, store) = build_topk_engine(); + + let context = engine + .build_query_execution_context_promql(TOPK_QUERY.to_string(), QUERY_TIME) + .expect("context should build"); + let window = &context.store_plan.values_query; + + let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); + for i in 1..=15u64 { + let srcip = format!("10.0.0.{i}"); + sketch.inner.update(&srcip, (i * 10) as f64); + } + + let output = + PrecomputedOutput::new(window.start_timestamp, window.end_timestamp, None, AGG_ID); + store + .insert_precomputed_output(output, Box::new(sketch)) + .expect("insert should succeed"); + + let (_, query_result) = engine + .handle_query_promql_native(format!("{TOPK_QUERY} + 0"), QUERY_TIME) + .expect("binary-expr-wrapped topk should still resolve via the native path"); + + let results = match query_result { + QueryResult::Vector(iv) => iv.values, + other => panic!("expected a vector result, got {other:?}"), + }; + + assert_eq!( + results.len(), + 10, + "topk(10, ...) + 0 must still truncate to 10 rows" + ); + for pair in results.windows(2) { + assert!( + pair[0].value >= pair[1].value, + "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].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)", + ); + } + } } From eb2fd259f0f3aea12e9d2d489dac56a1598825d7 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Fri, 21 Aug 2026 13:52:27 -0400 Subject: [PATCH 07/10] feat(query-engine): cut over PromQL binary-expr instant queries to native execution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewires handle_binary_expr_promql to call the native evaluator/ combiners (built in the prior stage) instead of building a DataFusion plan. Deletes the old DataFusion-based handle_binary_expr_promql, build_arm_logical_plan, and the tokio::task::block_in_place(... block_on(...)) wrapper it needed (native execution is synchronous). Renames evaluate_arm_native/combine_vector_vector_native/ combine_scalar_native -> evaluate_binary_arm/combine_vector_vector/ combine_scalar now that native is the only implementation, and drops the now-redundant handle_binary_expr_promql_native/ handle_query_promql_native test-only wrappers. This was the only production code path still reachable through DataFusion (asap-query-engine's SimpleEngine now serves every query shape through the same native fetch/merge pipeline). DataFusion's CustomQueryPlanner/PrecomputedSummaryReadExec/SummaryMergeMultipleExec/ build_binary_vector_plan/build_scalar_plan and the datafusion dependency itself are left in place, per #567's scope — still exercised by their own dedicated tests and the unwired execute_plan prototype, just no longer reachable from production. Closes ProjectASAP/ASAPQuery#567 Stage 3 (final stage). Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/promql.rs | 187 ++-------- .../plan_execution_arithmetic_tests.rs | 3 +- .../src/tests/native_binary_instant_tests.rs | 335 +++++++++--------- 3 files changed, 205 insertions(+), 320 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index bdf615e4..970596ec 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -36,8 +36,8 @@ fn detect_scalar_arm<'a>( } } -/// Native vector-vector combiner for one binary-expr level (instant query): -/// joins two arms' results by label key and applies `op` per matching key, +/// Vector-vector combiner for one binary-expr level (instant query): joins +/// two arms' results by label key and applies `op` per matching key, /// dropping non-matches (inner-join semantics, matching DataFusion's /// `build_binary_vector_plan`). Positional `KeyByLabelValues` equality is /// safe *once `lhs_labels == rhs_labels` is confirmed*: label names are @@ -46,7 +46,7 @@ fn detect_scalar_arm<'a>( /// two arms don't share the same label set — mirroring DataFusion's /// `build_binary_vector_plan`, which fails to resolve a join column that /// only exists on one side. -fn combine_vector_vector_native( +fn combine_vector_vector( lhs_results: Vec, lhs_labels: &[String], rhs_results: Vec, @@ -75,10 +75,10 @@ fn combine_vector_vector_native( ) } -/// Native scalar combiner for one binary-expr level (instant query): -/// applies `op(scalar, value)` or `op(value, scalar)` per `scalar_on_left` -/// to every element of the vector arm's results. -fn combine_scalar_native( +/// Scalar combiner for one binary-expr level (instant query): applies +/// `op(scalar, value)` or `op(value, scalar)` per `scalar_on_left` to every +/// element of the vector arm's results. +fn combine_scalar( vector_results: Vec, scalar: f64, op: &promql_parser::parser::token::TokenType, @@ -395,7 +395,7 @@ impl SimpleEngine { /// Recursively unwraps `Paren`, then structurally resolves a leaf PromQL /// arm (i.e. not `Binary` or `NumberLiteral`) to its `QueryConfig` and /// base `QueryExecutionContext`. Shared leaf-resolution step for both - /// `build_arm_logical_plan` (instant) and `build_arm_range_context` (range). + /// `evaluate_binary_arm` (instant) and `build_arm_range_context` (range). /// /// Returns `None` for `Binary` arms (caller handles recursion) and /// `NumberLiteral` arms (caller handles scalars). @@ -418,115 +418,17 @@ impl SimpleEngine { } } - /// Recursively builds a DataFusion logical plan for one arm of a binary - /// arithmetic expression. - /// - /// - Leaf arm (supported PromQL pattern): resolved via `resolve_arm_leaf_context`, - /// returning its `to_logical_plan()` together with the output label names. - /// - Binary arm: recursively build both sub-arms and combine with - /// `build_binary_vector_plan`. - /// - Scalar literal: returns `None` (handled by the caller separately). - fn build_arm_logical_plan( - &self, - arm_ast: &promql_parser::parser::Expr, - time: f64, - ) -> Option<(datafusion::logical_expr::LogicalPlan, Vec)> { - use crate::engines::logical::plan_builder::build_binary_vector_plan; - use promql_parser::parser::Expr; - - match arm_ast { - Expr::NumberLiteral(_) => None, // caller handles scalars - Expr::Paren(paren) => self.build_arm_logical_plan(&paren.expr, time), - Expr::Binary(binary) => { - // Nested binary expression — recurse on both sides - let (lhs_plan, lhs_labels) = self.build_arm_logical_plan(&binary.lhs, time)?; - let (rhs_plan, _) = self.build_arm_logical_plan(&binary.rhs, time)?; - let combined = - build_binary_vector_plan(lhs_plan, rhs_plan, &binary.op, lhs_labels.clone()) - .ok()?; - Some((combined, lhs_labels)) - } - _ => { - let (ctx, label_names) = self.resolve_arm_leaf_context(arm_ast, time)?; - let plan = ctx.to_logical_plan().ok()?; - Some((plan, label_names)) - } - } - } - - /// Handles a binary arithmetic PromQL expression by building a combined - /// DataFusion plan (vector–vector join or scalar projection) and executing it. - /// - /// Returns `None` if any arm is not acceleratable (caller falls back to Prometheus). - fn handle_binary_expr_promql( - &self, - ast: &promql_parser::parser::Expr, - time: f64, - ) -> Option<(KeyByLabelNames, QueryResult)> { - use crate::engines::logical::plan_builder::{build_binary_vector_plan, build_scalar_plan}; - use promql_parser::parser::Expr; - - let query_time = Self::convert_query_time_to_data_time(time); - - let binary = match ast { - Expr::Binary(b) => b, - _ => 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 (vector_plan, label_names) = self.build_arm_logical_plan(vector_arm, time)?; - let combined = - build_scalar_plan(vector_plan, scalar, op, scalar_on_left, label_names.clone()) - .ok()?; - let results = tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on(self.execute_logical_plan( - combined, - label_names.clone(), - "", - &Statistic::Sum, - )) - }) - .ok()?; - return Some(( - KeyByLabelNames::new(label_names), - QueryResult::vector(results, query_time), - )); - } - - // Vector–vector - let (lhs_plan, lhs_labels) = self.build_arm_logical_plan(lhs, time)?; - let (rhs_plan, _) = self.build_arm_logical_plan(rhs, time)?; - let combined = build_binary_vector_plan(lhs_plan, rhs_plan, op, lhs_labels.clone()).ok()?; - let results = tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on(self.execute_logical_plan( - combined, - lhs_labels.clone(), - "", - &Statistic::Sum, - )) - }) - .ok()?; - let output_labels = KeyByLabelNames::new(lhs_labels); - Some((output_labels, QueryResult::vector(results, query_time))) - } - /// Recursively evaluates one arm of a binary arithmetic expression via - /// the native pipeline (`execute_query_pipeline`), instead of building a - /// DataFusion plan. Mirrors `build_arm_logical_plan`'s shape exactly, - /// including its limitation: nested `Binary` arms are combined as - /// vector-vector only — a scalar inside a nested arm (e.g. `(a+5)*b`) - /// is not supported, matching today's DataFusion path. + /// the native pipeline (`execute_query_pipeline`). /// - /// - Leaf arm: resolved via `resolve_arm_leaf_context`, executed through - /// `execute_query_pipeline`. - /// - Binary arm: recursively evaluate both sides then combine with - /// `combine_vector_vector_native`. + /// - Leaf arm (supported PromQL pattern): resolved via `resolve_arm_leaf_context`, + /// executed through `execute_query_pipeline`. + /// - Binary arm: recursively evaluate both sub-arms and combine with + /// `combine_vector_vector`. Nested `Binary` arms are combined as + /// vector-vector only — a scalar inside a nested arm (e.g. `(a+5)*b`) + /// is not supported (tracked separately, same as before this cutover). /// - Scalar literal: returns `None` (handled by the caller separately). - fn evaluate_arm_native( + fn evaluate_binary_arm( &self, arm_ast: &promql_parser::parser::Expr, time: f64, @@ -535,12 +437,12 @@ impl SimpleEngine { match arm_ast { Expr::NumberLiteral(_) => None, // caller handles scalars - Expr::Paren(paren) => self.evaluate_arm_native(&paren.expr, time), + Expr::Paren(paren) => self.evaluate_binary_arm(&paren.expr, time), Expr::Binary(binary) => { // Nested binary expression — recurse on both sides - let (lhs_results, lhs_labels) = self.evaluate_arm_native(&binary.lhs, time)?; - let (rhs_results, rhs_labels) = self.evaluate_arm_native(&binary.rhs, time)?; - let combined = combine_vector_vector_native( + let (lhs_results, lhs_labels) = self.evaluate_binary_arm(&binary.lhs, time)?; + let (rhs_results, rhs_labels) = self.evaluate_binary_arm(&binary.rhs, time)?; + let combined = combine_vector_vector( lhs_results, &lhs_labels, rhs_results, @@ -551,9 +453,9 @@ impl SimpleEngine { } _ => { let (ctx, label_names) = self.resolve_arm_leaf_context(arm_ast, time)?; - // Unlike DataFusion's PrecomputedSummaryReadExec (which streams - // whatever rows exist, including zero, so a currently-empty arm - // still returns Some(empty vector)), execute_query_pipeline errors + // Unlike DataFusion's PrecomputedSummaryReadExec (which streamed + // whatever rows existed, including zero, so a currently-empty arm + // used to return Some(empty vector)), execute_query_pipeline errors // when the store has no precomputed outputs at all for this arm — // that propagates to None here, triggering a full Prometheus // fallback for the whole expression instead of an empty result @@ -567,10 +469,9 @@ impl SimpleEngine { .execute_query_pipeline(&ctx, true, true) .map_err(|e| { warn!( - "Native binary-expr arm for metric '{}' produced no results \ - ({}) — unlike DataFusion, this falls back to Prometheus for \ - the whole expression rather than returning an empty result \ - for just this arm", + "Binary-expr arm for metric '{}' produced no results ({}) — \ + falls back to Prometheus for the whole expression rather than \ + returning an empty result for just this arm", ctx.metric, e ); e @@ -581,13 +482,11 @@ impl SimpleEngine { } } - /// Native (non-DataFusion) implementation of binary-expr handling, - /// staged ahead of #567's cutover so it can be equivalence-tested - /// against the DataFusion path (`handle_binary_expr_promql`) before - /// that path is rewired to call this one. Not yet reachable from - /// `handle_query_promql` — see `handle_query_promql_native` for a - /// test-facing entrypoint. - pub fn handle_binary_expr_promql_native( + /// Handles a binary arithmetic PromQL expression via the native pipeline + /// (vector–vector join or scalar combine). + /// + /// Returns `None` if any arm is not acceleratable (caller falls back to Prometheus). + fn handle_binary_expr_promql( &self, ast: &promql_parser::parser::Expr, time: f64, @@ -606,8 +505,8 @@ impl SimpleEngine { let op = &binary.op; if let Some((scalar, vector_arm, scalar_on_left)) = detect_scalar_arm(lhs, rhs) { - let (vector_results, label_names) = self.evaluate_arm_native(vector_arm, time)?; - let combined = combine_scalar_native(vector_results, scalar, op, scalar_on_left); + let (vector_results, label_names) = self.evaluate_binary_arm(vector_arm, time)?; + let combined = combine_scalar(vector_results, scalar, op, scalar_on_left); return Some(( KeyByLabelNames::new(label_names), QueryResult::vector(combined, query_time), @@ -615,26 +514,14 @@ impl SimpleEngine { } // Vector–vector - let (lhs_results, lhs_labels) = self.evaluate_arm_native(lhs, time)?; - let (rhs_results, rhs_labels) = self.evaluate_arm_native(rhs, time)?; + let (lhs_results, lhs_labels) = self.evaluate_binary_arm(lhs, time)?; + let (rhs_results, rhs_labels) = self.evaluate_binary_arm(rhs, time)?; let combined = - combine_vector_vector_native(lhs_results, &lhs_labels, rhs_results, &rhs_labels, op)?; + combine_vector_vector(lhs_results, &lhs_labels, rhs_results, &rhs_labels, op)?; let output_labels = KeyByLabelNames::new(lhs_labels); Some((output_labels, QueryResult::vector(combined, query_time))) } - /// Parses `query` and dispatches to `handle_binary_expr_promql_native`. - /// Test-facing convenience wrapper — mirrors the binary-expr branch of - /// `handle_query_promql`, but for the native path. - pub fn handle_query_promql_native( - &self, - query: String, - time: f64, - ) -> Option<(KeyByLabelNames, QueryResult)> { - let ast = promql_parser::parser::parse(&query).ok()?; - self.handle_binary_expr_promql_native(&ast, time) - } - /// Applies a PromQL binary arithmetic operator to two f64 values. fn apply_range_binary_op( op: &promql_parser::parser::token::TokenType, @@ -657,7 +544,7 @@ impl SimpleEngine { /// arithmetic expression. /// /// Leaf resolution (Paren-unwrap + structural config lookup) is shared - /// with `build_arm_logical_plan` via `resolve_arm_leaf_context`. Note this + /// with `evaluate_binary_arm` via `resolve_arm_leaf_context`. Note this /// does not support nested `Binary` arms (e.g. `(a+b)*c` over a range) — /// tracked separately in #516. fn build_arm_range_context( diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs index 91d1356d..69a119ea 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs @@ -2,7 +2,8 @@ //! //! Verify that binary arithmetic queries (vector/vector and scalar/vector) //! produce numerically correct results when executed end-to-end through -//! `handle_binary_expr_promql` via DataFusion. +//! `handle_binary_expr_promql`, natively as of #567's Stage 3 cutover +//! (previously via DataFusion). #[cfg(test)] mod tests { diff --git a/asap-query-engine/src/tests/native_binary_instant_tests.rs b/asap-query-engine/src/tests/native_binary_instant_tests.rs index 1599d8c4..4f0bd30c 100644 --- a/asap-query-engine/src/tests/native_binary_instant_tests.rs +++ b/asap-query-engine/src/tests/native_binary_instant_tests.rs @@ -1,22 +1,21 @@ -//! Native instant binary-expr evaluator tests (issue #567, Stage 2). +//! PromQL binary-expr instant query tests, native execution (issue #567). //! -//! `handle_binary_expr_promql_native`/`evaluate_arm_native` are a new, -//! parallel implementation of PromQL binary-arithmetic instant queries, -//! built to eventually replace the DataFusion-backed -//! `handle_binary_expr_promql` (#567's Stage 3 cutover). These tests compare -//! the two paths directly — both are reachable today (`handle_query_promql` -//! for DataFusion, `handle_query_promql_native` for the new path) so -//! equivalence can be proven before anything is rewired. +//! `handle_query_promql`'s binary-arithmetic path (`handle_binary_expr_promql` +//! → `evaluate_binary_arm` → `combine_vector_vector`/`combine_scalar`) runs +//! natively as of #567's Stage 3 cutover — no more DataFusion involved. +//! These tests were originally written to compare the native path against +//! the (now-removed) DataFusion path before the cutover landed; they now +//! assert the native path's results directly. #[cfg(test)] mod tests { - use crate::data_model::{AggregationType, KeyByLabelValues}; + use crate::data_model::{AggregationType, KeyByLabelValues, WindowType}; use crate::engines::query_result::QueryResult; use crate::precompute_operators::sum_accumulator::SumAccumulator; use crate::precompute_operators::{CountMinSketchAccumulator, DeltaSetAggregatorAccumulator}; use crate::tests::test_utilities::engine_factories::{ - create_engine_dual_input, create_engine_single_pop, create_engine_three_metrics, - create_engine_two_metrics, + create_engine_dual_input, create_engine_multi_timestamp_with_window, + create_engine_single_pop, create_engine_three_metrics, create_engine_two_metrics, }; use crate::AggregateCore; @@ -39,11 +38,9 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - async fn native_vector_vector_all_ops_match_datafusion() { - // `^` excluded: DataFusion's build_binary_vector_plan maps T_POW to - // Operator::BitwiseXor "as a proxy" (plan_builder.rs comment) and - // fails to even produce a result for it today — nothing to compare - // native against for that operator. See the dedicated `^` test below. + async fn binary_expr_vector_vector_all_ops() { + // `^` excluded: needs its own accumulator setup (see the dedicated + // power-operator test below). for (op, expected) in [ ("+", 30.0), ("-", -10.0), @@ -71,33 +68,22 @@ mod tests { ); let query = format!("sum(metric_a) by (host) {op} sum(metric_b) by (host)"); - let (_, old_qr) = engine - .handle_query_promql(query.clone(), QUERY_TIME) - .unwrap_or_else(|| panic!("old path failed for op {op}")); - let (_, new_qr) = engine - .handle_query_promql_native(query.clone(), QUERY_TIME) - .unwrap_or_else(|| panic!("new path failed for op {op}")); - - let old_values = sorted(vector_values(old_qr)); - let new_values = sorted(vector_values(new_qr)); - assert_eq!(old_values.len(), 1, "op {op}"); - assert_eq!( - new_values, old_values, - "op {op}: native must match DataFusion" - ); + let (_, qr) = engine + .handle_query_promql(query, QUERY_TIME) + .unwrap_or_else(|| panic!("query failed for op {op}")); + + let values = vector_values(qr); + assert_eq!(values.len(), 1, "op {op}"); assert!( - (new_values[0].1 - expected).abs() < 1e-10, + (values[0].1 - expected).abs() < 1e-10, "op {op}: expected {expected}, got {}", - new_values[0].1 + values[0].1 ); } } #[tokio::test(flavor = "multi_thread")] - async fn native_power_operator_computes_correctly_datafusion_unsupported() { - // DataFusion's `^` support is broken today (see comment above); this - // documents that native computes it correctly via f64::powf, - // independent of a DataFusion comparison. + async fn binary_expr_power_operator_computes_correctly() { let engine = create_engine_two_metrics( "metric_a", AggregationType::Sum, @@ -118,15 +104,15 @@ mod tests { ); let query = "sum(metric_a) by (host) ^ sum(metric_b) by (host)"; - let (_, new_qr) = engine - .handle_query_promql_native(query.to_string(), QUERY_TIME) - .expect("new path failed"); - let new_values = vector_values(new_qr); - assert!((new_values[0].1 - 1024.0).abs() < 1e-6, "2^10 = 1024"); + let (_, qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("query failed"); + let values = vector_values(qr); + assert!((values[0].1 - 1024.0).abs() < 1e-6, "2^10 = 1024"); } #[tokio::test(flavor = "multi_thread")] - async fn native_vector_scalar_both_orderings_match_datafusion() { + async fn binary_expr_scalar_both_orderings() { let engine = create_engine_single_pop( "errors_total", AggregationType::Sum, @@ -142,22 +128,16 @@ mod tests { "sum(errors_total) by (host) * 100", "100 * sum(errors_total) by (host)", ] { - let (_, old_qr) = engine + let (_, qr) = engine .handle_query_promql(query.to_string(), QUERY_TIME) - .unwrap_or_else(|| panic!("old path failed for {query}")); - let (_, new_qr) = engine - .handle_query_promql_native(query.to_string(), QUERY_TIME) - .unwrap_or_else(|| panic!("new path failed for {query}")); - - let old_values = vector_values(old_qr); - let new_values = vector_values(new_qr); - assert_eq!(new_values, old_values, "{query}"); - assert!((new_values[0].1 - 700.0).abs() < 1e-10, "{query}"); + .unwrap_or_else(|| panic!("query failed for {query}")); + let values = vector_values(qr); + assert!((values[0].1 - 700.0).abs() < 1e-10, "{query}"); } } #[tokio::test(flavor = "multi_thread")] - async fn native_nested_binary_matches_datafusion() { + async fn binary_expr_nested_binary() { // (metric_a + metric_b) * metric_c, all Sum, host-a: (10+20)*3 = 90 let engine = create_engine_three_metrics( "metric_a", @@ -187,26 +167,21 @@ mod tests { ); let query = "(sum(metric_a) by (host) + sum(metric_b) by (host)) * sum(metric_c) by (host)"; - let (_, old_qr) = engine + let (_, qr) = engine .handle_query_promql(query.to_string(), QUERY_TIME) - .expect("old path failed"); - let (_, new_qr) = engine - .handle_query_promql_native(query.to_string(), QUERY_TIME) - .expect("new path failed"); - - let old_values = vector_values(old_qr); - let new_values = vector_values(new_qr); - assert_eq!(new_values, old_values); - assert!((new_values[0].1 - 90.0).abs() < 1e-10, "(10+20)*3 = 90"); + .expect("query failed"); + let values = vector_values(qr); + assert!((values[0].1 - 90.0).abs() < 1e-10, "(10+20)*3 = 90"); } #[tokio::test(flavor = "multi_thread")] - async fn native_binary_expr_no_data_falls_back_to_none() { + async fn binary_expr_no_data_falls_back_to_none() { // metric_a is configured (schema + pattern match) but has zero - // precomputed data. Accepted behavior change (#567): unlike - // DataFusion (which would return an empty-but-present result), - // native falls back to Prometheus (returns None) for the whole - // expression — see evaluate_arm_native's leaf branch. + // precomputed data. Accepted behavior change (#567, kept from + // DataFusion's now-removed empty-result behavior): falls back to + // Prometheus (returns None) for the whole expression — see + // evaluate_binary_arm's leaf branch, which warns loudly when this + // happens. let engine = create_engine_two_metrics( "metric_a", AggregationType::Sum, @@ -224,23 +199,17 @@ mod tests { ); let query = "sum(metric_a) by (host) + sum(metric_b) by (host)"; - let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); assert!( - old_result.is_some(), - "DataFusion returns an empty-but-present result, not None, for a currently-empty arm" - ); - - let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); - assert!( - new_result.is_none(), - "native falls back to Prometheus (None) for a currently-empty arm" + result.is_none(), + "arm with no current precomputed data falls back to Prometheus" ); } #[tokio::test(flavor = "multi_thread")] - async fn native_binary_expr_unsupported_arm_returns_none() { - // foo() is not a supported PromQL function -> arm lookup fails -> None, - // same as today's DataFusion path (dispatch_arithmetic_tests.rs). + async fn binary_expr_unsupported_arm_returns_none() { + // foo() is not a supported PromQL function -> arm lookup fails -> None + // (graceful fallback to Prometheus). let engine = create_engine_single_pop( "requests_total", AggregationType::Sum, @@ -256,18 +225,14 @@ mod tests { assert!(engine .handle_query_promql(query.to_string(), QUERY_TIME) .is_none()); - assert!(engine - .handle_query_promql_native(query.to_string(), QUERY_TIME) - .is_none()); } #[tokio::test(flavor = "multi_thread")] - async fn native_vector_vector_dual_population_matches_datafusion() { + async fn binary_expr_vector_vector_dual_population() { // event_frequency is dual-population (CountMinSketch values + - // DeltaSetAggregator keys) -- confirms the leaf swap - // (ctx.to_logical_plan() -> execute_query_pipeline) still resolves - // keys_query correctly. Wrapped in `+ 0` to route through the - // binary-expr handler at all. + // DeltaSetAggregator keys) -- confirms the leaf's keys_query + // resolution works through the binary-expr path. Wrapped in `+ 0` + // to route through the binary-expr handler at all. let cms = CountMinSketchAccumulator::new(2, 3); let mut keys = DeltaSetAggregatorAccumulator::new(); keys.add_key(KeyByLabelValues { @@ -286,34 +251,95 @@ mod tests { ); let query = "count(event_frequency) by (host, event) + 0"; - let (_, old_qr) = engine + let (_, qr) = engine .handle_query_promql(query.to_string(), QUERY_TIME) - .expect("old path failed"); - let (_, new_qr) = engine - .handle_query_promql_native(query.to_string(), QUERY_TIME) - .expect("new path failed"); + .expect("query failed"); + assert!(!sorted(vector_values(qr)).is_empty()); + } - assert_eq!(sorted(vector_values(new_qr)), sorted(vector_values(old_qr))); + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_sliding_window_end_to_end_merges_correctly() { + // Ties Stage 1's sliding-bucket merge fix (#570) to the actual + // production entrypoint this issue changes: 2 buckets for the same + // key under one Sliding exact window must both be merged, not just + // the first, when reached through a real binary-expr query. + let data = vec![ + ( + 1_000_000, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ), + ( + 1_000_000, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(5.0)) as Box, + ), + ]; + let leaf_query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + leaf_query, + 1_000, // window_size_ms, matches the fixed 1000ms bucket width + WindowType::Sliding, + ); + + let query = format!("{leaf_query} + 0"); + let (_, qr) = engine + .handle_query_promql(query, QUERY_TIME) + .expect("query failed"); + let values = vector_values(qr); + assert_eq!(values.len(), 1); + assert!( + (values[0].1 - 15.0).abs() < 1e-10, + "expected both sliding-window buckets merged into 15.0, got {}", + values[0].1 + ); } - // --- Regression tests: combine_vector_vector_native joins purely on - // positional KeyByLabelValues equality (raw Vec of values, no - // label names attached at all -- see key_by_label_values.rs). DataFusion's - // build_binary_vector_plan joins on named columns (lhs.{label} = - // rhs.{label}) instead, so it fails to resolve (-> None) whenever the two - // arms don't share the same label set. Native has no equivalent check, so - // it can silently join two arms grouped by *different* labels whenever - // their values happen to coincide. These tests are expected to be RED - // against the current implementation -- they exist to prove the - // divergence before combine_vector_vector_native is fixed. + #[tokio::test] + async fn binary_expr_works_on_current_thread_runtime() { + // Default (single-threaded) tokio runtime, not `flavor = "multi_thread"` + // like every other test in this file. The old DataFusion path's + // tokio::task::block_in_place(...block_on(...)) wrapper panics on a + // current-thread runtime — this test only passes because that + // wrapper is actually gone (#567 Stage 3 cutover), not because of + // any value it computes. + let engine = create_engine_single_pop( + "errors_total", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(7.0)) as Box, + )], + "sum(errors_total) by (host)", + ); + + let result = + engine.handle_query_promql("sum(errors_total) by (host) * 2".to_string(), QUERY_TIME); + assert!(result.is_some()); + } + + // --- Regression tests: combine_vector_vector must reject a join between + // arms with different label sets rather than silently matching on + // positional KeyByLabelValues equality alone (raw Vec of values, + // no label names attached -- see key_by_label_values.rs). This mirrors + // DataFusion's build_binary_vector_plan, which fails to resolve a join + // column that only exists on one side. These tests originally compared + // against the DataFusion path to prove the divergence before the fix; + // now that combine_vector_vector checks label-set equality directly and + // DataFusion is no longer in the production path (#567 Stage 3), they + // assert the fixed behavior directly. #[tokio::test(flavor = "multi_thread")] - async fn native_mismatched_label_names_with_colliding_values_diverges_from_datafusion() { + async fn binary_expr_mismatched_label_sets_with_colliding_values_returns_none() { // metric_a grouped by (host), metric_b grouped by (region) -- disjoint - // label sets -- but both happen to produce the value "us-east". - // DataFusion can't join lhs.host against a plan with no "host" column - // -> None. Native joins on the bare value vector ["us-east"] == - // ["us-east"] -> a spurious combined result. + // label sets -- but both happen to produce the value "us-east". A + // value-only join would spuriously match ["us-east"] == ["us-east"]; + // must return None instead. let engine = create_engine_two_metrics( "metric_a", AggregationType::Sum, @@ -334,28 +360,19 @@ mod tests { ); let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; - let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); - - match old_result { - None => assert!( - new_result.is_none(), - "BUG: DataFusion refuses to join arms with different label sets (None), \ - but native silently produced {new_result:?} by matching on value alone" - ), - Some((_, old_qr)) => assert_eq!( - new_result.map(|(_, qr)| sorted(vector_values(qr))), - Some(sorted(vector_values(old_qr))), - "native and DataFusion must agree when both produce a result" - ), - } + let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + + assert!( + result.is_none(), + "mismatched label sets must return None, not a spurious value-matched join: {result:?}" + ); } #[tokio::test(flavor = "multi_thread")] - async fn native_mismatched_label_names_multi_label_full_collision_diverges_from_datafusion() { - // Same bug with a 2-label grouping: (host, dc) vs (region, zone), but + async fn binary_expr_mismatched_label_sets_multi_label_full_collision_returns_none() { + // Same case with a 2-label grouping: (host, dc) vs (region, zone), but // the *entire* ordered value vector coincides ("us-east", "az1" on - // both sides) -- confirms the collision isn't a single-label fluke. + // both sides) -- confirms the check isn't a single-label fluke. let engine = create_engine_two_metrics( "metric_a", AggregationType::Sum, @@ -376,31 +393,20 @@ mod tests { ); let query = "sum(metric_a) by (host, dc) + sum(metric_b) by (region, zone)"; - let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); - - match old_result { - None => assert!( - new_result.is_none(), - "BUG: DataFusion refuses to join arms with different label sets (None), \ - but native silently produced {new_result:?} by matching on value alone" - ), - Some((_, old_qr)) => assert_eq!( - new_result.map(|(_, qr)| sorted(vector_values(qr))), - Some(sorted(vector_values(old_qr))), - "native and DataFusion must agree when both produce a result" - ), - } + let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + + assert!( + result.is_none(), + "mismatched label sets must return None, not a spurious value-matched join: {result:?}" + ); } #[tokio::test(flavor = "multi_thread")] - async fn native_mismatched_label_names_non_colliding_values_still_diverges_from_datafusion() { + async fn binary_expr_mismatched_label_sets_non_colliding_values_returns_none() { // Same disjoint label sets (host vs region), but this time the values - // don't collide ("us-east" vs "eu-west"). DataFusion still can't - // resolve the join (column doesn't exist, independent of the data) - // -> None. Native's join correctly finds no match on the *value*, but - // that's incidental -- it still returns Some(empty) instead of None, - // because it never checked whether the label sets matched at all. + // don't collide ("us-east" vs "eu-west") either -- must still return + // None because the label sets themselves don't match, not because no + // values happened to match. let engine = create_engine_two_metrics( "metric_a", AggregationType::Sum, @@ -421,23 +427,19 @@ mod tests { ); let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; - let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); - - assert_eq!( - old_result.is_none(), - new_result.is_none(), - "BUG: DataFusion returns {old_result:?} (no such column to join on), native \ - returns {new_result:?} -- native never checks that the two arms' label sets match" + let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + + assert!( + result.is_none(), + "mismatched label sets must return None even when no values happen to match: {result:?}" ); } #[tokio::test(flavor = "multi_thread")] - async fn native_same_label_set_no_matching_values_both_agree_empty() { + async fn binary_expr_same_label_set_no_matching_values_returns_empty() { // Control case: identical label sets (host on both sides) but - // disjoint values -- both paths should agree there's no match. This - // isolates the bug to *differing* label sets, not "no match" in - // general. + // disjoint values -- should resolve to an empty (not None) result. + // This isolates the label-set check from ordinary "no match" cases. let engine = create_engine_two_metrics( "metric_a", AggregationType::Sum, @@ -458,15 +460,10 @@ mod tests { ); let query = "sum(metric_a) by (host) + sum(metric_b) by (host)"; - let (_, old_qr) = engine + let (_, qr) = engine .handle_query_promql(query.to_string(), QUERY_TIME) - .expect("old path failed"); - let (_, new_qr) = engine - .handle_query_promql_native(query.to_string(), QUERY_TIME) - .expect("new path failed"); - - let old_values = vector_values(old_qr); - assert_eq!(old_values, Vec::new()); - assert_eq!(sorted(vector_values(new_qr)), sorted(old_values)); + .expect("query failed"); + + assert_eq!(vector_values(qr), Vec::new()); } } From 9d43776dc7ad72722490a093abe5b330e485697b Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Fri, 21 Aug 2026 17:24:27 -0400 Subject: [PATCH 08/10] fix(query-engine): address code-review cleanup findings on PR #567 Stage 3 cutover - move plan_execution_arithmetic_tests.rs out of tests/datafusion/ since it now exercises the native binary-expr path, not DataFusion - fix warn! in evaluate_binary_arm that mislabeled any execute_query_pipeline error as "produced no results" - mark orphaned execute_logical_plan #[allow(dead_code)] and update its doc comment, matching its sibling execute_plan - avoid an unnecessary Vec clone in combine_vector_vector - mark design-252 doc as superseded by the native cutover in #567 --- asap-query-engine/src/engines/simple_engine/mod.rs | 7 +++++-- asap-query-engine/src/engines/simple_engine/promql.rs | 6 +++--- asap-query-engine/src/tests/datafusion/mod.rs | 1 - asap-query-engine/src/tests/mod.rs | 1 + ...c_tests.rs => native_binary_arithmetic_plan_tests.rs} | 0 docs/design-252-arithmetic-operators.md | 9 +++++++++ 6 files changed, 18 insertions(+), 6 deletions(-) rename asap-query-engine/src/tests/{datafusion/plan_execution_arithmetic_tests.rs => native_binary_arithmetic_plan_tests.rs} (100%) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 4cf0fbd2..90f4d74a 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -823,8 +823,11 @@ impl SimpleEngine { /// Executes a pre-built DataFusion logical plan and returns results. /// - /// This is the shared execution kernel used by both `execute_plan` (for single-metric - /// queries) and the binary arithmetic dispatch path. + /// This was the shared execution kernel for `execute_plan` and the DataFusion-based + /// binary arithmetic dispatch path; the latter was cut over to a native implementation + /// in #567, leaving this unused in production. Kept alongside `execute_plan` as part + /// of the still-exercised DataFusion path (see its dedicated tests). + #[allow(dead_code)] pub async fn execute_logical_plan( &self, logical_plan: datafusion::logical_expr::LogicalPlan, diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 970596ec..4aed87fd 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -66,9 +66,9 @@ fn combine_vector_vector( lhs_results .into_iter() .filter_map(|lhs_elem| { - rhs_map.get(&lhs_elem.labels).map(|&rhs_val| { + rhs_map.get(&lhs_elem.labels).map(move |&rhs_val| { let value = SimpleEngine::apply_range_binary_op(op, lhs_elem.value, rhs_val); - InstantVectorElement::new(lhs_elem.labels.clone(), value) + InstantVectorElement::new(lhs_elem.labels, value) }) }) .collect(), @@ -469,7 +469,7 @@ impl SimpleEngine { .execute_query_pipeline(&ctx, true, true) .map_err(|e| { warn!( - "Binary-expr arm for metric '{}' produced no results ({}) — \ + "Binary-expr arm for metric '{}' failed ({}) — \ falls back to Prometheus for the whole expression rather than \ returning an empty result for just this arm", ctx.metric, e diff --git a/asap-query-engine/src/tests/datafusion/mod.rs b/asap-query-engine/src/tests/datafusion/mod.rs index 143cc0a4..fd946643 100644 --- a/asap-query-engine/src/tests/datafusion/mod.rs +++ b/asap-query-engine/src/tests/datafusion/mod.rs @@ -7,7 +7,6 @@ pub mod accumulator_serde_tests; pub mod dispatch_arithmetic_tests; pub mod plan_builder_binary_tests; pub mod plan_builder_regression_tests; -pub mod plan_execution_arithmetic_tests; pub mod plan_execution_dual_input_tests; pub mod plan_execution_temporal_tests; pub mod plan_execution_tests; diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index d5b31eec..6ac84281 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -3,6 +3,7 @@ pub mod clickhouse_forwarding_tests; pub mod datafusion; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests; +pub mod native_binary_arithmetic_plan_tests; pub mod native_binary_instant_tests; pub mod native_pipeline_merge_tests; pub mod prometheus_forwarding_tests; diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs b/asap-query-engine/src/tests/native_binary_arithmetic_plan_tests.rs similarity index 100% rename from asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs rename to asap-query-engine/src/tests/native_binary_arithmetic_plan_tests.rs diff --git a/docs/design-252-arithmetic-operators.md b/docs/design-252-arithmetic-operators.md index 541f2fc2..fc5f22d5 100644 --- a/docs/design-252-arithmetic-operators.md +++ b/docs/design-252-arithmetic-operators.md @@ -1,5 +1,14 @@ # Design: PromQL Arithmetic Operator Acceleration (Issue #252) +> **Superseded (#567):** the DataFusion execution path described below +> (`execute_plan`/`execute_logical_plan`, the `Join + Projection` plan) was +> replaced by a purely native execution path in `handle_binary_expr_promql` / +> `evaluate_binary_arm` / `combine_vector_vector` / `combine_scalar` +> (`asap-query-engine/src/engines/simple_engine/promql.rs`). The DataFusion +> code is kept only for its own dedicated tests and is no longer reachable +> from production. This doc is retained as the historical record of the +> original design decision. + ## Problem ASAPQuery accelerates PromQL queries by pre-computing sketches over streaming data and serving answers from those sketches at query time, bypassing the underlying TSDB for supported query patterns. From bc580c2b54f129eddc3390c0d66d1ace2b531f30 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 22 Aug 2026 22:30:12 -0400 Subject: [PATCH 09/10] fixup(query-engine): update topk-in-binary-expr test for post-cutover API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit handle_query_promql_native was removed by the #567 Stage 3 cutover (native is now the only path, folded into handle_query_promql) — this commit was squashed into the branch during rebase onto the cutover. Co-Authored-By: Claude Sonnet 5 --- asap-query-engine/src/engines/simple_engine/promql.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 4aed87fd..c11bb509 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -1572,8 +1572,8 @@ mod topk_pipeline_tests { .expect("insert should succeed"); let (_, query_result) = engine - .handle_query_promql_native(format!("{TOPK_QUERY} + 0"), QUERY_TIME) - .expect("binary-expr-wrapped topk should still resolve via the native path"); + .handle_query_promql(format!("{TOPK_QUERY} + 0"), QUERY_TIME) + .expect("binary-expr-wrapped topk should still resolve"); let results = match query_result { QueryResult::Vector(iv) => iv.values, From b9287d85713a22e5257f5a457fc658578548fe26 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 22 Aug 2026 22:45:17 -0400 Subject: [PATCH 10/10] fix(query-engine): reject cross-label-set joins in range binary-expr; fix stale dead-code comment handle_binary_expr_range_promql's vector-vector join matched purely on positional KeyByLabelValues equality (rhs labels discarded), unlike the instant-query combine_vector_vector fixed earlier in this stack (#572) to reject a join between arms grouped by different label sets. Two arms grouped by disjoint labels (e.g. (host) vs (region)) whose values happened to coincide could silently join into a wrong-but-plausible result across the whole range. Now rejects the join (returns None) when lhs_labels != rhs_labels, mirroring the instant-query guard. Extends create_range_engine_two_metrics (range_query_arithmetic_tests.rs) to take per-metric grouping labels, and adds a regression test proving the divergence. Also corrects execute_logical_plan's doc comment: it claimed to be "part of the still-exercised DataFusion path (see its dedicated tests)", but it has zero callers anywhere in the repo, including tests -- unlike its sibling execute_plan, which genuinely is still called by DataFusion-path tests. Co-Authored-By: Claude Sonnet 5 --- .../src/engines/simple_engine/mod.rs | 10 ++-- .../src/engines/simple_engine/promql.rs | 12 ++++- .../range_query_arithmetic_tests.rs | 51 +++++++++++++++++-- 3 files changed, 63 insertions(+), 10 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 90f4d74a..f4e4c514 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -823,10 +823,12 @@ impl SimpleEngine { /// Executes a pre-built DataFusion logical plan and returns results. /// - /// This was the shared execution kernel for `execute_plan` and the DataFusion-based - /// binary arithmetic dispatch path; the latter was cut over to a native implementation - /// in #567, leaving this unused in production. Kept alongside `execute_plan` as part - /// of the still-exercised DataFusion path (see its dedicated tests). + /// This was the entry point for the DataFusion-based binary arithmetic + /// dispatch path, cut over to a native implementation in #567. Unlike its + /// sibling `execute_plan` (still called by DataFusion-path tests), this + /// function has zero callers anywhere in the repo, including tests — it + /// is genuinely dead code, kept only in case the native cutover needs to + /// be reverted. #[allow(dead_code)] pub async fn execute_logical_plan( &self, diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index c11bb509..8b37f843 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -651,9 +651,17 @@ impl SimpleEngine { return Some((KeyByLabelNames::new(labels), QueryResult::matrix(combined))); } - // Vector-vector: evaluate both arms, join by label key, apply op per matching timestamp + // Vector-vector: evaluate both arms, join by label key, apply op per matching timestamp. + // Reject mismatched label sets up front — same guard as the instant-query + // combine_vector_vector, and for the same reason: positional + // KeyByLabelValues equality below is only safe once the label *names* + // match (they're canonically sorted by KeyByLabelNames::new(), so two + // arms with the same label set always order their values the same way). let (lhs_ctx, lhs_labels) = self.build_arm_range_context(lhs, start, end, step)?; - let (rhs_ctx, _) = self.build_arm_range_context(rhs, start, end, step)?; + let (rhs_ctx, rhs_labels) = self.build_arm_range_context(rhs, start, end, step)?; + 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()?; diff --git a/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs b/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs index a2fca523..cdf67201 100644 --- a/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs @@ -36,18 +36,22 @@ mod tests { /// `engine_factories::create_engine_two_metrics` (single timestamp, instant /// queries only), this inserts one bucket per `(timestamp, value)` pair so /// range queries have more than one output sample to join across. + #[allow(clippy::too_many_arguments)] fn create_range_engine_two_metrics( metric_a: &str, + labels_a: Vec<&str>, data_a: TimeSeriesData, query_a: &str, metric_b: &str, + labels_b: Vec<&str>, data_b: TimeSeriesData, query_b: &str, ) -> SimpleEngine { - let labels = vec!["host".to_string()]; + let labels_a: Vec = labels_a.iter().map(|s| s.to_string()).collect(); + let labels_b: Vec = labels_b.iter().map(|s| s.to_string()).collect(); let mut aggregation_configs = HashMap::new(); - for (id, metric) in [(1u64, metric_a), (2u64, metric_b)] { + for (id, metric, labels) in [(1u64, metric_a, &labels_a), (2u64, metric_b, &labels_b)] { aggregation_configs.insert( id, AggregationConfig { @@ -91,8 +95,8 @@ mod tests { } let promql_schema = PromQLSchema::new() - .add_metric(metric_a.to_string(), KeyByLabelNames::new(labels.clone())) - .add_metric(metric_b.to_string(), KeyByLabelNames::new(labels)); + .add_metric(metric_a.to_string(), KeyByLabelNames::new(labels_a)) + .add_metric(metric_b.to_string(), KeyByLabelNames::new(labels_b)); let inference_config = InferenceConfig { schema: SchemaConfig::PromQL(promql_schema), @@ -143,9 +147,11 @@ mod tests { let data_requests = host_a_series([(1000, 200.0), (2000, 300.0)]); let engine = create_range_engine_two_metrics( "errors_total", + vec!["host"], data_errors, "sum(errors_total) by (host)", "requests_total", + vec!["host"], data_requests, "sum(requests_total) by (host)", ); @@ -174,9 +180,11 @@ mod tests { let data_b = host_a_series([(1000, 20.0), (2000, 25.0)]); let engine = create_range_engine_two_metrics( "metric_a", + vec!["host"], data_a, "sum(metric_a) by (host)", "metric_b", + vec!["host"], data_b, "sum(metric_b) by (host)", ); @@ -199,10 +207,12 @@ mod tests { let data_a = host_a_series([(1000, 5.0), (2000, 6.0)]); let engine = create_range_engine_two_metrics( "metric_a", + vec!["host"], data_a, "sum(metric_a) by (host)", // second metric not used but the helper requires it; empty data. "dummy", + vec!["host"], vec![], "sum(dummy) by (host)", ); @@ -225,9 +235,11 @@ mod tests { let data_a = host_a_series([(1000, 0.9), (2000, 0.75)]); let engine = create_range_engine_two_metrics( "metric_a", + vec!["host"], data_a, "sum(metric_a) by (host)", "dummy", + vec!["host"], vec![], "sum(dummy) by (host)", ); @@ -243,4 +255,35 @@ mod tests { assert!((by_ts[&1000] - 0.1).abs() < 1e-10); assert!((by_ts[&2000] - 0.25).abs() < 1e-10); } + + // Regression test: handle_binary_expr_range_promql's vector-vector join + // used to match purely on positional KeyByLabelValues equality (rhs + // labels discarded), unlike the instant-query combine_vector_vector, + // which rejects a join between arms grouped by different label sets. Two + // arms grouped by disjoint labels ((host) vs (region)) that happen to + // produce the same value could silently join into a wrong-but-plausible + // result across the whole range. + #[tokio::test(flavor = "multi_thread")] + async fn test_range_vector_vector_mismatched_label_sets_return_none() { + let data_a = host_a_series([(1000, 10.0), (2000, 15.0)]); + let data_b = host_a_series([(1000, 10.0), (2000, 15.0)]); + let engine = create_range_engine_two_metrics( + "metric_a", + vec!["host"], + data_a, + "sum(metric_a) by (host)", + "metric_b", + vec!["region"], + data_b, + "sum(metric_b) by (region)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; + let result = engine.handle_range_query_promql(query.to_string(), 1.0, 2.0, 1.0); + assert!( + result.is_none(), + "BUG: arms grouped by different label sets must not join, even when their \ + values coincide, got {result:?}" + ); + } }