diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 65da4653..9d333640 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -1,4 +1,5 @@ pub(crate) mod merge_utils; +pub(crate) mod query_plan; pub mod query_result; pub mod simple_engine; pub(crate) mod sliding_window_composition; diff --git a/asap-query-engine/src/engines/query_plan.rs b/asap-query-engine/src/engines/query_plan.rs new file mode 100644 index 00000000..4499d6ad --- /dev/null +++ b/asap-query-engine/src/engines/query_plan.rs @@ -0,0 +1,350 @@ +//! Request-specific native query DAGs. + +use crate::engines::simple_engine::{RangeQueryExecutionContext, StoreQueryParams}; +use asap_types::enums::WindowType; +use promql_utilities::query_logics::enums::Statistic; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct NodeId(usize); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum StoreReadStrategy { + WindowGrid, + SlidingExactCover, +} + +#[derive(Debug, Clone)] +pub(crate) enum QueryPlanNode { + StoreRead { + query: StoreQueryParams, + strategy: StoreReadStrategy, + }, + ComposeWindows { + input: NodeId, + output_timestamps: Vec, + lookback_ms: u64, + window_size_ms: u64, + bucket_step_ms: u64, + }, + ResolveKeys { + values: NodeId, + keys: Option, + }, + Estimate { + input: NodeId, + statistic: Statistic, + query_kwargs: std::collections::HashMap, + }, + LimitTopK { + input: NodeId, + k: String, + }, + Format { + input: NodeId, + include_metric_name: bool, + }, +} + +#[derive(Debug, Clone)] +pub(crate) struct QueryPlan { + nodes: Vec, + root: NodeId, +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct PlanOptions { + pub limit_topk: bool, + pub format_output: bool, +} + +impl QueryPlan { + pub(crate) fn compile_range( + context: &RangeQueryExecutionContext, + options: PlanOptions, + ) -> Result { + let mut nodes = Vec::new(); + let values_read = Self::push_read( + &mut nodes, + &context.base.store_plan.values_query, + context.window_type, + ); + let values = Self::push_compose( + &mut nodes, + values_read, + &context.output_timestamps, + context.query_range_ms, + context.window_size_ms, + context.tumbling_window_ms, + ); + let keys = context.base.store_plan.keys_query.as_ref().map(|query| { + let read = Self::push_read( + &mut nodes, + query, + context.keys_window_type.unwrap_or(context.window_type), + ); + Self::push_compose( + &mut nodes, + read, + &context.output_timestamps, + context.keys_lookback_ms.unwrap_or(context.query_range_ms), + context + .keys_window_size_ms + .unwrap_or(context.window_size_ms), + context + .keys_tumbling_window_ms + .unwrap_or(context.tumbling_window_ms), + ) + }); + let resolved = Self::push(&mut nodes, QueryPlanNode::ResolveKeys { values, keys }); + let mut root = Self::push( + &mut nodes, + QueryPlanNode::Estimate { + input: resolved, + statistic: context.base.metadata.statistic_to_compute, + query_kwargs: context.base.metadata.query_kwargs.clone(), + }, + ); + if options.limit_topk && context.base.metadata.statistic_to_compute == Statistic::Topk { + let k = context + .base + .metadata + .query_kwargs + .get("k") + .cloned() + .ok_or_else(|| "Topk query is missing required `k` parameter".to_string())?; + k.parse::() + .map_err(|_| "Topk query has an invalid `k` parameter".to_string())?; + root = Self::push(&mut nodes, QueryPlanNode::LimitTopK { input: root, k }); + } + if options.format_output { + root = Self::push( + &mut nodes, + QueryPlanNode::Format { + input: root, + include_metric_name: context.base.metadata.keep_metric_name, + }, + ); + } + Ok(Self { nodes, root }) + } + + fn push(nodes: &mut Vec, node: QueryPlanNode) -> NodeId { + let id = NodeId(nodes.len()); + nodes.push(node); + id + } + + fn push_read( + nodes: &mut Vec, + query: &StoreQueryParams, + window_type: WindowType, + ) -> NodeId { + let strategy = match window_type { + WindowType::Tumbling => StoreReadStrategy::WindowGrid, + WindowType::Sliding => StoreReadStrategy::SlidingExactCover, + }; + Self::push( + nodes, + QueryPlanNode::StoreRead { + query: query.clone(), + strategy, + }, + ) + } + + fn push_compose( + nodes: &mut Vec, + input: NodeId, + output_timestamps: &[u64], + lookback_ms: u64, + window_size_ms: u64, + bucket_step_ms: u64, + ) -> NodeId { + Self::push( + nodes, + QueryPlanNode::ComposeWindows { + input, + output_timestamps: output_timestamps.to_vec(), + lookback_ms, + window_size_ms, + bucket_step_ms, + }, + ) + } + + pub(crate) fn explain(&self) -> String { + let mut lines = Vec::with_capacity(self.nodes.len() + 1); + for (index, node) in self.nodes.iter().enumerate() { + let line = match node { + QueryPlanNode::StoreRead { query, strategy } => format!( + "n{index} StoreRead({strategy:?}, {}#{}, [{}, {}])", + query.metric, query.aggregation_id, query.start_timestamp, query.end_timestamp + ), + QueryPlanNode::ComposeWindows { input, output_timestamps, lookback_ms, window_size_ms, bucket_step_ms } => format!( + "n{index} ComposeWindows(n{}, outputs={:?}, lookback={lookback_ms}ms, window={window_size_ms}ms, step={bucket_step_ms}ms)", + input.0, output_timestamps + ), + QueryPlanNode::ResolveKeys { values, keys } => format!( + "n{index} ResolveKeys(values=n{}, keys={})", + values.0, + keys.map(|id| format!("n{}", id.0)).unwrap_or_else(|| "self".to_string()) + ), + QueryPlanNode::Estimate { input, statistic, query_kwargs } => { + let mut kwargs: Vec<_> = query_kwargs.iter().collect(); + kwargs.sort_unstable_by_key(|(key, _)| *key); + format!("n{index} Estimate(n{}, {statistic}, {kwargs:?})", input.0) + }, + QueryPlanNode::LimitTopK { input, k } => format!("n{index} LimitTopK(n{}, k={k})", input.0), + QueryPlanNode::Format { input, include_metric_name } => format!( + "n{index} Format(n{}, include_metric_name={include_metric_name})", input.0 + ), + }; + lines.push(line); + } + lines.push(format!("root: n{}", self.root.0)); + lines.join("\n") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data_model::AggregationIdInfo; + use crate::engines::simple_engine::{QueryExecutionContext, QueryMetadata, StoreQueryPlan}; + use promql_utilities::data_model::KeyByLabelNames; + use promql_utilities::query_logics::enums::AggregationType; + use std::collections::HashMap; + + fn context() -> RangeQueryExecutionContext { + RangeQueryExecutionContext { + base: QueryExecutionContext { + metric: "requests".into(), + metadata: QueryMetadata { + query_output_labels: KeyByLabelNames::empty(), + statistic_to_compute: Statistic::Sum, + query_kwargs: HashMap::new(), + keep_metric_name: false, + }, + store_plan: StoreQueryPlan { + values_query: StoreQueryParams { + metric: "requests".into(), + aggregation_id: 7, + start_timestamp: 0, + end_timestamp: 1_000, + }, + keys_query: None, + }, + agg_info: AggregationIdInfo { + aggregation_id_for_key: 7, + aggregation_id_for_value: 7, + aggregation_type_for_key: AggregationType::Sum, + aggregation_type_for_value: AggregationType::Sum, + }, + value_window_type: WindowType::Tumbling, + do_merge: false, + spatial_filter: String::new(), + query_time: 1_000, + grouping_labels: KeyByLabelNames::empty(), + aggregated_labels: KeyByLabelNames::empty(), + }, + output_timestamps: vec![1_000], + query_range_ms: 1_000, + buckets_per_step: 1, + lookback_bucket_count: 1, + tumbling_window_ms: 1_000, + window_type: WindowType::Tumbling, + window_size_ms: 1_000, + keys_window_type: None, + keys_window_size_ms: None, + keys_lookback_ms: None, + keys_tumbling_window_ms: None, + } + } + + #[test] + fn separate_key_branch_fans_into_key_resolution() { + let mut context = context(); + context.base.store_plan.keys_query = Some(StoreQueryParams { + metric: "requests".into(), + aggregation_id: 8, + start_timestamp: 0, + end_timestamp: 1_000, + }); + context.keys_window_type = Some(WindowType::Sliding); + + let explanation = QueryPlan::compile_range( + &context, + PlanOptions { + limit_topk: false, + format_output: false, + }, + ) + .unwrap() + .explain(); + + assert!(explanation.contains("n4 ResolveKeys(values=n1, keys=n3)")); + assert!(explanation.contains("n2 StoreRead(SlidingExactCover, requests#8")); + assert!(explanation.ends_with("root: n5")); + } + + #[test] + fn range_plan_keeps_every_output_timestamp() { + let mut context = context(); + context.output_timestamps = vec![1_000, 2_000, 3_000]; + + let explanation = QueryPlan::compile_range( + &context, + PlanOptions { + limit_topk: false, + format_output: false, + }, + ) + .unwrap() + .explain(); + + assert!(explanation.contains("outputs=[1000, 2000, 3000]")); + } + + #[test] + fn topk_formatting_is_the_plan_root() { + let mut context = context(); + context.base.metadata.statistic_to_compute = Statistic::Topk; + context + .base + .metadata + .query_kwargs + .insert("k".to_string(), "3".to_string()); + context.base.metadata.keep_metric_name = true; + + let explanation = QueryPlan::compile_range( + &context, + PlanOptions { + limit_topk: true, + format_output: true, + }, + ) + .unwrap() + .explain(); + + assert!(explanation.contains("LimitTopK(n3, k=3)")); + assert!(explanation.contains("Format(n4, include_metric_name=true)")); + assert!(explanation.ends_with("root: n5")); + } + + #[test] + fn rejects_topk_without_a_limit() { + let mut context = context(); + context.base.metadata.statistic_to_compute = Statistic::Topk; + + let error = QueryPlan::compile_range( + &context, + PlanOptions { + limit_topk: true, + format_output: false, + }, + ) + .expect_err("topk plan without k must fail loudly"); + + assert_eq!(error, "Topk query is missing required `k` parameter"); + } +} diff --git a/asap-query-engine/src/engines/query_plan/README.md b/asap-query-engine/src/engines/query_plan/README.md new file mode 100644 index 00000000..ec0dfe2a --- /dev/null +++ b/asap-query-engine/src/engines/query_plan/README.md @@ -0,0 +1,18 @@ +# Native query plan + +`query_plan` compiles a resolved native query into a request-specific DAG. It is +debug-only for now: the existing executor remains the source of execution. + +Nodes are connected by `n` inputs: + +- `StoreRead` — input: none; fetches one aggregation over its timestamp bounds. +- `ComposeWindows` — input: one `StoreRead`; composes buckets for each output timestamp. +- `ResolveKeys` — inputs: a value composition and optional keys composition; resolves output keys. + Without a keys input, the value accumulator supplies its own keys. +- `Estimate` — input: `ResolveKeys`; computes the requested statistic with its parameters. +- `LimitTopK` — input: `Estimate`; keeps the highest-ranked `k` candidates. +- `Format` — input: `Estimate` or `LimitTopK`; applies protocol-specific output labels. + +`StoreRead` is either a tumbling grid scan or a sliding exact-cover scan. A +range plan has one read branch shared across all output timestamps. Binary +expression DAGs are intentionally deferred to the binary-execution PR. diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index bbb569c1..0daaa6d7 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -6,6 +6,7 @@ use crate::data_model::{ AggregationIdInfo, InferenceConfig, KeyByLabelValues, QueryBounds, QueryConfig, QueryLanguage, StreamingConfig, }; +use crate::engines::query_plan::{PlanOptions, QueryPlan}; use crate::engines::query_result::{InstantVectorElement, QueryResult}; use crate::engines::sliding_window_composition::{plan_exact_cover, SlidingWindowSpec}; // use crate::stores::promsketch_store::{ @@ -1187,7 +1188,7 @@ impl SimpleEngine { ) })?; - let range_results = self.execute_range_query_pipeline( + let range_results = self.execute_observed_range_query_pipeline( &range_context, enable_topk_limiting, enable_topk_formatting, @@ -1901,6 +1902,23 @@ impl SimpleEngine { /// (#581 stage E.3), before insertion into the final result map. /// Formatting (metric-name label prefix) is a separate, smaller pass /// afterward, once per group rather than once per timestep. + fn execute_observed_range_query_pipeline( + &self, + context: &RangeQueryExecutionContext, + enable_topk_limiting: bool, + enable_topk_formatting: bool, + ) -> Result, String> { + let plan = QueryPlan::compile_range( + context, + PlanOptions { + limit_topk: enable_topk_limiting, + format_output: enable_topk_formatting, + }, + )?; + debug!(plan = %plan.explain(), "Compiled native query plan"); + self.execute_range_query_pipeline(context, enable_topk_limiting, enable_topk_formatting) + } + fn execute_range_query_pipeline( &self, context: &RangeQueryExecutionContext, diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index d114d627..bd1d88a1 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -742,7 +742,9 @@ impl SimpleEngine { // Binary arms need Topk limiting, but must remain in the // unformatted intermediate label representation until after the // arithmetic operation. - let results = self.execute_range_query_pipeline(&ctx, true, false).ok()?; + let results = self + .execute_observed_range_query_pipeline(&ctx, true, false) + .ok()?; let combined: Vec = results .into_iter() .map(|mut elem| { @@ -772,10 +774,10 @@ impl SimpleEngine { } // Binary arms need Topk limiting, but not final presentation formatting. let lhs_results = self - .execute_range_query_pipeline(&lhs_ctx, true, false) + .execute_observed_range_query_pipeline(&lhs_ctx, true, false) .ok()?; let rhs_results = self - .execute_range_query_pipeline(&rhs_ctx, true, false) + .execute_observed_range_query_pipeline(&rhs_ctx, true, false) .ok()?; // Build lookup: label_key -> {timestamp -> value} for rhs @@ -1342,7 +1344,7 @@ impl SimpleEngine { // instant's handle_query_promql -- both flags are no-ops unless this // query's statistic is Topk. let results: Vec = self - .execute_range_query_pipeline(&context, true, true) + .execute_observed_range_query_pipeline(&context, true, true) .map_err(|e| { warn!("Range query execution failed: {}", e); e