From 4c59e9bf9d95108c8c416dcd3778ef8c5a43b2e5 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 23 Aug 2026 15:01:27 -0400 Subject: [PATCH 1/3] test(query-engine): add red tests for range-query keys_query/is_exact_query bugs (#580) --- asap-query-engine/src/tests/mod.rs | 1 + .../src/tests/native_range_query_tests.rs | 370 ++++++++++++++++++ 2 files changed, 371 insertions(+) create mode 100644 asap-query-engine/src/tests/native_range_query_tests.rs diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 6ac8428..8fbd71e 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -6,6 +6,7 @@ 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 native_range_query_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_range_query_tests.rs b/asap-query-engine/src/tests/native_range_query_tests.rs new file mode 100644 index 0000000..055f65d --- /dev/null +++ b/asap-query-engine/src/tests/native_range_query_tests.rs @@ -0,0 +1,370 @@ +//! Range-query pipeline correctness tests (issue #580). +//! +//! `execute_range_query_pipeline` (`simple_engine/mod.rs`), the pipeline +//! behind `handle_range_query_promql`, has two known bugs relative to the +//! instant-query pipeline (`execute_and_merge_store_queries`): +//! +//! 1. It never reads `keys_query`, so dual-population metrics (separate +//! value/key aggregations, values keyed `None`, grouping coming entirely +//! from the keys aggregation's `get_keys()`) silently return an empty +//! result over a range instead of the expanded key set. +//! 2. `finish_range_context` (`promql.rs`) unconditionally forces +//! `is_exact_query = false`, ignoring the aggregation's real `WindowType`, +//! so Sliding-window range queries don't fetch/merge the way the instant +//! path does. +//! +//! These tests are RED against current code: they mirror instant-query +//! precedents that already pass (`native_binary_instant_tests.rs`'s +//! `binary_expr_vector_vector_dual_population` and +//! `binary_expr_sliding_window_end_to_end_merges_correctly`) but drive the +//! range entrypoint instead. + +#[cfg(test)] +mod tests { + use crate::data_model::{ + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + KeyByLabelValues, PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, + SchemaConfig, StreamingConfig, WindowType, + }; + use crate::engines::query_result::{QueryResult, RangeVectorElement}; + use crate::engines::simple_engine::SimpleEngine; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::precompute_operators::{CountMinSketchAccumulator, DeltaSetAggregatorAccumulator}; + use crate::stores::simple_map_store::SimpleMapStore; + use crate::stores::Store; + use crate::tests::test_utilities::engine_factories::{ + create_engine_dual_input, create_engine_multi_timestamp_with_window, + }; + use crate::AggregateCore; + use promql_utilities::data_model::KeyByLabelNames; + use std::collections::HashMap; + use std::sync::Arc; + + const WINDOW_MS: u64 = 1000; + + fn matrix_values(qr: QueryResult) -> Vec { + match qr { + QueryResult::Matrix(m) => m.values, + _ => panic!("Expected matrix (range vector) result"), + } + } + + /// Dual-population counterpart to + /// datafusion::range_query_arithmetic_tests::create_range_engine_two_metrics: + /// one metric, separate value/key aggregations, data spread across + /// multiple 1s tumbling-window buckets so a range query has more than + /// one output step to expand keys for. + #[allow(clippy::too_many_arguments)] + fn create_range_engine_dual_input( + metric: &str, + value_agg_type: AggregationType, + key_agg_type: AggregationType, + grouping_labels: Vec<&str>, + aggregated_labels: Vec<&str>, + value_data: Vec<(u64, Option>, Box)>, + keys_data: Vec<(u64, Option>, Box)>, + promql_query: &str, + ) -> SimpleEngine { + let grouping_label_strings: Vec = + grouping_labels.iter().map(|s| s.to_string()).collect(); + let aggregated_label_strings: Vec = + aggregated_labels.iter().map(|s| s.to_string()).collect(); + let all_labels: Vec = grouping_label_strings + .iter() + .chain(aggregated_label_strings.iter()) + .cloned() + .collect(); + + let mut aggregation_configs = HashMap::new(); + aggregation_configs.insert( + 1u64, + AggregationConfig { + aggregation_id: 1, + aggregation_type: value_agg_type, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::new(grouping_label_strings.clone()), + aggregated_labels: KeyByLabelNames::empty(), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: WINDOW_MS, + slide_interval_ms: WINDOW_MS, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: metric.to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + aggregation_configs.insert( + 2u64, + AggregationConfig { + aggregation_id: 2, + aggregation_type: key_agg_type, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::new(grouping_label_strings), + aggregated_labels: KeyByLabelNames::new(aggregated_label_strings), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: WINDOW_MS, + slide_interval_ms: WINDOW_MS, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: metric.to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ); + + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs, + }); + + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + + for (agg_id, data) in [(1u64, value_data), (2u64, keys_data)] { + for (timestamp, label_values_opt, acc) in data { + let key = label_values_opt.map(|labels| KeyByLabelValues { labels }); + let output = PrecomputedOutput::new(timestamp - WINDOW_MS, timestamp, key, agg_id); + store.insert_precomputed_output(output, acc).unwrap(); + } + } + + let promql_schema = + PromQLSchema::new().add_metric(metric.to_string(), KeyByLabelNames::new(all_labels)); + + let query_config = QueryConfig::new(promql_query.to_string()) + .add_aggregation(AggregationReference::new(1, None)) + .add_aggregation(AggregationReference::new(2, None)); + + let inference_config = InferenceConfig { + schema: SchemaConfig::PromQL(promql_schema), + query_configs: vec![query_config], + cleanup_policy: CleanupPolicy::NoCleanup, + }; + + SimpleEngine::new( + store, + inference_config, + streaming_config, + WINDOW_MS, + QueryLanguage::promql, + ) + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_query_dual_population_returns_key_expansion() { + // Same fixture as native_binary_instant_tests::binary_expr_vector_vector_dual_population, + // but queried as a plain range instead of an instant/binary-expr query. + // Data lands at data-time 1_000_000ms (see create_engine_dual_input). + 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)"; + let result = engine.handle_range_query_promql(query.to_string(), 1000.0, 1000.5, 1.0); + let (_, qr) = result.expect("range query failed"); + assert!( + !matrix_values(qr).is_empty(), + "expected keys_query expansion to produce at least one series over the range, \ + matching the instant-query dual-population path" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_query_sliding_window_merges_both_buckets() { + // Same fixture as + // native_binary_instant_tests::binary_expr_sliding_window_end_to_end_merges_correctly, + // but queried as a plain range instead of an instant/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 query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + 1_000, // window_size_ms, matches the fixed 1000ms bucket width + WindowType::Sliding, + ); + + let result = engine.handle_range_query_promql(query.to_string(), 1000.0, 1000.5, 1.0); + let (_, qr) = result.expect("range query failed"); + let elements = matrix_values(qr); + assert_eq!(elements.len(), 1, "expected one merged series for host-a"); + assert_eq!(elements[0].samples.len(), 1); + assert!( + (elements[0].samples[0].value - 15.0).abs() < 1e-10, + "expected both sliding-window buckets merged into 15.0, got {}", + elements[0].samples[0].value + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_query_sliding_window_merges_three_buckets_same_timestamp() { + // Generalizes range_query_sliding_window_merges_both_buckets beyond + // exactly 2 colliding buckets, mirroring + // native_pipeline_merge_tests::sliding_bucket_count_mismatch_still_returns_merged_result. + 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, + ), + ( + 1_000_000, + 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, + 1_000, + WindowType::Sliding, + ); + + let result = engine.handle_range_query_promql(query.to_string(), 1000.0, 1000.5, 1.0); + let (_, qr) = result.expect("range query failed"); + let elements = matrix_values(qr); + assert_eq!(elements.len(), 1, "expected one merged series for host-a"); + assert_eq!(elements[0].samples.len(), 1); + assert!( + (elements[0].samples[0].value - 18.0).abs() < 1e-10, + "expected all 3 sliding-window buckets merged into 18.0, got {}", + elements[0].samples[0].value + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_query_sliding_window_single_bucket_regression() { + // No-collision counterpart to the merge tests above: a single + // Sliding bucket per output step must still return its value + // unchanged once is_exact_query correctly honors WindowType::Sliding + // for range queries. Mirrors + // native_pipeline_merge_tests::sliding_single_bucket_returns_its_value. + let data = vec![( + 1_000_000, + 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, + 1_000, + WindowType::Sliding, + ); + + let result = engine.handle_range_query_promql(query.to_string(), 1000.0, 1000.5, 1.0); + let (_, qr) = result.expect("range query failed"); + let elements = matrix_values(qr); + assert_eq!(elements.len(), 1); + assert_eq!(elements[0].samples.len(), 1); + assert!( + (elements[0].samples[0].value - 42.0).abs() < 1e-10, + "expected the single sliding-window bucket's value unchanged, got {}", + elements[0].samples[0].value + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn range_query_dual_population_expands_keys_across_multiple_steps() { + // Extends range_query_dual_population_returns_key_expansion across + // two output timestamps instead of one, so a fix that only expands + // keys for the first step (rather than every step in the loop) + // still fails this. + let cms_1 = CountMinSketchAccumulator::new(2, 3); + let cms_2 = CountMinSketchAccumulator::new(2, 3); + let mut keys_1 = DeltaSetAggregatorAccumulator::new(); + keys_1.add_key(KeyByLabelValues { + labels: vec!["host-a".to_string(), "evt-1".to_string()], + }); + let mut keys_2 = DeltaSetAggregatorAccumulator::new(); + keys_2.add_key(KeyByLabelValues { + labels: vec!["host-a".to_string(), "evt-1".to_string()], + }); + + let engine = create_range_engine_dual_input( + "event_frequency", + AggregationType::CountMinSketch, + AggregationType::DeltaSetAggregator, + vec![], + vec!["host", "event"], + vec![ + (1000, None, Box::new(cms_1) as Box), + (2000, None, Box::new(cms_2) as Box), + ], + vec![ + (1000, None, Box::new(keys_1) as Box), + (2000, None, Box::new(keys_2) as Box), + ], + "count(event_frequency) by (host, event)", + ); + + let query = "count(event_frequency) by (host, event)"; + let result = engine.handle_range_query_promql(query.to_string(), 1.0, 2.0, 1.0); + let (_, qr) = result.expect("range query failed"); + let elements = matrix_values(qr); + assert!( + !elements.is_empty(), + "expected keys_query expansion to produce at least one series over the range" + ); + let timestamps: std::collections::HashSet = elements + .iter() + .flat_map(|e| e.samples.iter().map(|s| s.timestamp)) + .collect(); + assert_eq!( + timestamps, + std::collections::HashSet::from([1000, 2000]), + "expected keys expansion at every output step, not just the first, got {:?}", + timestamps + ); + } +} From fc6080eb19b2dae185ff65742589c4d273259551 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 23 Aug 2026 15:37:54 -0400 Subject: [PATCH 2/3] fix(query-engine): range queries expand keys_query and merge same-timestamp buckets (#580) execute_range_query_pipeline never read keys_query, so dual-population metrics returned nothing over a range. It also collapsed same-start- timestamp buckets in a HashMap, dropping all but one when a Sliding aggregation legitimately returns more than one bucket per window (#567/#570). Mirrors the keys fetch/merge/expand pattern already used by execute_and_merge_store_queries/collect_results_separate_keys. --- .../src/engines/simple_engine/mod.rs | 135 ++++++++++-------- .../src/tests/native_range_query_tests.rs | 26 ++-- 2 files changed, 87 insertions(+), 74 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index f4e4c51..d6e11ee 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1456,6 +1456,22 @@ impl SimpleEngine { all_data.values().map(|v| v.len()).sum::() ); + // Dual-population metrics (separate key/value aggregations) need the + // keys side merged once up front, then expanded per value group below + // — mirrors execute_and_merge_store_queries/collect_results_separate_keys + // (#580: this pipeline previously never read keys_query at all). + let merged_keys: Option = + if let Some(keys_params) = &context.base.store_plan.keys_query { + let keys_map = self.execute_store_query(keys_params)?; + Some(self.merge_precomputed_outputs( + &keys_map, + true, + context.base.agg_info.aggregation_type_for_key, + )) + } else { + None + }; + let mut results: HashMap = HashMap::new(); // Determine accumulator type for merger selection @@ -1467,6 +1483,8 @@ impl SimpleEngine { let end_ms = context.range_params.end; let buckets_per_step = context.buckets_per_step; let lookback_bucket_count = context.lookback_bucket_count; + let tumbling_window_ms = context.tumbling_window_ms; + let lookback_ms = (lookback_bucket_count as u64) * tumbling_window_ms; let window_mode = if buckets_per_step <= lookback_bucket_count { "sliding (slide <= size)" @@ -1479,42 +1497,42 @@ impl SimpleEngine { start_ms, end_ms, step_ms, - context.tumbling_window_ms, + tumbling_window_ms, buckets_per_step, lookback_bucket_count, window_mode ); - // Process each key independently - for (key_opt, timestamped_buckets) in &all_data { - let key = match key_opt { - Some(k) => k.clone(), - None => continue, // Skip None keys for now + // Process each value group independently + for (group_key, timestamped_buckets) in &all_data { + // Resolve which output label-keys this value group serves: its + // own key for single-population metrics, or every key the merged + // keys aggregation expands it to for dual-population metrics. + let expansion_keys: Vec = match &merged_keys { + Some(keys_map) => match keys_map.get(group_key).and_then(|kp| kp.get_keys()) { + Some(ks) => ks, + None => continue, + }, + None => match group_key { + Some(k) => vec![k.clone()], + None => continue, // Skip None keys for now + }, }; - // Build lookup: bucket_start_timestamp -> bucket for O(1) access - let bucket_map: HashMap = timestamped_buckets - .iter() - .map(|((start, _), bucket)| (*start, bucket.as_ref())) - .collect(); + // Build lookup: bucket_start_timestamp -> all buckets sharing + // that start. A Sliding aggregation can legitimately return more + // than one bucket per start timestamp (#567/#570) — every one of + // them must be merged, not just the last one collected here. + let mut bucket_map: HashMap> = HashMap::new(); + for ((start, _), bucket) in timestamped_buckets { + bucket_map.entry(*start).or_default().push(bucket.as_ref()); + } debug!( - "Key {:?}: built bucket_map with {} entries, timestamps: {:?}", - key, + "Group {:?}: built bucket_map with {} start-timestamps, expands to {} key(s)", + group_key, bucket_map.len(), - bucket_map.keys().collect::>() - ); - - // Create result element for this key - let mut element = RangeVectorElement::new(key.clone()); - - // Calculate window parameters - let tumbling_window_ms = context.tumbling_window_ms; - let lookback_ms = (lookback_bucket_count as u64) * tumbling_window_ms; - - debug!( - "Key {:?}: range [{}, {}], step={}, lookback_ms={}, tumbling_window_ms={}", - key, start_ms, end_ms, step_ms, lookback_ms, tumbling_window_ms + expansion_keys.len() ); // Iterate by OUTPUT timestamp, not by bucket index @@ -1529,10 +1547,10 @@ impl SimpleEngine { let mut t = window_start; while t < current_time { - if let Some(bucket) = bucket_map.get(&t) { - window_buckets.push((*bucket).clone_boxed_core()); + if let Some(buckets) = bucket_map.get(&t) { + window_buckets.extend(buckets.iter().map(|b| b.clone_boxed_core())); } - // If bucket missing at timestamp t, just skip it (partial data is okay) + // If no bucket at timestamp t, just skip it (partial data is okay) t += tumbling_window_ms; } @@ -1543,56 +1561,47 @@ impl SimpleEngine { match merger.get_merged() { Ok(merged) => { - // Query statistic and emit sample at current_time - match self.query_precompute_for_statistic( - merged.as_ref(), - &context.base.metadata.statistic_to_compute, - &Some(key.clone()), - &context.base.metadata.query_kwargs, - ) { - Ok(value) => { - debug!( - "Key {:?}: emitting sample (t={}, value={})", - key, current_time, value - ); - element.add_sample(current_time, value); - } - Err(e) => { - debug!( - "Failed to query statistic at t={} for key {:?}: {}", - current_time, key, e - ); + // Query statistic and emit a sample at current_time + // for every expanded key sharing this value group. + for key in &expansion_keys { + match self.query_precompute_for_statistic( + merged.as_ref(), + &context.base.metadata.statistic_to_compute, + &Some(key.clone()), + &context.base.metadata.query_kwargs, + ) { + Ok(value) => { + results + .entry(key.clone()) + .or_insert_with(|| RangeVectorElement::new(key.clone())) + .add_sample(current_time, value); + } + Err(e) => { + debug!( + "Failed to query statistic at t={} for key {:?}: {}", + current_time, key, e + ); + } } } } Err(e) => { debug!( - "Failed to get merged result at t={} for key {:?}: {}", - current_time, key, e + "Failed to get merged result at t={} for group {:?}: {}", + current_time, group_key, e ); } } } else { // No data at all for this window - skip sample debug!( - "Key {:?}: skipping sample at {} - no data in window [{}, {})", - key, current_time, window_start, current_time + "Group {:?}: skipping sample at {} - no data in window [{}, {})", + group_key, current_time, window_start, current_time ); } current_time += step_ms; } - - debug!( - "Key {:?}: finished with {} samples", - key, - element.samples.len() - ); - - // Only include keys with samples - if !element.samples.is_empty() { - results.insert(key, element); - } } // Convert to Vec diff --git a/asap-query-engine/src/tests/native_range_query_tests.rs b/asap-query-engine/src/tests/native_range_query_tests.rs index 055f65d..891e3c7 100644 --- a/asap-query-engine/src/tests/native_range_query_tests.rs +++ b/asap-query-engine/src/tests/native_range_query_tests.rs @@ -32,9 +32,7 @@ mod tests { use crate::precompute_operators::{CountMinSketchAccumulator, DeltaSetAggregatorAccumulator}; use crate::stores::simple_map_store::SimpleMapStore; use crate::stores::Store; - use crate::tests::test_utilities::engine_factories::{ - create_engine_dual_input, create_engine_multi_timestamp_with_window, - }; + use crate::tests::test_utilities::engine_factories::create_engine_multi_timestamp_with_window; use crate::AggregateCore; use promql_utilities::data_model::KeyByLabelNames; use std::collections::HashMap; @@ -49,6 +47,9 @@ mod tests { } } + /// One tumbling-window bucket: (bucket end timestamp ms, label values, accumulator). + type TimeSeriesData = Vec<(u64, Option>, Box)>; + /// Dual-population counterpart to /// datafusion::range_query_arithmetic_tests::create_range_engine_two_metrics: /// one metric, separate value/key aggregations, data spread across @@ -61,8 +62,8 @@ mod tests { key_agg_type: AggregationType, grouping_labels: Vec<&str>, aggregated_labels: Vec<&str>, - value_data: Vec<(u64, Option>, Box)>, - keys_data: Vec<(u64, Option>, Box)>, + value_data: TimeSeriesData, + keys_data: TimeSeriesData, promql_query: &str, ) -> SimpleEngine { let grouping_label_strings: Vec = @@ -164,28 +165,31 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn range_query_dual_population_returns_key_expansion() { - // Same fixture as native_binary_instant_tests::binary_expr_vector_vector_dual_population, + // Same dual-population shape as native_binary_instant_tests::binary_expr_vector_vector_dual_population, // but queried as a plain range instead of an instant/binary-expr query. - // Data lands at data-time 1_000_000ms (see create_engine_dual_input). + // Uses create_range_engine_dual_input (not create_engine_dual_input): + // the latter places its bucket at (timestamp, timestamp), a zero-width + // window tuned for the instant-query fetch path, which never lines up + // with the range pipeline's bucket_map (keyed by timestamp - window_size). 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( + let engine = create_range_engine_dual_input( "event_frequency", AggregationType::CountMinSketch, AggregationType::DeltaSetAggregator, vec![], vec!["host", "event"], - vec![(None, Box::new(cms))], - vec![(None, Box::new(keys))], + vec![(1000, None, Box::new(cms) as Box)], + vec![(1000, None, Box::new(keys) as Box)], "count(event_frequency) by (host, event)", ); let query = "count(event_frequency) by (host, event)"; - let result = engine.handle_range_query_promql(query.to_string(), 1000.0, 1000.5, 1.0); + let result = engine.handle_range_query_promql(query.to_string(), 1.0, 2.0, 1.0); let (_, qr) = result.expect("range query failed"); assert!( !matrix_values(qr).is_empty(), From 7b9f084d04602c271f1ab04336e97332b95164f9 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 23 Aug 2026 16:25:37 -0400 Subject: [PATCH 3/3] fix(query-engine): hard-error on unresolvable range keys, dedup keys fetch/merge (#582 review) - execute_range_query_pipeline now iterates the merged keys map (like collect_results_separate_keys) and fails the whole query on an unresolvable key set, instead of silently continue-ing past it. - Extracted fetch_and_merge_keys, shared by the instant and range paths, and pass the real do_merge flag through instead of hardcoding true for range. - merge_precomputed_outputs's spatial (non-merge) branch now warns and merges instead of asserting exactly 1 precompute per key, since do_merge=false no longer guarantees that under a range query's widened fetch (and won't once Sliding needs merging there too). --- .../src/engines/simple_engine/mod.rs | 172 +++++++++++------- 1 file changed, 105 insertions(+), 67 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index d6e11ee..e7f610c 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -594,37 +594,50 @@ impl SimpleEngine { ); // Query and merge keys if needed - let merged_keys = if let Some(keys_params) = &plan.keys_query { - let keys_store_query_start_time = Instant::now(); - let keys_map = self.execute_store_query(keys_params).map_err(|e| { - warn!("Error querying store for keys: {}", e); - e - })?; - debug!( - "[LATENCY] Keys store query (metric: {}, agg: {}): {}ms", - &keys_params.metric, - keys_params.aggregation_id, - keys_store_query_start_time.elapsed().as_millis() - ); - debug!("Keys query returned {} unique keys", keys_map.len()); + let merged_keys = self.fetch_and_merge_keys( + &plan.keys_query, + agg_info.aggregation_type_for_key, + do_merge, + )?; - let keys_merge_start_time = Instant::now(); - let merged = self.merge_precomputed_outputs( - &keys_map, - do_merge, - agg_info.aggregation_type_for_key, - ); - debug!( - "[LATENCY] Keys merge operation: {:.2}ms, resulted in {} merged outputs", - keys_merge_start_time.elapsed().as_secs_f64() * 1000.0, - merged.len() - ); - Some(merged) - } else { - None + Ok((merged_values, merged_keys)) + } + + /// Fetches and merges the keys side of a dual-population query plan, if + /// present. Shared by `execute_and_merge_store_queries` (instant) and + /// `execute_range_query_pipeline` (range) — both build a `StoreQueryPlan` + /// that may carry a separate `keys_query` and need it merged the same way. + fn fetch_and_merge_keys( + &self, + keys_query: &Option, + key_aggregation_type: AggregationType, + do_merge: bool, + ) -> Result, String> { + let Some(keys_params) = keys_query else { + return Ok(None); }; - Ok((merged_values, merged_keys)) + let keys_store_query_start_time = Instant::now(); + let keys_map = self.execute_store_query(keys_params).map_err(|e| { + warn!("Error querying store for keys: {}", e); + e + })?; + debug!( + "[LATENCY] Keys store query (metric: {}, agg: {}): {}ms", + &keys_params.metric, + keys_params.aggregation_id, + keys_store_query_start_time.elapsed().as_millis() + ); + debug!("Keys query returned {} unique keys", keys_map.len()); + + let keys_merge_start_time = Instant::now(); + let merged = self.merge_precomputed_outputs(&keys_map, do_merge, key_aggregation_type); + debug!( + "[LATENCY] Keys merge operation: {:.2}ms, resulted in {} merged outputs", + keys_merge_start_time.elapsed().as_secs_f64() * 1000.0, + merged.len() + ); + Ok(Some(merged)) } /// Collects all results based on whether keys are separate or not @@ -1113,12 +1126,27 @@ impl SimpleEngine { } } } else { - assert_eq!( - precomputes.len(), - 1, - "Spatial queries should have exactly 1 precompute per key" - ); - merged.insert(key.clone(), precomputes[0].clone_boxed_core()); + // Spatial queries (do_merge=false) normally see exactly 1 + // precompute per key. A range query's widened fetch can + // surface more than expected even when do_merge was + // computed false for the base instant range, and future + // Sliding-window support may make >1 legitimate here too + // — warn and merge instead of asserting/panicking. + if precomputes.len() != 1 { + warn!( + "Spatial query expected 1 precompute per key {:?}, found {}. Merging anyway.", + key, + precomputes.len() + ); + } + match self.merge_accumulators(precomputes) { + Ok(merged_accumulator) => { + merged.insert(key.clone(), merged_accumulator); + } + Err(e) => { + warn!("Failed to merge accumulators for key {:?}: {}", key, e); + } + } } } } @@ -1460,17 +1488,11 @@ impl SimpleEngine { // keys side merged once up front, then expanded per value group below // — mirrors execute_and_merge_store_queries/collect_results_separate_keys // (#580: this pipeline previously never read keys_query at all). - let merged_keys: Option = - if let Some(keys_params) = &context.base.store_plan.keys_query { - let keys_map = self.execute_store_query(keys_params)?; - Some(self.merge_precomputed_outputs( - &keys_map, - true, - context.base.agg_info.aggregation_type_for_key, - )) - } else { - None - }; + let merged_keys = self.fetch_and_merge_keys( + &context.base.store_plan.keys_query, + context.base.agg_info.aggregation_type_for_key, + context.base.do_merge, + )?; let mut results: HashMap = HashMap::new(); @@ -1503,22 +1525,39 @@ impl SimpleEngine { window_mode ); - // Process each value group independently - for (group_key, timestamped_buckets) in &all_data { - // Resolve which output label-keys this value group serves: its - // own key for single-population metrics, or every key the merged - // keys aggregation expands it to for dual-population metrics. - let expansion_keys: Vec = match &merged_keys { - Some(keys_map) => match keys_map.get(group_key).and_then(|kp| kp.get_keys()) { - Some(ks) => ks, - None => continue, - }, - None => match group_key { - Some(k) => vec![k.clone()], - None => continue, // Skip None keys for now - }, - }; + // Resolve, for every value group, which output label-keys it serves: + // its own key for single-population metrics, or every key the merged + // keys aggregation expands it to for dual-population metrics. Mirrors + // collect_results_separate_keys exactly, including its error + // semantics — an unresolvable key set fails the whole range query + // (so callers fall back to Prometheus) instead of silently returning + // a partial result. See #582 review. + let groups: Vec<( + &Vec, + Vec, + )> = match &merged_keys { + Some(keys_map) => keys_map + .iter() + .map(|(group_key, keys_precompute)| { + let timestamped_buckets = all_data + .get(group_key) + .ok_or_else(|| format!("No value for key: {:?}", group_key))?; + let expansion_keys = keys_precompute + .get_keys() + .ok_or_else(|| "Keys required for separate aggregation".to_string())?; + Ok((timestamped_buckets, expansion_keys)) + }) + .collect::, String>>()?, + None => all_data + .iter() + .filter_map(|(group_key, buckets)| { + group_key.as_ref().map(|k| (buckets, vec![k.clone()])) + }) + .collect(), + }; + // Process each value group independently + for (timestamped_buckets, expansion_keys) in groups { // Build lookup: bucket_start_timestamp -> all buckets sharing // that start. A Sliding aggregation can legitimately return more // than one bucket per start timestamp (#567/#570) — every one of @@ -1529,10 +1568,9 @@ impl SimpleEngine { } debug!( - "Group {:?}: built bucket_map with {} start-timestamps, expands to {} key(s)", - group_key, + "Group with {} start-timestamps, expands to keys: {:?}", bucket_map.len(), - expansion_keys.len() + expansion_keys ); // Iterate by OUTPUT timestamp, not by bucket index @@ -1587,16 +1625,16 @@ impl SimpleEngine { } Err(e) => { debug!( - "Failed to get merged result at t={} for group {:?}: {}", - current_time, group_key, e + "Failed to get merged result at t={} for keys {:?}: {}", + current_time, expansion_keys, e ); } } } else { // No data at all for this window - skip sample debug!( - "Group {:?}: skipping sample at {} - no data in window [{}, {})", - group_key, current_time, window_start, current_time + "Keys {:?}: skipping sample at {} - no data in window [{}, {})", + expansion_keys, current_time, window_start, current_time ); }