From bcf5488144ecf50b7cd01c6e16d482fecf3a2fdb Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Sat, 18 Apr 2026 14:32:21 -0400 Subject: [PATCH] =?UTF-8?q?refactor(sketch-db):=20Phase=203b-2-a=20?= =?UTF-8?q?=E2=80=94=20extract=20forced-agg-id=20context=20builder?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pure refactor. Splits `SimpleEngine::build_query_execution_context_promql` into composable pieces so Phase 3b-2-b's timeline dispatch can build per-segment execution contexts against specific `agg_id`s without re-running the auto-resolver. Zero behavioural change to existing code paths. ## What's extracted * `parse_and_match_promql(query) -> Option<(QueryPatternType, PromQLMatchResult)>` — shared phase-1 (PromQL parse + pattern match). Previously inlined inside `build_query_execution_context_promql`. * `resolve_agg_info_promql(query, match_result, pattern_type) -> Option` — phase-2 auto-resolver (try `QueryConfig` exact match, fall back to capability matching with miss-notify). Previously inlined inside `build_query_execution_context_promql`. * `agg_info_from_forced_id(agg_id) -> Option` — the new building block that powers per-segment dispatch. Looks up the agg type in the current `StreamingConfig` and returns the single-aggregation `AggregationIdInfo` shape (key and value share the same id), matching the fallback branch in `get_aggregation_id_info`. Returns `None` if the agg_id has been removed from the config. ## What's new * `build_query_execution_context_promql_for_agg_id(query, time, forced_agg_id) -> Option` — parallel entry point that does phase-1 + `agg_info_from_forced_id` + phase-3 (`build_promql_execution_context_tail`), skipping the auto-resolver entirely. This is the hook Phase 3b-2-b calls per `TimelineSegment`. ## What's NOT changed * `build_query_execution_context_promql` — now delegates to the extracted helpers but body-by-body identical in effect. Caller-observable output unchanged. * `build_promql_execution_context_tail` — untouched (already agg-info-agnostic). * `execute_context` / `execute_range_query_pipeline` — untouched. * All existing query paths (`handle_query_promql`, `handle_range_query_promql`, etc.) — continue through the auto-resolver entry point with zero change. ## Test plan - [x] 4 new unit tests in `forced_agg_id_tests`: * Auto-resolve still works after refactor (regression gate). * Forced entry point produces a context for a known agg_id. * Unknown agg_id returns `None` without panic. * Forced and auto-resolved produce observably identical `AggregationIdInfo` for the common single-agg case. - [x] 696 lib tests pass (up from 692); existing `dispatch_arithmetic_tests::test_handle_query_promql_single_metric_still_works` and `plan_execution_temporal_tests::test_temporal_sum_over_time_*` all still green. - [x] `cargo clippy --workspace --all-targets --tests -- -D warnings` clean. - [x] `cargo fmt -- --check` clean. ## Next Phase 3b-2-b wires the per-segment dispatch: 1. At query entry, call `SchemaRegistry::timeline_for_metric(metric, start, end)`. 2. If exactly 1 segment, delegate to the existing auto-resolver (zero behaviour change for single-schema queries). 3. If > 1 segment: for each segment, call `build_query_execution_context_promql_for_agg_id` with the segment's `agg_id` and clipped range, execute, collect the scalar, and combine via `combine_statistic`. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/engines/simple_engine.rs | 265 ++++++++++++++++-- 1 file changed, 235 insertions(+), 30 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine.rs b/asap-query-engine/src/engines/simple_engine.rs index b12bd4fc..f2019414 100644 --- a/asap-query-engine/src/engines/simple_engine.rs +++ b/asap-query-engine/src/engines/simple_engine.rs @@ -2817,10 +2817,73 @@ impl SimpleEngine { time: f64, ) -> Option { let query_time = Self::convert_query_time_to_data_time(time); + let (query_pattern_type, match_result) = self.parse_and_match_promql(&query)?; + debug!("Found matching query config for: {}", query); + + let query_context_start_time = Instant::now(); + + // Resolve aggregation: try pre-configured query_configs first, + // fall back to capability matching. + let agg_info = self.resolve_agg_info_promql(&query, &match_result, query_pattern_type)?; + + let result = self.build_promql_execution_context_tail( + &match_result, + query_pattern_type, + query_time, + agg_info, + ); + + let query_context_duration = query_context_start_time.elapsed(); + debug!( + "[LATENCY] Query context build: {:.2}ms", + query_context_duration.as_secs_f64() * 1000.0 + ); + + result + } - // Parse PromQL AST using promql-parser crate + /// Like `build_query_execution_context_promql`, but skips the + /// auto-resolution of the `agg_id` and forces the provided + /// `forced_agg_id` instead. The same parse + pattern-match + /// pipeline runs; only the "which aggregation covers this + /// query" step is replaced. + /// + /// Phase 3b-2-a (refactor) — the caller for this entry point is + /// Phase 3b-2-b's per-segment dispatch: for each + /// `TimelineSegment` returned by + /// [`crate::stores::sketch_db::SchemaRegistry::timeline_for_metric`], + /// the dispatch builds a context targeting that segment's + /// `agg_id`, executes it against the clipped segment range, + /// collects the scalar, and combines across segments via + /// [`crate::engines::timeline_dispatch::combine_statistic`]. + /// + /// Returns `None` when the query can't be parsed / pattern-matched, + /// or when `forced_agg_id` isn't in the current `StreamingConfig`. + pub fn build_query_execution_context_promql_for_agg_id( + &self, + query: String, + time: f64, + forced_agg_id: u64, + ) -> Option { + let query_time = Self::convert_query_time_to_data_time(time); + let (query_pattern_type, match_result) = self.parse_and_match_promql(&query)?; + let agg_info = self.agg_info_from_forced_id(forced_agg_id)?; + self.build_promql_execution_context_tail( + &match_result, + query_pattern_type, + query_time, + agg_info, + ) + } + + /// Shared phase-1 of PromQL context construction: parse the + /// query string, run it against every registered pattern, and + /// return the matched `(QueryPatternType, PromQLMatchResult)`. + /// Centralised so the auto-resolution and forced-agg-id entry + /// points share identical parse semantics. + fn parse_and_match_promql(&self, query: &str) -> Option<(QueryPatternType, PromQLMatchResult)> { let parse_start_time = Instant::now(); - let ast = match promql_parser::parser::parse(&query) { + let ast = match promql_parser::parser::parse(query) { Ok(ast) => { let parse_duration = parse_start_time.elapsed(); debug!( @@ -2856,57 +2919,76 @@ impl SimpleEngine { } } - let (query_pattern_type, match_result) = match found_match { + match found_match { Some((pt, result)) => { let pattern_match_duration = pattern_match_start_time.elapsed(); debug!( "Pattern matching took: {:.2}ms", pattern_match_duration.as_secs_f64() * 1000.0 ); - (pt, result) + Some((pt, result)) } None => { warn!("No matching pattern found for query: {}", query); - return None; + None } - }; - - debug!("Found matching query config for: {}", query); - - let query_context_start_time = Instant::now(); + } + } - // Resolve aggregation: try pre-configured query_configs first, fall back to capability matching. - let agg_info: AggregationIdInfo = if let Some(config) = self.find_query_config(&query) { + /// Resolve which aggregation covers a PromQL query: try the + /// `QueryConfig` exact-string match first, then fall back to + /// capability-based matching (with controller miss-notification + /// if wired). Extracted from `build_query_execution_context_promql` + /// so the Phase 3b-2-b per-segment dispatch can choose NOT to + /// auto-resolve (it has a forced agg_id from the timeline). + fn resolve_agg_info_promql( + &self, + query: &str, + match_result: &PromQLMatchResult, + query_pattern_type: QueryPatternType, + ) -> Option { + if let Some(config) = self.find_query_config(query) { self.get_aggregation_id_info(config) .map_err(|e| { warn!("{}", e); e }) - .ok()? + .ok() } else { warn!( "No query_config entry for PromQL query '{}'. Attempting capability-based matching.", query ); let requirements = - self.build_query_requirements_promql(&match_result, query_pattern_type); - self.find_compatible_aggregation_with_miss_notify(&requirements)? - }; - - let result = self.build_promql_execution_context_tail( - &match_result, - query_pattern_type, - query_time, - agg_info, - ); - - let query_context_duration = query_context_start_time.elapsed(); - debug!( - "[LATENCY] Query context build: {:.2}ms", - query_context_duration.as_secs_f64() * 1000.0 - ); + self.build_query_requirements_promql(match_result, query_pattern_type); + self.find_compatible_aggregation_with_miss_notify(&requirements) + } + } - result + /// Build an `AggregationIdInfo` from a single forced `agg_id`, + /// for the Phase 3b-2-b per-segment dispatch. Uses the same + /// "one agg covers both key and value" shape as the single- + /// aggregation branch in `get_aggregation_id_info` (line + /// ~1881), so downstream dispatch treats this agg identically + /// to a single-aggregation `QueryConfig` match. + /// + /// Returns `None` if the agg_id isn't in the current + /// `StreamingConfig` — either a stale controller posted a + /// backfill for a removed agg, or the timeline contains a + /// retired entry whose config was evicted. Either way the + /// per-segment dispatch will skip this segment as + /// non-answerable. + fn agg_info_from_forced_id(&self, agg_id: u64) -> Option { + let streaming_config = self.streaming_config_snapshot(); + let agg_type = streaming_config + .get_aggregation_config(agg_id) + .map(|c| c.aggregation_type)?; + Some(AggregationIdInfo { + aggregation_id_for_key: agg_id, + aggregation_id_for_value: agg_id, + aggregation_type_for_key: agg_type, + aggregation_type_for_value: agg_type, + }) } /// Merge precomputed outputs (extracts buckets from timestamped data) @@ -5343,3 +5425,126 @@ mod aux_pushdown_tests { ); } } + +// ── Phase 3b-2-a: build_query_execution_context_promql_for_agg_id tests ── + +#[cfg(test)] +mod forced_agg_id_tests { + use super::*; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::tests::test_utilities::engine_factories::create_engine_single_pop; + + /// Sanity: the refactored auto-resolve path still produces a + /// valid context for a simple sum_over_time query — ensures + /// the `parse_and_match_promql` / `resolve_agg_info_promql` + /// extraction hasn't broken existing behaviour. + #[test] + fn auto_resolve_still_works_after_refactor() { + let data = vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )]; + let query = "sum_over_time(http_requests[60s])"; + let engine = create_engine_single_pop( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + ); + let ctx = engine.build_query_execution_context_promql(query.to_string(), 1000.0); + assert!(ctx.is_some(), "auto-resolve should yield a context"); + } + + /// The new forced-agg-id entry point produces a context when + /// the forced `agg_id` matches the one the auto-resolver + /// would have picked. Basic smoke test; §7 timeline dispatch + /// tests in Phase 3b-2-b will exercise it against multiple + /// agg_ids. + #[test] + fn forced_agg_id_produces_context_for_known_id() { + let data = vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )]; + let query = "sum_over_time(http_requests[60s])"; + let engine = create_engine_single_pop( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + ); + let ctx = engine.build_query_execution_context_promql_for_agg_id( + query.to_string(), + 1000.0, + 1, // `create_engine_single_pop` wires agg_id=1 + ); + assert!(ctx.is_some(), "forced valid agg_id should yield a context"); + let ctx = ctx.unwrap(); + assert_eq!(ctx.agg_info.aggregation_id_for_value, 1); + assert_eq!(ctx.agg_info.aggregation_id_for_key, 1); + } + + /// Unknown `agg_id` returns `None` without panicking or + /// polluting the auto-resolver state. Phase 3b-2-b relies on + /// this to gracefully skip timeline segments whose agg_id + /// disappeared from the StreamingConfig mid-query. + #[test] + fn forced_agg_id_returns_none_for_unknown_id() { + let data = vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )]; + let query = "sum_over_time(http_requests[60s])"; + let engine = create_engine_single_pop( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + ); + let ctx = engine.build_query_execution_context_promql_for_agg_id( + query.to_string(), + 1000.0, + 9999, // not in StreamingConfig + ); + assert!(ctx.is_none(), "unknown agg_id → None"); + } + + /// Forced and auto-resolved contexts should be observably + /// equivalent for the common one-agg case (where the + /// auto-resolver would have picked the same id). The + /// invariant that matters for Phase 3b-2-b: dispatching + /// through the forced path against the single covering + /// segment yields the same answer as the existing path. + #[test] + fn forced_and_auto_resolve_produce_same_agg_info_for_single_agg() { + let data = vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )]; + let query = "sum_over_time(http_requests[60s])"; + let engine = create_engine_single_pop( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + ); + let auto = engine + .build_query_execution_context_promql(query.to_string(), 1000.0) + .unwrap(); + let forced = engine + .build_query_execution_context_promql_for_agg_id(query.to_string(), 1000.0, 1) + .unwrap(); + assert_eq!( + auto.agg_info.aggregation_id_for_value, + forced.agg_info.aggregation_id_for_value + ); + assert_eq!( + auto.agg_info.aggregation_type_for_value, + forced.agg_info.aggregation_type_for_value + ); + } +}