From ed788a3a1f08feb788ef1764723268304c1e3189 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 16:12:42 -0400 Subject: [PATCH 1/2] refactor(query-engine): unify instant key-resolution into a shared resolver Adds resolve_and_query_group, one function both collect_results_separate_keys and collect_results_same_aggregation now call to resolve a value group's expansion keys and query a statistic for each. Consolidates three lenient skip-and-warn cases into one place instead of three separately-maintained copies: a group with keys data but no value data (#597), a keys accumulator that can't produce a resolvable key set, and a single key's stat query failing. The latter two used to hard-fail the entire instant query; they now skip just the affected group/key, matching range's existing behavior. TDD: two new RED tests pin the desired behavior (instant_query_dual_population_unresolvable_key_set_is_skipped_not_fatal, instant_query_dual_population_key_missing_from_value_accumulator_is_skipped_not_fatal) and now pass. Full suite (577 tests) green. Range's inline KeysSource logic still needs rewiring onto this same resolver -- follow-up commit. Part of #581. --- .../src/engines/simple_engine/mod.rs | 156 ++++++++++++------ .../src/tests/native_binary_instant_tests.rs | 138 +++++++++++++++- 2 files changed, 239 insertions(+), 55 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 6b6ff20..4735255 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -656,6 +656,89 @@ impl SimpleEngine { } /// Collects all results based on whether keys are separate or not + /// Resolves a value group's expansion keys and queries `statistic` for + /// each, returning `(key, value)` pairs. Shared by both + /// `collect_results_*` (instant, called once per group) and + /// `execute_range_query_pipeline` (range, called once per group per + /// output step) -- the single place "how do keys get resolved and + /// queried" is decided, so the two pipelines can't drift apart on it + /// again (#570, #582, #587, #597 were all instances of exactly that + /// drift). See #581. + /// + /// - `value_precompute`: `None` means a dual-population group whose keys + /// accumulator has data but whose value accumulator doesn't -- skipped + /// with a warning, not a hard failure (#597). + /// - `keys_precompute`: `Some` for dual-population groups (a separate + /// keys aggregation exists) -- its `get_keys()` supplies the expansion + /// keys, and `value_precompute`'s own `get_keys()` is never consulted + /// (#587). `get_keys()` returning `None` (e.g. a DeltaSetAggregator + /// invariant violation) skips the group with a warning, not a hard + /// failure. `None` for single-population groups -- `value_precompute`'s + /// own `get_keys()` takes priority if present (e.g. a top-k heap); + /// otherwise exactly one row is emitted using `fallback_key` verbatim + /// (the store-level group key, which may itself be `None` for a fully + /// ungrouped query). + /// - A resolved key whose `query_precompute_for_statistic` call fails + /// (e.g. keys/value skew for a dual-population metric) is skipped with + /// a warning; the rest of the group's keys still return. + fn resolve_and_query_group( + &self, + value_precompute: Option<&dyn AggregateCore>, + keys_precompute: Option<&dyn AggregateCore>, + fallback_key: &Option, + statistic: &Statistic, + query_kwargs: &HashMap, + ) -> Vec<(Option, f64)> { + let Some(value_precompute) = value_precompute else { + warn!( + "Group {:?} has keys data but no value data -- skipping this group instead of \ + failing the whole query (#597)", + fallback_key + ); + return Vec::new(); + }; + + let resolved_keys: Vec> = match keys_precompute { + Some(kp) => match kp.get_keys() { + Some(keys) => keys.into_iter().map(Some).collect(), + None => { + warn!( + "Group {:?}'s keys accumulator produced no resolvable key set -- \ + skipping this group instead of failing the whole query", + fallback_key + ); + return Vec::new(); + } + }, + None => match value_precompute.get_keys() { + Some(keys) => keys.into_iter().map(Some).collect(), + None => vec![fallback_key.clone()], + }, + }; + + resolved_keys + .into_iter() + .filter_map(|key| { + match self.query_precompute_for_statistic( + value_precompute, + statistic, + &key, + query_kwargs, + ) { + Ok(value) => Some((key, value)), + Err(e) => { + warn!( + "Failed to query statistic for key {:?} in group {:?}: {} -- \ + skipping this key instead of failing the whole query", + key, fallback_key, e + ); + None + } + } + }) + .collect() + } + fn collect_all_results( &self, merged_values: &HashMap, Box>, @@ -1253,35 +1336,16 @@ impl SimpleEngine { ) -> Result, f64>, String> { let mut unformatted_results = HashMap::new(); - for (key, precompute) in merged_keys { - let keys_for_this_precompute = precompute - .get_keys() - .ok_or_else(|| "Keys required for separate aggregation".to_string())?; - - // A group with keys data but no matching value data is skipped - // instead of failing the whole query, mirroring the range - // query's #583 behavior (previously `.ok_or_else(...)?` here - // hard-failed everything for one missing group; see #597). - let Some(value_precompute) = merged_values.get(key) else { - warn!( - "Instant query: group {:?} has keys data but no value data -- \ - skipping this group instead of failing the whole query (#597)", - key - ); - continue; - }; - - for key_for_this_precompute in keys_for_this_precompute { - let value = self - .query_precompute_for_statistic( - value_precompute.as_ref(), - statistic, - &Some(key_for_this_precompute.clone()), - query_kwargs, - ) - .map_err(|e| format!("Query failed: {}", e))?; - - unformatted_results.insert(Some(key_for_this_precompute.clone()), value); + for (group_key, keys_precompute) in merged_keys { + let value_precompute = merged_values.get(group_key).map(|b| b.as_ref()); + for (key, value) in self.resolve_and_query_group( + value_precompute, + Some(keys_precompute.as_ref()), + group_key, + statistic, + query_kwargs, + ) { + unformatted_results.insert(key, value); } } @@ -1304,31 +1368,15 @@ impl SimpleEngine { ) -> Result, f64>, String> { let mut unformatted_results = HashMap::new(); - for (key, precompute) in merged_outputs { - if let Some(unwrapped_keys) = precompute.get_keys() { - for key_for_this_precompute in unwrapped_keys { - let value = self - .query_precompute_for_statistic( - precompute.as_ref(), - statistic, - &Some(key_for_this_precompute.clone()), - query_kwargs, - ) - .map_err(|e| format!("Query failed: {}", e))?; - - unformatted_results.insert(Some(key_for_this_precompute.clone()), value); - } - } else { - let value = self - .query_precompute_for_statistic( - precompute.as_ref(), - statistic, - &None, - query_kwargs, - ) - .map_err(|e| format!("Query failed: {}", e))?; - - unformatted_results.insert(key.clone(), value); + for (group_key, value_precompute) in merged_outputs { + for (key, value) in self.resolve_and_query_group( + Some(value_precompute.as_ref()), + None, + group_key, + statistic, + query_kwargs, + ) { + unformatted_results.insert(key, value); } } 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 3746908..1967068 100644 --- a/asap-query-engine/src/tests/native_binary_instant_tests.rs +++ b/asap-query-engine/src/tests/native_binary_instant_tests.rs @@ -12,7 +12,9 @@ mod tests { 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::precompute_operators::{ + CountMinSketchAccumulator, DeltaSetAggregatorAccumulator, MultipleSumAccumulator, + }; use crate::tests::test_utilities::engine_factories::{ create_engine_dual_input, create_engine_multi_timestamp_with_window, create_engine_single_pop, create_engine_three_metrics, create_engine_two_metrics, @@ -338,6 +340,140 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread")] + async fn instant_query_dual_population_unresolvable_key_set_is_skipped_not_fatal() { + // collect_results_separate_keys used to hard-fail the ENTIRE instant + // query if the keys precompute's get_keys() returned None: + // `.ok_or_else(|| "Keys required for separate aggregation")?`. + // region=broken's DeltaSetAggregator has the same key in both + // `added` and `removed` -- the invariant get_keys() checks for -- + // so it resolves to None. Per #581, this now skips just that group + // with a warning instead of failing the whole query. + let cms_normal = CountMinSketchAccumulator::new(2, 3); + let cms_broken = CountMinSketchAccumulator::new(2, 3); + + let mut keys_normal = DeltaSetAggregatorAccumulator::new(); + keys_normal.add_key(KeyByLabelValues { + labels: vec![ + "normal".to_string(), + "host-a".to_string(), + "evt-1".to_string(), + ], + }); + let broken_key = KeyByLabelValues { + labels: vec![ + "broken".to_string(), + "host-z".to_string(), + "evt-1".to_string(), + ], + }; + let mut keys_broken = DeltaSetAggregatorAccumulator::new(); + keys_broken.add_key(broken_key.clone()); + keys_broken.remove_key(broken_key); + + let engine = create_engine_dual_input( + "event_frequency", + AggregationType::CountMinSketch, + AggregationType::DeltaSetAggregator, + vec!["region"], + vec!["host", "event"], + vec![ + ( + Some(vec!["normal".to_string()]), + Box::new(cms_normal) as Box, + ), + ( + Some(vec!["broken".to_string()]), + Box::new(cms_broken) as Box, + ), + ], + vec![ + ( + Some(vec!["normal".to_string()]), + Box::new(keys_normal) as Box, + ), + ( + Some(vec!["broken".to_string()]), + Box::new(keys_broken) as Box, + ), + ], + "count(event_frequency) by (region, host, event)", + ); + + let query = "count(event_frequency) by (region, host, event) + 0"; + let (_, qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect( + "instant query should succeed by skipping the unresolvable region=broken group", + ); + let values = vector_values(qr); + + assert!( + values + .iter() + .any(|(labels, _)| labels.contains(&"normal".to_string())), + "region=normal has a resolvable key set and should be unaffected" + ); + assert!( + !values + .iter() + .any(|(labels, _)| labels.contains(&"broken".to_string())), + "region=broken's key set is unresolvable (added/removed invariant violated) -- \ + must be silently skipped, not appear or fail the query" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn instant_query_dual_population_key_missing_from_value_accumulator_is_skipped_not_fatal() + { + // collect_results_separate_keys used to hard-fail the ENTIRE instant + // query if a single resolved key's query_precompute_for_statistic + // call failed: `.map_err(...)?`. Keys and value data come from + // independently-computed accumulators for a dual-population metric, + // so they CAN skew: a key the DeltaSetAggregator (keys side) knows + // about may have no entry in the MultipleSum (value side) + // accumulator at all. Per #581, that one key is now skipped with a + // warning instead of failing every other key in the same query. + let present_key = KeyByLabelValues { + labels: vec!["host-a".to_string(), "evt-1".to_string()], + }; + let missing_key = KeyByLabelValues { + labels: vec!["host-b".to_string(), "evt-2".to_string()], + }; + + let mut value = MultipleSumAccumulator::new(); + value.add_sum(present_key.clone(), 42.0); + // Deliberately no entry for `missing_key`. + + let mut keys = DeltaSetAggregatorAccumulator::new(); + keys.add_key(present_key.clone()); + keys.add_key(missing_key.clone()); + + let engine = create_engine_dual_input( + "event_frequency", + AggregationType::MultipleSum, + AggregationType::DeltaSetAggregator, + vec![], + vec!["host", "event"], + vec![(None, Box::new(value) as Box)], + vec![(None, Box::new(keys) as Box)], + "sum(event_frequency) by (host, event)", + ); + + let query = "sum(event_frequency) by (host, event) + 0"; + let (_, qr) = engine.handle_query_promql(query.to_string(), QUERY_TIME).expect( + "instant query should succeed by skipping the one key missing from the value accumulator", + ); + let values = sorted(vector_values(qr)); + + assert_eq!( + values, + vec![(vec!["host-a".to_string(), "evt-1".to_string()], 42.0)], + "host-b/evt-2 has keys data but no entry in the value accumulator -- must be \ + silently skipped, not appear or fail the query" + ); + } + #[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 From 82a4d601892c13248753792105e2a9eeb44be360 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 17:10:38 -0400 Subject: [PATCH 2/2] refactor(query-engine): rewire range pipeline onto the shared key resolver execute_range_query_pipeline's per-step loop now calls resolve_and_query_group (added in the previous commit for the instant pipeline) instead of its own separate per-key resolution and query loop. KeysSource::Fixed's payload changes from Vec (always 0 or 1 elements) to Option, matching resolve_and_query_group's fallback_key parameter directly. Behavior-preserving: same early-exit triggers for an empty/unmergeable keys window, same fallback-key semantics, same "drop the fully-unlabeled case" outcome (RangeVectorElement can't represent a None key). One deliberate, already-agreed change: a per-key query failure now logs at warn! via the shared resolver instead of range's previous debug!, matching instant's policy -- the row is still dropped either way, only the log level differs. Full suite (577 tests) green, including every existing range dual-population/ sliding-window/delta-set test -- no behavior regression. B (instant/range key-resolution unification) is now complete: both pipelines call the same resolve_and_query_group. Part of #581. --- .../src/engines/simple_engine/mod.rs | 173 +++++++----------- 1 file changed, 69 insertions(+), 104 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 4735255..a68af71 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1688,7 +1688,7 @@ impl SimpleEngine { // over two raw Option fields in the first place — just applied all // the way through instead of partway. enum KeysSource<'a> { - Fixed(Vec), + Fixed(Option), PerStep { bucket_map: HashMap>, lookback_ms: u64, @@ -1748,12 +1748,7 @@ impl SimpleEngine { // this list. None => all_data .iter() - .map(|(group_key, buckets)| { - ( - buckets, - KeysSource::Fixed(group_key.clone().into_iter().collect()), - ) - }) + .map(|(group_key, buckets)| (buckets, KeysSource::Fixed(group_key.clone()))) .collect(), }; @@ -1775,12 +1770,12 @@ impl SimpleEngine { while current_time <= end_ms { // #583: dual-population groups resolve their expansion keys // from the keys aggregation, per step — not a single - // snapshot reused for every step — and (#587) never from - // the value accumulator's own get_keys(). If nothing - // resolves at this step, skip it before ever touching the - // value merge below. Fixed (single-population) groups defer - // key resolution until after the value merge (#584/#587). - let per_step_keys: Option> = match &keys_source { + // snapshot reused for every step. If nothing resolves at + // this step, skip it before ever touching the value merge + // below (avoids wasted merge work on steps outside the + // key's lifetime). Fixed (single-population) groups have no + // separate keys accumulator to merge here at all. + let keys_precompute: Option> = match &keys_source { KeysSource::PerStep { bucket_map: keys_bucket_map, lookback_ms: keys_lookback_ms, @@ -1794,46 +1789,25 @@ impl SimpleEngine { *keys_tumbling_window_ms, ); - let expansion_keys = if keys_window_buckets.is_empty() { - Vec::new() - } else { - let mut key_merger = create_window_merger(key_accumulator_type); - key_merger.initialize(keys_window_buckets); - match key_merger.get_merged() { - Ok(merged) => match merged.get_keys() { - Some(keys) => keys, - None => { - // e.g. a DeltaSetAggregator "remove" with no - // matching "add" resolved in this window -- - // distinct from (and louder than) the routine, - // expected "no buckets in this window at all" - // case below, since it means a merge DID - // happen but couldn't resolve a key set. - warn!( - "Keys merge at t={} produced an unresolved key \ - set (get_keys() returned None) -- skipping \ - this step for this group", - current_time - ); - Vec::new() - } - }, - Err(e) => { - warn!("Failed to merge keys at t={}: {}", current_time, e); - Vec::new() - } - } - }; - - if expansion_keys.is_empty() { + if keys_window_buckets.is_empty() { debug!( - "No expansion keys resolved at t={} — skipping this step for this group", + "No keys data in window at t={} — skipping this step for this group", current_time ); current_time += step_ms; continue; } - Some(expansion_keys) + + let mut key_merger = create_window_merger(key_accumulator_type); + key_merger.initialize(keys_window_buckets); + match key_merger.get_merged() { + Ok(merged_keys) => Some(merged_keys), + Err(e) => { + warn!("Failed to merge keys at t={}: {}", current_time, e); + current_time += step_ms; + continue; + } + } } KeysSource::Fixed(_) => None, }; @@ -1846,66 +1820,57 @@ impl SimpleEngine { let window_buckets = Self::scan_window(&bucket_map, window_start, current_time, tumbling_window_ms); - if !window_buckets.is_empty() { - // Merge available buckets - let mut merger = create_window_merger(*accumulator_type); - merger.initialize(window_buckets); - - match merger.get_merged() { - Ok(merged) => { - // See the note above KeysSource: dual-population - // (per_step_keys already resolved, non-empty) - // always uses that; single-population lets the - // value accumulator's own get_keys() (read after - // merging this window, since e.g. a top-k heap's - // keys depend on the window's data) take - // priority, falling back to the store-level - // group key otherwise. - let resolved_keys = match &keys_source { - KeysSource::PerStep { .. } => per_step_keys.expect( - "PerStep always sets per_step_keys above, or continues", - ), - KeysSource::Fixed(fallback_keys) => { - merged.get_keys().unwrap_or_else(|| fallback_keys.clone()) - } - }; - // Query statistic and emit a sample at current_time - // for every expanded key sharing this value group. - for key in &resolved_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={} (per_step_keys={:?}): {}", - current_time, per_step_keys, e - ); - } - } - } else { + if window_buckets.is_empty() { // No data at all for this window - skip sample debug!( - "Skipping sample at {} (per_step_keys={:?}) - no data in window [{}, {})", - current_time, per_step_keys, window_start, current_time + "Skipping sample at {} - no data in window [{}, {})", + current_time, window_start, current_time ); + current_time += step_ms; + continue; + } + + let mut merger = create_window_merger(*accumulator_type); + merger.initialize(window_buckets); + + let merged = match merger.get_merged() { + Ok(merged) => merged, + Err(e) => { + debug!("Failed to get merged result at t={}: {}", current_time, e); + current_time += step_ms; + continue; + } + }; + + let fallback_key = match &keys_source { + KeysSource::Fixed(fallback_key) => fallback_key.clone(), + KeysSource::PerStep { .. } => None, + }; + + // See the note above KeysSource: dual-population always + // resolves via keys_precompute; single-population lets the + // value accumulator's own get_keys() (read after merging + // this window, since e.g. a top-k heap's keys depend on the + // window's data) take priority, falling back to + // fallback_key otherwise. Same resolver instant uses + // (resolve_and_query_group) -- see #581. + for (key, value) in self.resolve_and_query_group( + Some(merged.as_ref()), + keys_precompute.as_deref(), + &fallback_key, + &context.base.metadata.statistic_to_compute, + &context.base.metadata.query_kwargs, + ) { + // A fully unlabeled result (fallback_key was None and + // the value accumulator has no self-keys) has no + // RangeVectorElement representation (labels: + // KeyByLabelValues, not Option) -- matches today's + // behavior of producing no sample for this combination. + let Some(key) = key else { continue }; + results + .entry(key.clone()) + .or_insert_with(|| RangeVectorElement::new(key)) + .add_sample(current_time, value); } current_time += step_ms;