From 690282deb40e364f056db4c8e0aff96e33dc2b38 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 26 Sep 2026 19:26:47 -0400 Subject: [PATCH 01/22] feat(planner): add explicit query execution plans --- .../rs/asap_types/src/inference_config.rs | 44 +++++++++++- .../rs/asap_types/src/query_config.rs | 52 ++++++++++++++ asap-planner-rs/src/generator.rs | 10 +++ asap-planner-rs/src/optimizer/translator.rs | 7 +- asap-planner-rs/tests/integration.rs | 67 +++++++++++++++++++ 5 files changed, 175 insertions(+), 5 deletions(-) diff --git a/asap-common/dependencies/rs/asap_types/src/inference_config.rs b/asap-common/dependencies/rs/asap_types/src/inference_config.rs index 6bd9a530..fd2f0c93 100644 --- a/asap-common/dependencies/rs/asap_types/src/inference_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/inference_config.rs @@ -7,7 +7,7 @@ use std::io::BufReader; use crate::aggregation_reference::AggregationReference; use crate::enums::{CleanupPolicy, QueryLanguage}; use crate::promql_schema::PromQLSchema; -use crate::query_config::QueryConfig; +use crate::query_config::{QueryConfig, QueryTimeAggregation}; use elastic_dsl_utilities::{ElasticIndexSchema, ElasticMappingSchema}; use promql_utilities::data_model::KeyByLabelNames; use sql_utilities::sqlhelper::{SQLSchema, Table}; @@ -250,6 +250,18 @@ impl InferenceConfig { .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing query field"))? .to_string(); + let planned_subquery = query_data + .get("planned_subquery") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing planned_subquery field"))? + .to_string(); + let query_time_aggregations = query_data + .get("query_time_aggregations") + .ok_or_else(|| anyhow::anyhow!("Missing query_time_aggregations field")) + .and_then(|value| { + serde_yaml::from_value::>(value.clone()) + .map_err(anyhow::Error::from) + })?; let aggregations = if let Some(aggregations_data) = query_data.get("aggregations").and_then(|v| v.as_sequence()) @@ -290,7 +302,9 @@ impl InferenceConfig { Vec::new() }; - let config = QueryConfig::new(query).with_aggregations(aggregations); + let config = + QueryConfig::with_plan(query, planned_subquery, query_time_aggregations) + .with_aggregations(aggregations); configs.push(config); } configs @@ -300,3 +314,29 @@ impl InferenceConfig { Ok(query_configs) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_a_query_without_a_planned_subquery() { + let data: Value = serde_yaml::from_str( + r#" +cleanup_policy: + name: no_cleanup +metrics: {} +queries: + - query: "sum(metric)" + query_time_aggregations: [] + aggregations: [] +"#, + ) + .unwrap(); + + let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql) + .expect_err("query plans must name their planned subquery"); + + assert!(error.to_string().contains("planned_subquery")); + } +} diff --git a/asap-common/dependencies/rs/asap_types/src/query_config.rs b/asap-common/dependencies/rs/asap_types/src/query_config.rs index 670443c4..31ae4a60 100644 --- a/asap-common/dependencies/rs/asap_types/src/query_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/query_config.rs @@ -5,13 +5,65 @@ use crate::aggregation_reference::AggregationReference; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QueryConfig { pub query: String, + pub planned_subquery: String, + pub query_time_aggregations: Vec, pub aggregations: Vec, } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum QueryTimeAggregationOperator { + Sum, + Count, + Avg, + Min, + Max, + Quantile, + Topk, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum QueryTimeGroupingMode { + All, + By, + Without, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryTimeGrouping { + pub mode: QueryTimeGroupingMode, + pub labels: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum QueryTimeAggregationParameter { + Integer(u64), + Float(f64), +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryTimeAggregation { + pub operator: QueryTimeAggregationOperator, + pub grouping: QueryTimeGrouping, + pub parameter: Option, +} + impl QueryConfig { pub fn new(query: String) -> Self { + Self::with_plan(query.clone(), query, Vec::new()) + } + + pub fn with_plan( + query: String, + planned_subquery: String, + query_time_aggregations: Vec, + ) -> Self { Self { query, + planned_subquery, + query_time_aggregations, aggregations: Vec::new(), } } diff --git a/asap-planner-rs/src/generator.rs b/asap-planner-rs/src/generator.rs index 3193cd21..d266b7b2 100644 --- a/asap-planner-rs/src/generator.rs +++ b/asap-planner-rs/src/generator.rs @@ -25,8 +25,10 @@ pub(crate) const KEY_METRICS: &str = "metrics"; pub(crate) const KEY_NAME: &str = "name"; pub(crate) const KEY_NUM_AGG_TO_RETAIN: &str = "num_aggregates_to_retain"; pub(crate) const KEY_PARAMETERS: &str = "parameters"; +pub(crate) const KEY_PLANNED_SUBQUERY: &str = "planned_subquery"; pub(crate) const KEY_QUERIES: &str = "queries"; pub(crate) const KEY_QUERY: &str = "query"; +pub(crate) const KEY_QUERY_TIME_AGGREGATIONS: &str = "query_time_aggregations"; pub(crate) const KEY_READ_COUNT_THRESHOLD: &str = "read_count_threshold"; pub(crate) const KEY_SLIDE_INTERVAL_MS: &str = "slideIntervalMs"; pub(crate) const KEY_SPATIAL_FILTER: &str = "spatialFilter"; @@ -170,6 +172,14 @@ pub fn build_queries_yaml( YamlValue::String(KEY_QUERY.to_string()), YamlValue::String(query_str.clone()), ); + q_map.insert( + YamlValue::String(KEY_PLANNED_SUBQUERY.to_string()), + YamlValue::String(query_str.clone()), + ); + q_map.insert( + YamlValue::String(KEY_QUERY_TIME_AGGREGATIONS.to_string()), + YamlValue::Sequence(Vec::new()), + ); YamlValue::Mapping(q_map) }) .collect() diff --git a/asap-planner-rs/src/optimizer/translator.rs b/asap-planner-rs/src/optimizer/translator.rs index 8f5a07a5..aae75e36 100644 --- a/asap-planner-rs/src/optimizer/translator.rs +++ b/asap-planner-rs/src/optimizer/translator.rs @@ -38,9 +38,10 @@ fn build_inference_config(solution: &OptimizerSolution) -> InferenceConfig { let agg_ref = AggregationReference::new(aggregation_id, Some(retain)); for query_string in &assignment.aqe.query_strings { - inference - .query_configs - .push(QueryConfig::new(query_string.clone()).add_aggregation(agg_ref.clone())); + inference.query_configs.push( + QueryConfig::with_plan(query_string.clone(), query_string.clone(), vec![]) + .add_aggregation(agg_ref.clone()), + ); } } diff --git a/asap-planner-rs/tests/integration.rs b/asap-planner-rs/tests/integration.rs index 7c864429..08658785 100644 --- a/asap-planner-rs/tests/integration.rs +++ b/asap-planner-rs/tests/integration.rs @@ -487,6 +487,73 @@ fn topk_produces_count_min_sketch_with_heap() { ); } +#[test] +fn topk_query_emits_an_explicit_empty_query_time_pipeline() { + let query = "topk(3, sum_over_time(http_requests_total[5m]))"; + let controller = Controller::from_yaml_with_schema( + &format!( + r#" +query_groups: + - id: 1 + queries: + - "{query}" + repetition_delay_ms: 60000 +"# + ), + http_requests_schema(), + default_opts(), + ) + .unwrap(); + + let output = controller.generate().unwrap(); + let inference: serde_yaml::Value = + serde_yaml::from_str(&output.to_inference_yaml_string().unwrap()).unwrap(); + let planned_query = &inference["queries"][0]; + + assert_eq!(planned_query["query"].as_str(), Some(query)); + assert_eq!(planned_query["planned_subquery"].as_str(), Some(query)); + assert_eq!( + planned_query["query_time_aggregations"].as_sequence(), + Some(&vec![]), + "a fully planned query has no query-time aggregation pipeline" + ); +} + +#[test] +fn nested_topk_uses_the_inner_sum_as_its_planned_subquery() { + let query = "topk(3, sum by (job) (http_requests_total))"; + let anchor = "sum by (job) (http_requests_total)"; + let controller = Controller::from_yaml_with_schema( + &format!( + r#" +query_groups: + - id: 1 + queries: + - "{query}" + repetition_delay_ms: 60000 +"# + ), + http_requests_schema(), + default_opts(), + ) + .unwrap(); + + let output = controller.generate().unwrap(); + let inference: serde_yaml::Value = + serde_yaml::from_str(&output.to_inference_yaml_string().unwrap()).unwrap(); + let planned_query = &inference["queries"][0]; + + assert_eq!(planned_query["planned_subquery"].as_str(), Some(anchor)); + assert_eq!( + planned_query["query_time_aggregations"][0]["operator"].as_str(), + Some("topk") + ); + assert_eq!( + planned_query["query_time_aggregations"][0]["parameter"].as_u64(), + Some(3) + ); +} + #[test] fn topk_over_sum_over_time_produces_value_weighted_heap() { // https://github.com/ProjectASAP/asap-internal/issues/699 — topk wrapping From 8d0de9de4116391ac40a689ebf4426f3fee3c0c0 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 27 Sep 2026 08:11:31 -0400 Subject: [PATCH 02/22] feat(planner): emit nested aggregation plans --- asap-planner-rs/src/elastic_dsl/generator.rs | 19 ++-- asap-planner-rs/src/generator.rs | 33 +++++-- asap-planner-rs/src/planner/promql.rs | 99 ++++++++++++++++++++ asap-planner-rs/src/promql/generator.rs | 56 +++++++++-- asap-planner-rs/src/sql/generator.rs | 19 ++-- asap-planner-rs/tests/integration.rs | 38 ++++++++ asap-query-engine/src/utils/file_io.rs | 12 +++ 7 files changed, 245 insertions(+), 31 deletions(-) diff --git a/asap-planner-rs/src/elastic_dsl/generator.rs b/asap-planner-rs/src/elastic_dsl/generator.rs index df43b5f1..fcfae9af 100644 --- a/asap-planner-rs/src/elastic_dsl/generator.rs +++ b/asap-planner-rs/src/elastic_dsl/generator.rs @@ -7,7 +7,7 @@ use std::collections::HashMap; use crate::config::input::ElasticDSLControllerConfig; use crate::error::ControllerError; use crate::generator::{ - build_aggregation_entry, build_queries_yaml, GeneratorOutput, KEY_AGGREGATIONS, + build_aggregation_entry, build_queries_yaml, GeneratorOutput, QueryPlanEntry, KEY_AGGREGATIONS, KEY_CLEANUP_POLICY, KEY_NAME, KEY_QUERIES, }; use crate::planner::agg_config::IntermediateAggConfig; @@ -82,8 +82,8 @@ pub fn generate_elastic_plan( // Dedup map: identifying_key -> IntermediateAggConfig let mut dedup_map: IndexMap = IndexMap::new(); - // query_string -> Vec<(key, cleanup_param)> - let mut query_keys_map: IndexMap)>> = IndexMap::new(); + // query_string -> explicit physical query plan + let mut query_plan_map: IndexMap = IndexMap::new(); // index -> schema builder derived from the queries targeting that index let mut index_schema_builders: IndexMap = IndexMap::new(); @@ -127,7 +127,10 @@ pub fn generate_elastic_plan( keys_for_query.push((key.clone(), cleanup_param)); dedup_map.entry(key).or_insert(config_item); } - query_keys_map.insert(query_string.clone(), keys_for_query); + query_plan_map.insert( + query_string.clone(), + QueryPlanEntry::fully_planned(query_string.clone(), keys_for_query), + ); } } @@ -140,7 +143,7 @@ pub fn generate_elastic_plan( let streaming_yaml = build_elastic_streaming_yaml(&dedup_map, &id_map)?; let inference_yaml = build_elastic_inference_yaml( cleanup_policy, - &query_keys_map, + &query_plan_map, &id_map, &index_schema_builders, )?; @@ -150,7 +153,7 @@ pub fn generate_elastic_plan( streaming_yaml, inference_yaml, aggregation_count: dedup_map.len(), - query_count: query_keys_map.len(), + query_count: query_plan_map.len(), }) } @@ -174,7 +177,7 @@ fn build_elastic_streaming_yaml( fn build_elastic_inference_yaml( cleanup_policy: CleanupPolicy, - query_keys_map: &IndexMap)>>, + query_plan_map: &IndexMap, id_map: &HashMap, index_schema_builders: &IndexMap, ) -> Result { @@ -191,7 +194,7 @@ fn build_elastic_inference_yaml( ); root.insert( YamlValue::String(KEY_QUERIES.to_string()), - YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_keys_map, id_map)), + YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_plan_map, id_map)), ); root.insert( YamlValue::String("indices".to_string()), diff --git a/asap-planner-rs/src/generator.rs b/asap-planner-rs/src/generator.rs index d266b7b2..ba045540 100644 --- a/asap-planner-rs/src/generator.rs +++ b/asap-planner-rs/src/generator.rs @@ -4,6 +4,7 @@ use serde_yaml::Value as YamlValue; use std::collections::HashMap; use asap_types::enums::CleanupPolicy; +use asap_types::query_config::QueryTimeAggregation; use promql_utilities::data_model::KeyByLabelNames; use crate::planner::agg_config::IntermediateAggConfig; @@ -40,6 +41,24 @@ pub(crate) const KEY_VALUE_COLUMNS: &str = "value_columns"; pub(crate) const KEY_WINDOW_SIZE_MS: &str = "windowSizeMs"; pub(crate) const KEY_WINDOW_TYPE: &str = "windowType"; +/// The physical query anchor and query-time work for one configured query. +#[derive(Debug, Clone)] +pub struct QueryPlanEntry { + pub aggregation_keys: Vec<(String, Option)>, + pub planned_subquery: String, + pub query_time_aggregations: Vec, +} + +impl QueryPlanEntry { + pub fn fully_planned(query: String, aggregation_keys: Vec<(String, Option)>) -> Self { + Self { + aggregation_keys, + planned_subquery: query, + query_time_aggregations: Vec::new(), + } + } +} + pub fn key_by_labels_to_yaml(labels: &KeyByLabelNames) -> YamlValue { YamlValue::Sequence( labels @@ -127,13 +146,14 @@ pub fn build_aggregation_entry(id: u32, cfg: &IntermediateAggConfig) -> YamlValu pub fn build_queries_yaml( cleanup_policy: CleanupPolicy, - query_keys_map: &IndexMap)>>, + query_plan_map: &IndexMap, id_map: &HashMap, ) -> Vec { - query_keys_map + query_plan_map .iter() - .map(|(query_str, keys)| { - let aggregations: Vec = keys + .map(|(query_str, plan)| { + let aggregations: Vec = plan + .aggregation_keys .iter() .map(|(key, cleanup_param)| { let agg_id = id_map[key]; @@ -174,11 +194,12 @@ pub fn build_queries_yaml( ); q_map.insert( YamlValue::String(KEY_PLANNED_SUBQUERY.to_string()), - YamlValue::String(query_str.clone()), + YamlValue::String(plan.planned_subquery.clone()), ); q_map.insert( YamlValue::String(KEY_QUERY_TIME_AGGREGATIONS.to_string()), - YamlValue::Sequence(Vec::new()), + serde_yaml::to_value(&plan.query_time_aggregations) + .expect("query-time aggregation pipeline should serialize to YAML"), ); YamlValue::Mapping(q_map) }) diff --git a/asap-planner-rs/src/planner/promql.rs b/asap-planner-rs/src/planner/promql.rs index 240fbfcd..ee5b92a3 100644 --- a/asap-planner-rs/src/planner/promql.rs +++ b/asap-planner-rs/src/planner/promql.rs @@ -1,4 +1,8 @@ use asap_types::enums::CleanupPolicy; +use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeAggregationParameter, + QueryTimeGrouping, QueryTimeGroupingMode, +}; use asap_types::query_requirements::build_query_requirements_promql; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; @@ -27,6 +31,74 @@ pub enum BinaryArm { Scalar(f64), } +/// The supported physical subquery and the aggregation stages evaluated after it. +#[derive(Debug, Clone)] +pub struct NestedAggregationPlan { + pub planned_subquery: String, + pub query_time_aggregations: Vec, +} + +fn aggregation_stage( + aggregate: &promql_parser::parser::AggregateExpr, +) -> Option { + let operator = match aggregate.op.to_string().as_str() { + "sum" => QueryTimeAggregationOperator::Sum, + "count" => QueryTimeAggregationOperator::Count, + "avg" => QueryTimeAggregationOperator::Avg, + "min" => QueryTimeAggregationOperator::Min, + "max" => QueryTimeAggregationOperator::Max, + "quantile" => QueryTimeAggregationOperator::Quantile, + "topk" => QueryTimeAggregationOperator::Topk, + _ => return None, + }; + + let grouping = match &aggregate.modifier { + Some(promql_parser::parser::LabelModifier::Include(labels)) if !labels.is_empty() => { + QueryTimeGrouping { + mode: QueryTimeGroupingMode::By, + labels: labels.labels.clone(), + } + } + Some(promql_parser::parser::LabelModifier::Exclude(labels)) if !labels.is_empty() => { + QueryTimeGrouping { + mode: QueryTimeGroupingMode::Without, + labels: labels.labels.clone(), + } + } + _ => QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + }; + + let parameter = match operator { + QueryTimeAggregationOperator::Topk => { + let promql_parser::parser::Expr::NumberLiteral(number) = aggregate.param.as_deref()? + else { + return None; + }; + if number.val < 0.0 || number.val.fract() != 0.0 { + return None; + } + Some(QueryTimeAggregationParameter::Integer(number.val as u64)) + } + QueryTimeAggregationOperator::Quantile => { + let promql_parser::parser::Expr::NumberLiteral(number) = aggregate.param.as_deref()? + else { + return None; + }; + Some(QueryTimeAggregationParameter::Float(number.val)) + } + _ => None, + }; + + Some(QueryTimeAggregation { + operator, + grouping, + parameter, + }) +} + /// Convert an AST expression to a `BinaryArm`. Scalar literals become /// `BinaryArm::Scalar`; everything else is serialized to a query string. /// Outer parentheses are stripped so nested binary arms can be re-parsed @@ -163,6 +235,33 @@ impl SingleQueryProcessor { ) } + /// Find the largest existing planner-supported subtree below outer + /// aggregations and serialize the peeled stages in execution order. + pub fn nested_aggregation_plan(&self) -> Option { + let mut expression = promql_parser::parser::parse(&self.query).ok()?; + let mut outer_to_inner_stages = Vec::new(); + + loop { + let promql_parser::parser::Expr::Aggregate(aggregate) = expression else { + return None; + }; + outer_to_inner_stages.push(aggregation_stage(&aggregate)?); + expression = *aggregate.expr; + + let planned_subquery = format!("{}", expression); + if self + .make_arm_processor(planned_subquery.clone()) + .is_supported() + { + outer_to_inner_stages.reverse(); + return Some(NestedAggregationPlan { + planned_subquery, + query_time_aggregations: outer_to_inner_stages, + }); + } + } + } + /// Check if query should be processed (supported pattern) pub fn is_supported(&self) -> bool { if let Ok(ast) = promql_parser::parser::parse(&self.query) { diff --git a/asap-planner-rs/src/promql/generator.rs b/asap-planner-rs/src/promql/generator.rs index 297ce683..0481b435 100644 --- a/asap-planner-rs/src/promql/generator.rs +++ b/asap-planner-rs/src/promql/generator.rs @@ -9,7 +9,8 @@ use crate::config::input::ControllerConfig; use crate::error::ControllerError; use crate::generator::{ build_aggregation_entry, build_queries_yaml, key_by_labels_to_yaml, GeneratorOutput, - PuntedQuery, KEY_AGGREGATIONS, KEY_CLEANUP_POLICY, KEY_METRICS, KEY_NAME, KEY_QUERIES, + PuntedQuery, QueryPlanEntry, KEY_AGGREGATIONS, KEY_CLEANUP_POLICY, KEY_METRICS, KEY_NAME, + KEY_QUERIES, }; use crate::planner::agg_config::IntermediateAggConfig; use crate::planner::promql::{BinaryArm, SingleQueryProcessor}; @@ -57,8 +58,8 @@ pub fn generate_plan( // Deduplication map: identifying_key -> (agg_config, assigned_id_placeholder) let mut dedup_map: IndexMap = IndexMap::new(); - // query_string -> Vec<(key, cleanup_param)> - let mut query_keys_map: IndexMap)>> = IndexMap::new(); + // query_string -> physical anchor and query-time execution pipeline + let mut query_plan_map: IndexMap = IndexMap::new(); let mut punted_queries: Vec = Vec::new(); let mut windowing_errors: Vec = Vec::new(); @@ -97,7 +98,10 @@ pub fn generate_plan( keys_for_query.push((key.clone(), cleanup_param)); dedup_map.entry(key).or_insert(config); } - query_keys_map.insert(query_string.clone(), keys_for_query); + query_plan_map.insert( + query_string.clone(), + QueryPlanEntry::fully_planned(query_string.clone(), keys_for_query), + ); } Err(ControllerError::UnknownMetric(ref metric)) => { tracing::warn!( @@ -111,6 +115,38 @@ pub fn generate_plan( } Err(e) => return Err(e), } + } else if let Some(nested_plan) = processor.nested_aggregation_plan() { + let anchor_processor = + processor.make_arm_processor(nested_plan.planned_subquery.clone()); + match anchor_processor.get_streaming_aggregation_configs() { + Ok((configs, cleanup_param)) => { + let mut aggregation_keys = Vec::new(); + for config in configs { + let key = config.identifying_key(); + aggregation_keys.push((key.clone(), cleanup_param)); + dedup_map.entry(key).or_insert(config); + } + query_plan_map.insert( + query_string.clone(), + QueryPlanEntry { + aggregation_keys, + planned_subquery: nested_plan.planned_subquery, + query_time_aggregations: nested_plan.query_time_aggregations, + }, + ); + } + Err(ControllerError::UnknownMetric(ref metric)) => { + tracing::warn!( + query = %query_string, + metric = %metric, + "skipping query referencing unknown metric" + ); + } + Err(ControllerError::Windowing(error)) => { + windowing_errors.push(format!("query '{query_string}': {error}")); + } + Err(error) => return Err(error), + } } else { let mut pending_dedup_map = IndexMap::new(); if let Some(arm_entries) = collect_binary_leaf_entries( @@ -124,7 +160,9 @@ pub fn generate_plan( // Binary arithmetic: register each leaf arm in dedup_map and query_keys_map for (arm_query, keys_for_arm) in arm_entries { // Use `entry` so a standalone query that duplicates an arm wins - query_keys_map.entry(arm_query).or_insert(keys_for_arm); + query_plan_map.entry(arm_query.clone()).or_insert_with(|| { + QueryPlanEntry::fully_planned(arm_query, keys_for_arm) + }); } } } @@ -149,14 +187,14 @@ pub fn generate_plan( // Build inference_config YAML let inference_yaml = - build_inference_yaml(cleanup_policy, &query_keys_map, &id_map, &metric_schema)?; + build_inference_yaml(cleanup_policy, &query_plan_map, &id_map, &metric_schema)?; Ok(GeneratorOutput { punted_queries, streaming_yaml, inference_yaml, aggregation_count: dedup_map.len(), - query_count: query_keys_map.len(), + query_count: query_plan_map.len(), }) } @@ -283,7 +321,7 @@ fn build_streaming_yaml( fn build_inference_yaml( cleanup_policy: CleanupPolicy, - query_keys_map: &IndexMap)>>, + query_plan_map: &IndexMap, id_map: &HashMap, metric_schema: &asap_types::PromQLSchema, ) -> Result { @@ -293,7 +331,7 @@ fn build_inference_yaml( YamlValue::String(cleanup_policy.to_string()), ); - let queries = build_queries_yaml(cleanup_policy, query_keys_map, id_map); + let queries = build_queries_yaml(cleanup_policy, query_plan_map, id_map); // Build metrics section let mut metrics_map = serde_yaml::Mapping::new(); diff --git a/asap-planner-rs/src/sql/generator.rs b/asap-planner-rs/src/sql/generator.rs index 17c33aab..020d49ab 100644 --- a/asap-planner-rs/src/sql/generator.rs +++ b/asap-planner-rs/src/sql/generator.rs @@ -7,7 +7,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use crate::config::input::SQLControllerConfig; use crate::error::ControllerError; use crate::generator::{ - build_aggregation_entry, build_queries_yaml, GeneratorOutput, KEY_AGGREGATIONS, + build_aggregation_entry, build_queries_yaml, GeneratorOutput, QueryPlanEntry, KEY_AGGREGATIONS, KEY_CLEANUP_POLICY, KEY_METADATA_COLUMNS, KEY_NAME, KEY_QUERIES, KEY_TABLES, KEY_TIME_COLUMN, KEY_VALUE_COLUMNS, }; @@ -79,8 +79,8 @@ pub fn generate_sql_plan( // Dedup map: identifying_key -> IntermediateAggConfig let mut dedup_map: IndexMap = IndexMap::new(); - // query_string -> Vec<(key, cleanup_param)> - let mut query_keys_map: IndexMap)>> = IndexMap::new(); + // query_string -> explicit physical query plan + let mut query_plan_map: IndexMap = IndexMap::new(); let mut windowing_errors: Vec = Vec::new(); for qg in &config.query_groups { @@ -112,7 +112,10 @@ pub fn generate_sql_plan( keys_for_query.push((key.clone(), cleanup_param)); dedup_map.entry(key).or_insert(config_item); } - query_keys_map.insert(query_string.clone(), keys_for_query); + query_plan_map.insert( + query_string.clone(), + QueryPlanEntry::fully_planned(query_string.clone(), keys_for_query), + ); } } @@ -131,14 +134,14 @@ pub fn generate_sql_plan( let streaming_yaml = build_sql_streaming_yaml(config, &dedup_map, &id_map)?; let inference_yaml = - build_sql_inference_yaml(config, cleanup_policy, &query_keys_map, &id_map)?; + build_sql_inference_yaml(config, cleanup_policy, &query_plan_map, &id_map)?; Ok(GeneratorOutput { punted_queries: Vec::new(), streaming_yaml, inference_yaml, aggregation_count: dedup_map.len(), - query_count: query_keys_map.len(), + query_count: query_plan_map.len(), }) } @@ -205,7 +208,7 @@ fn build_sql_streaming_yaml( fn build_sql_inference_yaml( config: &SQLControllerConfig, cleanup_policy: CleanupPolicy, - query_keys_map: &IndexMap)>>, + query_plan_map: &IndexMap, id_map: &HashMap, ) -> Result { let mut cleanup_map = serde_yaml::Mapping::new(); @@ -221,7 +224,7 @@ fn build_sql_inference_yaml( ); root.insert( YamlValue::String(KEY_QUERIES.to_string()), - YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_keys_map, id_map)), + YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_plan_map, id_map)), ); root.insert( YamlValue::String(KEY_TABLES.to_string()), diff --git a/asap-planner-rs/tests/integration.rs b/asap-planner-rs/tests/integration.rs index 08658785..ca0d8cb6 100644 --- a/asap-planner-rs/tests/integration.rs +++ b/asap-planner-rs/tests/integration.rs @@ -554,6 +554,44 @@ query_groups: ); } +#[test] +fn nested_aggregations_are_emitted_from_inner_to_outer() { + let query = "max by (job) (topk(3, sum by (job) (http_requests_total)))"; + let anchor = "sum by (job) (http_requests_total)"; + let controller = Controller::from_yaml_with_schema( + &format!( + r#" +query_groups: + - id: 1 + queries: + - "{query}" + repetition_delay_ms: 60000 +"# + ), + http_requests_schema(), + default_opts(), + ) + .unwrap(); + + let output = controller.generate().unwrap(); + let inference: serde_yaml::Value = + serde_yaml::from_str(&output.to_inference_yaml_string().unwrap()).unwrap(); + let planned_query = &inference["queries"][0]; + let pipeline = planned_query["query_time_aggregations"] + .as_sequence() + .unwrap(); + + assert_eq!(planned_query["planned_subquery"].as_str(), Some(anchor)); + assert_eq!(pipeline.len(), 2); + assert_eq!(pipeline[0]["operator"].as_str(), Some("topk")); + assert_eq!(pipeline[1]["operator"].as_str(), Some("max")); + assert_eq!(pipeline[1]["grouping"]["mode"].as_str(), Some("by")); + assert_eq!( + pipeline[1]["grouping"]["labels"].as_sequence().unwrap()[0], + "job" + ); +} + #[test] fn topk_over_sum_over_time_produces_value_weighted_heap() { // https://github.com/ProjectASAP/asap-internal/issues/699 — topk wrapping diff --git a/asap-query-engine/src/utils/file_io.rs b/asap-query-engine/src/utils/file_io.rs index 4866088a..dd129d49 100644 --- a/asap-query-engine/src/utils/file_io.rs +++ b/asap-query-engine/src/utils/file_io.rs @@ -87,14 +87,20 @@ queries: - aggregation_id: 1 num_aggregates_to_retain: 6 query: quantile_over_time(0.5, fake_metric_total[1m]) + planned_subquery: quantile_over_time(0.5, fake_metric_total[1m]) + query_time_aggregations: [] - aggregations: - aggregation_id: 1 num_aggregates_to_retain: 6 query: quantile_over_time(0.95, fake_metric_total[1m]) + planned_subquery: quantile_over_time(0.95, fake_metric_total[1m]) + query_time_aggregations: [] - aggregations: - aggregation_id: 1 num_aggregates_to_retain: 6 query: quantile_over_time(0.99, fake_metric_total[1m]) + planned_subquery: quantile_over_time(0.99, fake_metric_total[1m]) + query_time_aggregations: [] "#; let mut inference_temp_file = NamedTempFile::new().unwrap(); @@ -159,14 +165,20 @@ queries: - aggregation_id: 1 num_aggregates_to_retain: 6 query: quantile_over_time(0.5, fake_metric_total[1m]) + planned_subquery: quantile_over_time(0.5, fake_metric_total[1m]) + query_time_aggregations: [] - aggregations: - aggregation_id: 1 num_aggregates_to_retain: 6 query: quantile_over_time(0.95, fake_metric_total[1m]) + planned_subquery: quantile_over_time(0.95, fake_metric_total[1m]) + query_time_aggregations: [] - aggregations: - aggregation_id: 1 num_aggregates_to_retain: 6 query: quantile_over_time(0.99, fake_metric_total[1m]) + planned_subquery: quantile_over_time(0.99, fake_metric_total[1m]) + query_time_aggregations: [] "#; let mut temp_file = NamedTempFile::new().unwrap(); From c42ef986ee851f9d1628e4d33899012b1b614e7b Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 27 Sep 2026 08:30:00 -0400 Subject: [PATCH 03/22] feat(query-engine): validate query-time aggregation configs --- .../rs/asap_types/src/inference_config.rs | 51 ++++++++++++++++ .../rs/asap_types/src/query_config.rs | 60 +++++++++++++++++++ 2 files changed, 111 insertions(+) diff --git a/asap-common/dependencies/rs/asap_types/src/inference_config.rs b/asap-common/dependencies/rs/asap_types/src/inference_config.rs index fd2f0c93..cb674149 100644 --- a/asap-common/dependencies/rs/asap_types/src/inference_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/inference_config.rs @@ -305,6 +305,9 @@ impl InferenceConfig { let config = QueryConfig::with_plan(query, planned_subquery, query_time_aggregations) .with_aggregations(aggregations); + config + .validate_execution_plan() + .map_err(|error| anyhow::anyhow!("Invalid query execution plan: {error}"))?; configs.push(config); } configs @@ -339,4 +342,52 @@ queries: assert!(error.to_string().contains("planned_subquery")); } + + #[test] + fn rejects_a_query_without_a_query_time_pipeline() { + let data: Value = serde_yaml::from_str( + r#" +cleanup_policy: + name: no_cleanup +metrics: {} +queries: + - query: "sum(metric)" + planned_subquery: "sum(metric)" + aggregations: [] +"#, + ) + .unwrap(); + + let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql) + .expect_err("query plans must declare their query-time pipeline"); + + assert!(error.to_string().contains("query_time_aggregations")); + } + + #[test] + fn rejects_an_invalid_query_time_aggregation() { + let data: Value = serde_yaml::from_str( + r#" +cleanup_policy: + name: no_cleanup +metrics: {} +queries: + - query: "quantile(1.5, sum(metric))" + planned_subquery: "sum(metric)" + query_time_aggregations: + - operator: quantile + grouping: + mode: all + labels: [] + parameter: 1.5 + aggregations: [] +"#, + ) + .unwrap(); + + let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql) + .expect_err("invalid pipeline parameters must fail during config loading"); + + assert!(error.to_string().contains("quantile")); + } } diff --git a/asap-common/dependencies/rs/asap_types/src/query_config.rs b/asap-common/dependencies/rs/asap_types/src/query_config.rs index 31ae4a60..eef9140f 100644 --- a/asap-common/dependencies/rs/asap_types/src/query_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/query_config.rs @@ -50,6 +50,56 @@ pub struct QueryTimeAggregation { pub parameter: Option, } +impl QueryTimeAggregation { + pub fn validate(&self) -> Result<(), String> { + match self.grouping.mode { + QueryTimeGroupingMode::All if !self.grouping.labels.is_empty() => { + return Err("all grouping cannot name labels".to_string()); + } + QueryTimeGroupingMode::By | QueryTimeGroupingMode::Without + if self.grouping.labels.is_empty() => + { + return Err("by and without grouping must name at least one label".to_string()); + } + _ => {} + } + + if self.grouping.labels.iter().any(|label| label.is_empty()) { + return Err("grouping labels cannot be empty".to_string()); + } + let unique_label_count = self + .grouping + .labels + .iter() + .collect::>() + .len(); + if unique_label_count != self.grouping.labels.len() { + return Err("grouping labels must be unique".to_string()); + } + + match (&self.operator, &self.parameter) { + ( + QueryTimeAggregationOperator::Topk, + Some(QueryTimeAggregationParameter::Integer(_)), + ) => {} + (QueryTimeAggregationOperator::Topk, _) => { + return Err("topk requires an integer parameter".to_string()); + } + ( + QueryTimeAggregationOperator::Quantile, + Some(QueryTimeAggregationParameter::Float(phi)), + ) if phi.is_finite() && (0.0..=1.0).contains(phi) => {} + (QueryTimeAggregationOperator::Quantile, _) => { + return Err("quantile requires a finite parameter between 0 and 1".to_string()); + } + (_, None) => {} + _ => return Err("this aggregation does not accept a parameter".to_string()), + } + + Ok(()) + } +} + impl QueryConfig { pub fn new(query: String) -> Self { Self::with_plan(query.clone(), query, Vec::new()) @@ -68,6 +118,16 @@ impl QueryConfig { } } + pub fn validate_execution_plan(&self) -> Result<(), String> { + if self.planned_subquery.trim().is_empty() { + return Err("planned_subquery cannot be empty".to_string()); + } + for aggregation in &self.query_time_aggregations { + aggregation.validate()?; + } + Ok(()) + } + pub fn add_aggregation(mut self, aggregation: AggregationReference) -> Self { self.aggregations.push(aggregation); self From 827414cf6d3df575a611696cd38fd7bb43c279b3 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 27 Sep 2026 13:01:10 -0400 Subject: [PATCH 04/22] test(planner): cover nested aggregation stages --- asap-planner-rs/tests/integration.rs | 75 ++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/asap-planner-rs/tests/integration.rs b/asap-planner-rs/tests/integration.rs index ca0d8cb6..1c8e3aea 100644 --- a/asap-planner-rs/tests/integration.rs +++ b/asap-planner-rs/tests/integration.rs @@ -592,6 +592,81 @@ query_groups: ); } +#[test] +fn nested_aggregation_operators_emit_query_time_stages() { + let cases = [ + ("sum by (job) (sum by (job) (http_requests_total))", "sum"), + ( + "count by (job) (sum by (job) (http_requests_total))", + "count", + ), + ("avg by (job) (sum by (job) (http_requests_total))", "avg"), + ("min by (job) (sum by (job) (http_requests_total))", "min"), + ("max by (job) (sum by (job) (http_requests_total))", "max"), + ( + "quantile by (job) (0.5, sum by (job) (http_requests_total))", + "quantile", + ), + ( + "topk by (job) (3, sum by (job) (http_requests_total))", + "topk", + ), + ]; + + for (query, operator) in cases { + let controller = Controller::from_yaml_with_schema( + &format!( + r#" +query_groups: + - id: 1 + queries: + - "{query}" + repetition_delay_ms: 60000 +"# + ), + http_requests_schema(), + default_opts(), + ) + .unwrap(); + + let output = controller.generate().unwrap(); + let inference: serde_yaml::Value = + serde_yaml::from_str(&output.to_inference_yaml_string().unwrap()).unwrap(); + let stage = &inference["queries"][0]["query_time_aggregations"][0]; + + assert_eq!(stage["operator"].as_str(), Some(operator)); + assert_eq!(stage["grouping"]["mode"].as_str(), Some("by")); + assert_eq!(stage["grouping"]["labels"].as_sequence().unwrap()[0], "job"); + } +} + +#[test] +fn nested_aggregation_without_grouping_is_preserved() { + let query = "max without (instance) (sum by (job) (http_requests_total))"; + let controller = Controller::from_yaml_with_schema( + &format!( + r#" +query_groups: + - id: 1 + queries: + - "{query}" + repetition_delay_ms: 60000 +"# + ), + http_requests_schema(), + default_opts(), + ) + .unwrap(); + + let output = controller.generate().unwrap(); + let inference: serde_yaml::Value = + serde_yaml::from_str(&output.to_inference_yaml_string().unwrap()).unwrap(); + let grouping = &inference["queries"][0]["query_time_aggregations"][0]["grouping"]; + + assert_eq!(grouping["mode"].as_str(), Some("without")); + assert_eq!(grouping["labels"].as_sequence().unwrap()[0], "instance"); +} + #[test] fn topk_over_sum_over_time_produces_value_weighted_heap() { // https://github.com/ProjectASAP/asap-internal/issues/699 — topk wrapping From ad9182ff68a48019acadffebb2b59c7ea78482c6 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 27 Sep 2026 21:13:49 -0400 Subject: [PATCH 05/22] feat(query-engine): execute query-time aggregation pipelines --- asap-query-engine/src/engines/mod.rs | 1 + .../src/engines/query_time_aggregation.rs | 362 ++++++++++++++++++ .../src/engines/simple_engine/promql.rs | 90 ++++- .../src/tests/structural_matching_tests.rs | 67 +++- 4 files changed, 516 insertions(+), 4 deletions(-) create mode 100644 asap-query-engine/src/engines/query_time_aggregation.rs diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 27898c14..5d1c8033 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod merge_utils; pub(crate) mod query_plan; pub mod query_result; +pub(crate) mod query_time_aggregation; pub mod simple_engine; pub(crate) mod sliding_window_composition; pub mod window_merger; diff --git a/asap-query-engine/src/engines/query_time_aggregation.rs b/asap-query-engine/src/engines/query_time_aggregation.rs new file mode 100644 index 00000000..59ab3ebd --- /dev/null +++ b/asap-query-engine/src/engines/query_time_aggregation.rs @@ -0,0 +1,362 @@ +use crate::data_model::KeyByLabelValues; +use crate::engines::query_result::{InstantVectorElement, RangeVectorElement}; +use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeAggregationParameter, + QueryTimeGroupingMode, +}; +use promql_utilities::data_model::KeyByLabelNames; +use std::collections::BTreeMap; + +const METRIC_NAME_LABEL: &str = "__name__"; + +fn output_labels( + input_labels: &KeyByLabelNames, + aggregation: &QueryTimeAggregation, +) -> Result { + if matches!(aggregation.operator, QueryTimeAggregationOperator::Topk) { + return Ok(input_labels.clone()); + } + + let labels = match aggregation.grouping.mode { + QueryTimeGroupingMode::All => Vec::new(), + QueryTimeGroupingMode::By => aggregation.grouping.labels.clone(), + QueryTimeGroupingMode::Without => input_labels + .labels + .iter() + .filter(|label| { + label.as_str() != METRIC_NAME_LABEL && !aggregation.grouping.labels.contains(label) + }) + .cloned() + .collect(), + }; + for label in &labels { + if !input_labels.labels.contains(label) { + return Err(format!( + "query-time aggregation references unknown label '{label}'" + )); + } + } + Ok(KeyByLabelNames::new(labels)) +} + +fn label_indices( + input_labels: &KeyByLabelNames, + output_labels: &KeyByLabelNames, +) -> Result, String> { + output_labels + .labels + .iter() + .map(|label| { + input_labels + .labels + .iter() + .position(|candidate| candidate == label) + .ok_or_else(|| format!("query-time aggregation references unknown label '{label}'")) + }) + .collect() +} + +fn select_labels(labels: &KeyByLabelValues, indices: &[usize]) -> Result { + indices + .iter() + .map(|index| { + labels + .get(*index) + .cloned() + .ok_or_else(|| "result labels do not match the configured label schema".to_string()) + }) + .collect::, _>>() + .map(KeyByLabelValues::new_with_labels) +} + +fn parameter_integer(aggregation: &QueryTimeAggregation) -> Result { + match aggregation.parameter { + Some(QueryTimeAggregationParameter::Integer(value)) => usize::try_from(value) + .map_err(|_| "topk parameter exceeds this platform's usize range".to_string()), + _ => Err("topk requires an integer parameter".to_string()), + } +} + +fn parameter_float(aggregation: &QueryTimeAggregation) -> Result { + match aggregation.parameter { + Some(QueryTimeAggregationParameter::Float(value)) if value.is_finite() => Ok(value), + _ => Err("quantile requires a finite parameter".to_string()), + } +} + +fn aggregate_values(aggregation: &QueryTimeAggregation, values: &[f64]) -> Result { + if values.iter().any(|value| !value.is_finite()) { + return Err("query-time aggregations cannot process non-finite values".to_string()); + } + + let value = match aggregation.operator { + QueryTimeAggregationOperator::Sum => values.iter().sum(), + QueryTimeAggregationOperator::Count => values.len() as f64, + QueryTimeAggregationOperator::Avg => values.iter().sum::() / values.len() as f64, + QueryTimeAggregationOperator::Min => values + .iter() + .copied() + .min_by(f64::total_cmp) + .ok_or_else(|| "cannot aggregate an empty vector".to_string())?, + QueryTimeAggregationOperator::Max => values + .iter() + .copied() + .max_by(f64::total_cmp) + .ok_or_else(|| "cannot aggregate an empty vector".to_string())?, + QueryTimeAggregationOperator::Quantile => { + let phi = parameter_float(aggregation)?; + if !(0.0..=1.0).contains(&phi) { + return Err("quantile parameter must be between 0 and 1".to_string()); + } + let mut sorted = values.to_vec(); + sorted.sort_by(f64::total_cmp); + let rank = phi * (sorted.len() - 1) as f64; + let lower = rank.floor() as usize; + let upper = rank.ceil() as usize; + sorted[lower] + (sorted[upper] - sorted[lower]) * (rank - lower as f64) + } + QueryTimeAggregationOperator::Topk => { + return Err("topk must be evaluated as a ranking stage".to_string()); + } + }; + if !value.is_finite() { + return Err("query-time aggregation produced a non-finite value".to_string()); + } + Ok(value) +} + +fn apply_stage( + input_labels: &KeyByLabelNames, + input: Vec, + aggregation: &QueryTimeAggregation, +) -> Result<(KeyByLabelNames, Vec), String> { + let output_labels = output_labels(input_labels, aggregation)?; + let grouping_labels = match aggregation.operator { + QueryTimeAggregationOperator::Topk => match aggregation.grouping.mode { + QueryTimeGroupingMode::All => KeyByLabelNames::empty(), + QueryTimeGroupingMode::By => KeyByLabelNames::new(aggregation.grouping.labels.clone()), + QueryTimeGroupingMode::Without => KeyByLabelNames::new( + input_labels + .labels + .iter() + .filter(|label| !aggregation.grouping.labels.contains(label)) + .cloned() + .collect(), + ), + }, + _ => output_labels.clone(), + }; + let indices = label_indices(input_labels, &grouping_labels)?; + let mut groups: BTreeMap, Vec> = BTreeMap::new(); + for element in input { + if !element.value.is_finite() { + return Err("query-time aggregations cannot process non-finite values".to_string()); + } + let key = select_labels(&element.labels, &indices)?; + groups.entry(key.labels).or_default().push(element); + } + + if matches!(aggregation.operator, QueryTimeAggregationOperator::Topk) { + let k = parameter_integer(aggregation)?; + let mut results = Vec::new(); + for mut group in groups.into_values() { + group.sort_by(|left, right| { + right + .value + .total_cmp(&left.value) + .then_with(|| left.labels.labels.cmp(&right.labels.labels)) + }); + group.truncate(k); + results.extend(group); + } + results.sort_by(|left, right| left.labels.labels.cmp(&right.labels.labels)); + return Ok((output_labels, results)); + } + + let mut results = Vec::new(); + for (key, group) in groups { + let value = aggregate_values( + aggregation, + &group + .iter() + .map(|element| element.value) + .collect::>(), + )?; + results.push(InstantVectorElement::new( + KeyByLabelValues::new_with_labels(key), + value, + )); + } + Ok((output_labels, results)) +} + +pub(crate) fn apply_instant_pipeline( + mut labels: KeyByLabelNames, + mut results: Vec, + pipeline: &[QueryTimeAggregation], +) -> Result<(KeyByLabelNames, Vec), String> { + for aggregation in pipeline { + (labels, results) = apply_stage(&labels, results, aggregation)?; + } + Ok((labels, results)) +} + +pub(crate) fn apply_range_pipeline( + labels: KeyByLabelNames, + results: Vec, + pipeline: &[QueryTimeAggregation], +) -> Result<(KeyByLabelNames, Vec), String> { + let mut per_timestamp: BTreeMap> = BTreeMap::new(); + for result in results { + for sample in result.samples { + per_timestamp + .entry(sample.timestamp) + .or_default() + .push(InstantVectorElement::new( + result.labels.clone(), + sample.value, + )); + } + } + + let input_labels = labels.clone(); + let mut output_labels = labels; + let mut output: BTreeMap, RangeVectorElement> = BTreeMap::new(); + for (timestamp, input) in per_timestamp { + let (stage_labels, stage_results) = + apply_instant_pipeline(input_labels.clone(), input, pipeline)?; + output_labels = stage_labels; + for result in stage_results { + output + .entry(result.labels.labels.clone()) + .or_insert_with(|| RangeVectorElement::new(result.labels.clone())) + .add_sample(timestamp, result.value); + } + } + Ok((output_labels, output.into_values().collect())) +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_types::query_config::{QueryTimeGrouping, QueryTimeGroupingMode}; + + fn stage(operator: QueryTimeAggregationOperator) -> QueryTimeAggregation { + QueryTimeAggregation { + operator, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::By, + labels: vec!["job".to_string()], + }, + parameter: None, + } + } + + fn input() -> (KeyByLabelNames, Vec) { + ( + KeyByLabelNames::new(vec!["instance".to_string(), "job".to_string()]), + vec![ + InstantVectorElement::new( + KeyByLabelValues::new_with_labels(vec!["a".into(), "api".into()]), + 1.0, + ), + InstantVectorElement::new( + KeyByLabelValues::new_with_labels(vec!["b".into(), "api".into()]), + 3.0, + ), + InstantVectorElement::new( + KeyByLabelValues::new_with_labels(vec!["c".into(), "worker".into()]), + 5.0, + ), + ], + ) + } + + #[test] + fn grouped_aggregations_transform_each_partition() { + for (operator, expected_api, expected_worker) in [ + (QueryTimeAggregationOperator::Sum, 4.0, 5.0), + (QueryTimeAggregationOperator::Count, 2.0, 1.0), + (QueryTimeAggregationOperator::Avg, 2.0, 5.0), + (QueryTimeAggregationOperator::Min, 1.0, 5.0), + (QueryTimeAggregationOperator::Max, 3.0, 5.0), + ] { + let (labels, results) = input(); + let (output_labels, output) = + apply_instant_pipeline(labels, results, &[stage(operator)]).unwrap(); + assert_eq!(output_labels.labels, vec!["job"]); + assert_eq!(output.len(), 2); + assert_eq!(output[0].labels.labels, vec!["api"]); + assert_eq!(output[0].value, expected_api); + assert_eq!(output[1].labels.labels, vec!["worker"]); + assert_eq!(output[1].value, expected_worker); + } + } + + #[test] + fn quantile_interpolates_sorted_values() { + let (labels, results) = input(); + let aggregation = QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Quantile, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Float(0.75)), + }; + + let (output_labels, output) = + apply_instant_pipeline(labels, results, &[aggregation]).unwrap(); + + assert!(output_labels.labels.is_empty()); + assert_eq!(output[0].value, 4.0); + } + + #[test] + fn topk_ranks_within_each_group_and_preserves_input_labels() { + let (labels, results) = input(); + let mut aggregation = stage(QueryTimeAggregationOperator::Topk); + aggregation.parameter = Some(QueryTimeAggregationParameter::Integer(1)); + + let (output_labels, output) = + apply_instant_pipeline(labels, results, &[aggregation]).unwrap(); + + assert_eq!(output_labels.labels, vec!["instance", "job"]); + assert_eq!(output.len(), 2); + assert_eq!(output[0].labels.labels, vec!["b", "api"]); + assert_eq!(output[0].value, 3.0); + assert_eq!(output[1].labels.labels, vec!["c", "worker"]); + assert_eq!(output[1].value, 5.0); + } + + #[test] + fn range_pipeline_applies_each_stage_at_each_timestamp() { + let (labels, instant) = input(); + let mut range = BTreeMap::, RangeVectorElement>::new(); + for element in instant { + let mut range_element = RangeVectorElement::new(element.labels.clone()); + range_element.add_sample(1_000, element.value); + range_element.add_sample(2_000, element.value * 2.0); + range.insert(element.labels.labels, range_element); + } + + let (output_labels, output) = apply_range_pipeline( + labels, + range.into_values().collect(), + &[stage(QueryTimeAggregationOperator::Sum)], + ) + .unwrap(); + + assert_eq!(output_labels.labels, vec!["job"]); + assert_eq!(output.len(), 2); + assert_eq!(output[0].labels.labels, vec!["api"]); + assert_eq!( + output[0] + .samples + .iter() + .map(|sample| sample.value) + .collect::>(), + vec![4.0, 8.0] + ); + } +} diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 65632b90..604fb020 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -10,6 +10,7 @@ use super::{ }; use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig}; use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVectorElement}; +use crate::engines::query_time_aggregation::{apply_instant_pipeline, apply_range_pipeline}; use asap_types::query_requirements::build_query_requirements_promql; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; @@ -278,11 +279,12 @@ impl SimpleEngine { &self, arm_ast: &promql_parser::parser::Expr, query_config: &QueryConfig, + planned_subquery: &str, time: f64, ) -> Option { let query_time = Self::convert_query_time_to_data_time(time); - let match_result = self.find_matching_controller_pattern(arm_ast, &query_config.query)?; + let match_result = self.find_matching_controller_pattern(arm_ast, planned_subquery)?; let agg_info = self .get_aggregation_id_info(query_config) @@ -293,7 +295,7 @@ impl SimpleEngine { .ok()?; self.build_promql_execution_context_tail( - &query_config.query, + planned_subquery, &match_result, query_time, agg_info, @@ -435,7 +437,12 @@ impl SimpleEngine { Expr::Paren(paren) => self.resolve_arm_leaf_context(&paren.expr, time), other => { let config = self.find_query_config_promql_structural(other)?; - let ctx = self.build_query_execution_context_from_ast(other, &config, time)?; + let ctx = self.build_query_execution_context_from_ast( + other, + &config, + &config.planned_subquery, + time, + )?; let label_names = binary_matching_label_names(ctx.metadata.query_output_labels.labels.clone()); Some((ctx, label_names)) @@ -1111,6 +1118,45 @@ impl SimpleEngine { return result; } + if let Some(config) = self.find_query_config(&query) { + if !config.query_time_aggregations.is_empty() { + let anchor_ast = match promql_parser::parser::parse(&config.planned_subquery) { + Ok(ast) => ast, + Err(error) => { + warn!( + query = %query, + planned_subquery = %config.planned_subquery, + "configured query-time aggregation anchor does not parse: {error}" + ); + return Ok(None); + } + }; + let Some(context) = self.build_query_execution_context_from_ast( + &anchor_ast, + &config, + &config.planned_subquery, + time, + ) else { + return Ok(None); + }; + let (anchor_labels, anchor_result) = + self.execute_context_result(context, false, false)?; + let QueryResult::Vector(anchor_values) = anchor_result else { + return Ok(None); + }; + let (labels, values) = apply_instant_pipeline( + anchor_labels, + anchor_values.values, + &config.query_time_aggregations, + ) + .map_err(QueryExecutionError::Native)?; + return Ok(Some(( + labels, + QueryResult::vector(values, Self::convert_query_time_to_data_time(time)), + ))); + } + } + let Some(context) = self.build_query_execution_context_from_parsed(&ast, &query, time) else { return Ok(None); @@ -1374,6 +1420,44 @@ impl SimpleEngine { return result; } + if let Some(config) = self.find_query_config(&query) { + if !config.query_time_aggregations.is_empty() { + let anchor_ast = match promql_parser::parser::parse(&config.planned_subquery) { + Ok(ast) => ast, + Err(error) => { + warn!( + query = %query, + planned_subquery = %config.planned_subquery, + "configured query-time aggregation anchor does not parse: {error}" + ); + return Ok(None); + } + }; + let Some(anchor_context) = self.build_query_execution_context_from_ast( + &anchor_ast, + &config, + &config.planned_subquery, + end, + ) else { + return Ok(None); + }; + let Some(context) = self.finish_range_context(anchor_context, start, end, step) + else { + return Ok(None); + }; + let anchor_results = self + .execute_observed_range_query_pipeline(&context, false, false) + .map_err(QueryExecutionError::Native)?; + let (labels, results) = apply_range_pipeline( + context.base.metadata.query_output_labels, + anchor_results, + &config.query_time_aggregations, + ) + .map_err(QueryExecutionError::Native)?; + return Ok(Some((labels, QueryResult::matrix(results)))); + } + } + let Some(context) = self.build_range_query_execution_context_from_parsed(&ast, &query, start, end, step) else { diff --git a/asap-query-engine/src/tests/structural_matching_tests.rs b/asap-query-engine/src/tests/structural_matching_tests.rs index d04d99aa..b9e9f8aa 100644 --- a/asap-query-engine/src/tests/structural_matching_tests.rs +++ b/asap-query-engine/src/tests/structural_matching_tests.rs @@ -6,9 +6,18 @@ #[cfg(test)] mod tests { - use crate::data_model::AggregationType; + use crate::data_model::{ + AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, PromQLSchema, + QueryConfig, SchemaConfig, + }; + use crate::engines::QueryResult; use crate::precompute_operators::sum_accumulator::SumAccumulator; use crate::tests::test_utilities::engine_factories::create_engine_single_pop; + use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeAggregationParameter, + QueryTimeGrouping, QueryTimeGroupingMode, + }; + use promql_utilities::data_model::KeyByLabelNames; #[test] fn test_structural_match_rate_query_finds_config() { @@ -32,6 +41,62 @@ mod tests { ); } + #[test] + fn exact_nested_query_config_executes_its_anchor_then_pipeline() { + let metric = "http_requests_total"; + let anchor = "sum by (job) (http_requests_total)"; + let query = "topk(1, sum by (job) (http_requests_total))"; + let engine = create_engine_single_pop( + metric, + AggregationType::Sum, + vec!["job"], + vec![ + ( + Some(vec!["api".to_string()]), + Box::new(SumAccumulator::with_sum(4.0)), + ), + ( + Some(vec!["worker".to_string()]), + Box::new(SumAccumulator::with_sum(9.0)), + ), + ], + anchor, + ); + let query_config = QueryConfig::with_plan( + query.to_string(), + anchor.to_string(), + vec![QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Topk, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Integer(1)), + }], + ) + .add_aggregation(AggregationReference::new(1, None)); + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + metric.to_string(), + KeyByLabelNames::new(vec!["job".to_string()]), + )), + query_configs: vec![query_config], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + + let (labels, result) = engine + .handle_query_promql(query.to_string(), 1_000.0) + .expect("the exact configured query should use its planned anchor"); + + assert_eq!(labels.labels, vec!["job"]); + let QueryResult::Vector(vector) = result else { + panic!("instant query should produce a vector"); + }; + assert_eq!(vector.values.len(), 1); + assert_eq!(vector.values[0].labels.labels, vec!["worker"]); + assert_eq!(vector.values[0].value, 9.0); + } + #[test] fn test_structural_match_wrong_metric_returns_none() { let engine = create_engine_single_pop( From 3b2763d30db5ab67c6f10730cb6fc570bbf6ef38 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 27 Sep 2026 21:34:46 -0400 Subject: [PATCH 06/22] test(query-engine): add nested aggregation e2e tracer --- .../tests/e2e_precompute_equivalence.rs | 85 ++++++++++++++++++- 1 file changed, 84 insertions(+), 1 deletion(-) diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 6a20ed3c..a56ba175 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -16,6 +16,7 @@ use serde_json::json; use std::collections::HashMap; use std::sync::Arc; +use asap_types::query_config::QueryTimeAggregation; use query_engine_rust::data_model::{ AggregationReference, InferenceConfig, PromQLSchema, QueryConfig, SchemaConfig, StreamingConfig, }; @@ -260,6 +261,16 @@ struct NativeDagScenario<'a> { impl NativeDagScenario<'_> { async fn build_engine(self) -> (SimpleEngine, String) { + let planned_subquery = self.query.to_string(); + self.build_engine_with_plan(&planned_subquery, Vec::new()) + .await + } + + async fn build_engine_with_plan( + self, + planned_subquery: &str, + query_time_aggregations: Vec, + ) -> (SimpleEngine, String) { let aggregation_ids: Vec = self .aggregation_configs .iter() @@ -305,7 +316,11 @@ impl NativeDagScenario<'_> { } let query_config = aggregation_ids.into_iter().fold( - QueryConfig::new(self.query.to_string()), + QueryConfig::with_plan( + self.query.to_string(), + planned_subquery.to_string(), + query_time_aggregations, + ), |config, aggregation_id| { config.add_aggregation(AggregationReference::new(aggregation_id, None)) }, @@ -362,6 +377,74 @@ fn assert_range_results_match( ); } +#[tokio::test] +async fn e2e_nested_topk_executes_after_its_planned_sum_anchor() { + use asap_types::query_config::{ + QueryTimeAggregationOperator, QueryTimeAggregationParameter, QueryTimeGrouping, + QueryTimeGroupingMode, + }; + + let metric = "nested_topk_requests"; + let anchor = "sum by (job) (nested_topk_requests)"; + let scenario = NativeDagScenario { + port: 19417, + metric, + query: "topk(1, sum by (job) (nested_topk_requests))", + aggregation_configs: vec![make_agg_config( + 17, + metric, + AggregationType::Sum, + "", + 1_000, + 0, + vec!["job"], + )], + schema_labels: vec!["instance".to_string(), "job".to_string()], + samples: vec![ + make_timeseries(metric, vec![("job", "api"), ("instance", "a")], 1_000, 2.0), + make_timeseries(metric, vec![("job", "api"), ("instance", "b")], 1_000, 3.0), + make_timeseries( + metric, + vec![("job", "worker"), ("instance", "c")], + 1_000, + 9.0, + ), + make_timeseries(metric, vec![], 3_000, 0.0), + ], + evaluation_time_seconds: 1.0, + base_interval_ms: 1_000, + }; + + let (engine, query) = scenario + .build_engine_with_plan( + anchor, + vec![QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Topk, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Integer(1)), + }], + ) + .await; + let (_, result) = engine + .handle_query_promql(query, 1.0) + .expect("nested query should execute through the planned anchor"); + let QueryResult::Vector(vector) = result else { + panic!("instant query should return a vector"); + }; + + assert_eq!(vector.values.len(), 1); + assert_eq!( + vector.values[0].labels.labels, + vec!["worker"], + "nested topk result: {:?}", + vector.values + ); + assert_eq!(vector.values[0].value, 9.0); +} + #[tokio::test] async fn e2e_sliding_precompute_outputs_compose_a_wider_query() { let port = 19402u16; From 77b6d9c6c5e6c898d30c116867752a0a394a40d1 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 27 Sep 2026 21:58:15 -0400 Subject: [PATCH 07/22] test(query-engine): cover nested topk e2e --- asap-query-engine/tests/e2e_precompute_equivalence.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index a56ba175..1cfa2ece 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -409,7 +409,13 @@ async fn e2e_nested_topk_executes_after_its_planned_sum_anchor() { 1_000, 9.0, ), - make_timeseries(metric, vec![], 3_000, 0.0), + make_timeseries(metric, vec![("job", "api"), ("instance", "a")], 3_000, 0.0), + make_timeseries( + metric, + vec![("job", "worker"), ("instance", "c")], + 3_000, + 0.0, + ), ], evaluation_time_seconds: 1.0, base_interval_ms: 1_000, From a95971c53769cad89e4a6478bbd8032bf03f207b Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 27 Sep 2026 22:06:02 -0400 Subject: [PATCH 08/22] test(query-engine): cover nested aggregation operators e2e --- .../tests/e2e_precompute_equivalence.rs | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 1cfa2ece..6a26419a 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -451,6 +451,111 @@ async fn e2e_nested_topk_executes_after_its_planned_sum_anchor() { assert_eq!(vector.values[0].value, 9.0); } +#[tokio::test] +async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { + use asap_types::query_config::{ + QueryTimeAggregationOperator, QueryTimeAggregationParameter, QueryTimeGrouping, + QueryTimeGroupingMode, + }; + + let metric = "nested_operator_matrix"; + let anchor = "sum by (job) (nested_operator_matrix)"; + let scenario = NativeDagScenario { + port: 19418, + metric, + query: anchor, + aggregation_configs: vec![make_agg_config( + 18, + metric, + AggregationType::Sum, + "", + 1_000, + 0, + vec!["job"], + )], + schema_labels: vec!["instance".to_string(), "job".to_string()], + samples: vec![ + make_timeseries(metric, vec![("job", "api"), ("instance", "a")], 1_000, 2.0), + make_timeseries(metric, vec![("job", "api"), ("instance", "b")], 1_000, 3.0), + make_timeseries( + metric, + vec![("job", "worker"), ("instance", "c")], + 1_000, + 9.0, + ), + make_timeseries(metric, vec![("job", "api"), ("instance", "a")], 3_000, 0.0), + make_timeseries( + metric, + vec![("job", "worker"), ("instance", "c")], + 3_000, + 0.0, + ), + ], + evaluation_time_seconds: 1.0, + base_interval_ms: 1_000, + }; + let (engine, _) = scenario.build_engine().await; + let operators = [ + ("sum", QueryTimeAggregationOperator::Sum, None), + ("count", QueryTimeAggregationOperator::Count, None), + ("avg", QueryTimeAggregationOperator::Avg, None), + ("min", QueryTimeAggregationOperator::Min, None), + ("max", QueryTimeAggregationOperator::Max, None), + ( + "quantile", + QueryTimeAggregationOperator::Quantile, + Some(QueryTimeAggregationParameter::Float(0.75)), + ), + ( + "topk", + QueryTimeAggregationOperator::Topk, + Some(QueryTimeAggregationParameter::Integer(3)), + ), + ]; + + for (name, operator, parameter) in &operators { + let stage = QueryTimeAggregation { + operator: operator.clone(), + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::By, + labels: vec!["job".to_string()], + }, + parameter: parameter.clone(), + }; + let query = match *name { + "quantile" => format!("quantile by (job) (0.75, {anchor})"), + "topk" => format!("topk by (job) (3, {anchor})"), + _ => format!("{name} by (job) ({anchor})"), + }; + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + metric.to_string(), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![ + "instance".to_string(), + "job".to_string(), + ]), + )), + query_configs: vec![QueryConfig::with_plan( + query.clone(), + anchor.to_string(), + vec![stage], + ) + .add_aggregation(AggregationReference::new(18, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + assert!( + engine.handle_query_promql(query.clone(), 1.0).is_some(), + "instant {name}" + ); + assert!( + engine + .handle_range_query_promql(query, 1.0, 2.0, 1.0) + .is_some(), + "range {name}" + ); + } +} + #[tokio::test] async fn e2e_sliding_precompute_outputs_compose_a_wider_query() { let port = 19402u16; From 467d1499eac3d7563723b85710944a507444b567 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 28 Sep 2026 09:26:05 -0400 Subject: [PATCH 09/22] test(query-engine): assert nested aggregation e2e values --- .../tests/e2e_precompute_equivalence.rs | 100 +++++++++++++----- 1 file changed, 76 insertions(+), 24 deletions(-) diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 6a26419a..f04feba1 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -496,36 +496,63 @@ async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { }; let (engine, _) = scenario.build_engine().await; let operators = [ - ("sum", QueryTimeAggregationOperator::Sum, None), - ("count", QueryTimeAggregationOperator::Count, None), - ("avg", QueryTimeAggregationOperator::Avg, None), - ("min", QueryTimeAggregationOperator::Min, None), - ("max", QueryTimeAggregationOperator::Max, None), + ( + "sum", + QueryTimeAggregationOperator::Sum, + None, + vec![(vec![], 14.0)], + ), + ( + "count", + QueryTimeAggregationOperator::Count, + None, + vec![(vec![], 2.0)], + ), + ( + "avg", + QueryTimeAggregationOperator::Avg, + None, + vec![(vec![], 7.0)], + ), + ( + "min", + QueryTimeAggregationOperator::Min, + None, + vec![(vec![], 5.0)], + ), + ( + "max", + QueryTimeAggregationOperator::Max, + None, + vec![(vec![], 9.0)], + ), ( "quantile", QueryTimeAggregationOperator::Quantile, Some(QueryTimeAggregationParameter::Float(0.75)), + vec![(vec![], 8.0)], ), ( "topk", QueryTimeAggregationOperator::Topk, Some(QueryTimeAggregationParameter::Integer(3)), + vec![(vec!["api"], 5.0), (vec!["worker"], 9.0)], ), ]; - for (name, operator, parameter) in &operators { + for (name, operator, parameter, expected) in operators { let stage = QueryTimeAggregation { - operator: operator.clone(), + operator, grouping: QueryTimeGrouping { - mode: QueryTimeGroupingMode::By, - labels: vec!["job".to_string()], + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), }, - parameter: parameter.clone(), + parameter, }; - let query = match *name { - "quantile" => format!("quantile by (job) (0.75, {anchor})"), - "topk" => format!("topk by (job) (3, {anchor})"), - _ => format!("{name} by (job) ({anchor})"), + let query = match name { + "quantile" => format!("quantile(0.75, {anchor})"), + "topk" => format!("topk(3, {anchor})"), + _ => format!("{name}({anchor})"), }; engine.update_inference_config(InferenceConfig { schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( @@ -543,16 +570,41 @@ async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { .add_aggregation(AggregationReference::new(18, None))], cleanup_policy: CleanupPolicy::NoCleanup, }); - assert!( - engine.handle_query_promql(query.clone(), 1.0).is_some(), - "instant {name}" - ); - assert!( - engine - .handle_range_query_promql(query, 1.0, 2.0, 1.0) - .is_some(), - "range {name}" - ); + let (_, instant) = engine + .handle_query_promql(query.clone(), 1.0) + .unwrap_or_else(|| panic!("instant {name}")); + let QueryResult::Vector(instant) = instant else { + panic!("instant {name} should return a vector"); + }; + let instant_values: Vec<(Vec, f64)> = instant + .values + .into_iter() + .map(|value| (value.labels.labels, value.value)) + .collect(); + let expected: Vec<(Vec, f64)> = expected + .into_iter() + .map(|(labels, value)| (labels.into_iter().map(str::to_string).collect(), value)) + .collect(); + assert_eq!(instant_values, expected, "instant {name}"); + + let (_, range) = engine + .handle_range_query_promql(query, 1.0, 2.0, 1.0) + .unwrap_or_else(|| panic!("range {name}")); + let QueryResult::Matrix(range) = range else { + panic!("range {name} should return a matrix"); + }; + let range_values: Vec<(Vec, f64)> = range + .values + .into_iter() + .map(|value| { + let sample = value + .samples + .last() + .expect("range result should have a final sample"); + (value.labels.labels, sample.value) + }) + .collect(); + assert_eq!(range_values, expected, "range {name}"); } } From cf48026c5913334bdf2f81d52b859c59dc900a58 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 28 Sep 2026 10:19:49 -0400 Subject: [PATCH 10/22] test(promql): add nested aggregation differential suite --- .../datasets/nested-aggregations.yaml | 328 +++++++++++++++ .../checked_in_fixture_validation_test.go | 1 + .../suites/nested-aggregations.yaml | 373 ++++++++++++++++++ 3 files changed, 702 insertions(+) create mode 100644 promql-compliance/datasets/nested-aggregations.yaml create mode 100644 promql-compliance/suites/nested-aggregations.yaml diff --git a/promql-compliance/datasets/nested-aggregations.yaml b/promql-compliance/datasets/nested-aggregations.yaml new file mode 100644 index 00000000..f4d93cd1 --- /dev/null +++ b/promql-compliance/datasets/nested-aggregations.yaml @@ -0,0 +1,328 @@ +name: nested-aggregations +series: + # Three jobs with four instances each. Values evolve at different rates so + # topk by (job) has meaningful, changing rankings across range steps. + - metric: data + labels: + job: frontend + instance: i-1 + samples: + - {offset_seconds: 0, value: 10} + - {offset_seconds: 60, value: 30} + - {offset_seconds: 120, value: 50} + - {offset_seconds: 180, value: 70} + - {offset_seconds: 240, value: 90} + - {offset_seconds: 300, value: 110} + - {offset_seconds: 360, value: 130} + - {offset_seconds: 420, value: 150} + - {offset_seconds: 480, value: 170} + - {offset_seconds: 540, value: 190} + - {offset_seconds: 600, value: 210} + - {offset_seconds: 660, value: 230} + - {offset_seconds: 720, value: 250} + - {offset_seconds: 780, value: 270} + - {offset_seconds: 840, value: 290} + - {offset_seconds: 900, value: 310} + - {offset_seconds: 960, value: 330} + - {offset_seconds: 1020, value: 350} + - {offset_seconds: 1080, value: 370} + - {offset_seconds: 1140, value: 390} + - {offset_seconds: 1200, value: 410} + - {offset_seconds: 1260, value: 430} + - metric: data + labels: + job: frontend + instance: i-2 + samples: + - {offset_seconds: 0, value: 20} + - {offset_seconds: 60, value: 40} + - {offset_seconds: 120, value: 60} + - {offset_seconds: 180, value: 80} + - {offset_seconds: 240, value: 100} + - {offset_seconds: 300, value: 120} + - {offset_seconds: 360, value: 140} + - {offset_seconds: 420, value: 160} + - {offset_seconds: 480, value: 180} + - {offset_seconds: 540, value: 200} + - {offset_seconds: 600, value: 220} + - {offset_seconds: 660, value: 240} + - {offset_seconds: 720, value: 260} + - {offset_seconds: 780, value: 280} + - {offset_seconds: 840, value: 300} + - {offset_seconds: 900, value: 320} + - {offset_seconds: 960, value: 340} + - {offset_seconds: 1020, value: 360} + - {offset_seconds: 1080, value: 380} + - {offset_seconds: 1140, value: 400} + - {offset_seconds: 1200, value: 420} + - {offset_seconds: 1260, value: 440} + - metric: data + labels: + job: frontend + instance: i-3 + samples: + - {offset_seconds: 0, value: 30} + - {offset_seconds: 60, value: 50} + - {offset_seconds: 120, value: 70} + - {offset_seconds: 180, value: 90} + - {offset_seconds: 240, value: 110} + - {offset_seconds: 300, value: 130} + - {offset_seconds: 360, value: 150} + - {offset_seconds: 420, value: 170} + - {offset_seconds: 480, value: 190} + - {offset_seconds: 540, value: 210} + - {offset_seconds: 600, value: 230} + - {offset_seconds: 660, value: 250} + - {offset_seconds: 720, value: 270} + - {offset_seconds: 780, value: 290} + - {offset_seconds: 840, value: 310} + - {offset_seconds: 900, value: 330} + - {offset_seconds: 960, value: 350} + - {offset_seconds: 1020, value: 370} + - {offset_seconds: 1080, value: 390} + - {offset_seconds: 1140, value: 410} + - {offset_seconds: 1200, value: 430} + - {offset_seconds: 1260, value: 450} + - metric: data + labels: + job: frontend + instance: i-4 + samples: + - {offset_seconds: 0, value: 40} + - {offset_seconds: 60, value: 60} + - {offset_seconds: 120, value: 80} + - {offset_seconds: 180, value: 100} + - {offset_seconds: 240, value: 120} + - {offset_seconds: 300, value: 140} + - {offset_seconds: 360, value: 160} + - {offset_seconds: 420, value: 180} + - {offset_seconds: 480, value: 200} + - {offset_seconds: 540, value: 220} + - {offset_seconds: 600, value: 240} + - {offset_seconds: 660, value: 260} + - {offset_seconds: 720, value: 280} + - {offset_seconds: 780, value: 300} + - {offset_seconds: 840, value: 320} + - {offset_seconds: 900, value: 340} + - {offset_seconds: 960, value: 360} + - {offset_seconds: 1020, value: 380} + - {offset_seconds: 1080, value: 400} + - {offset_seconds: 1140, value: 420} + - {offset_seconds: 1200, value: 440} + - {offset_seconds: 1260, value: 460} + - metric: data + labels: + job: backend + instance: i-1 + samples: + - {offset_seconds: 0, value: 160} + - {offset_seconds: 60, value: 155} + - {offset_seconds: 120, value: 150} + - {offset_seconds: 180, value: 145} + - {offset_seconds: 240, value: 140} + - {offset_seconds: 300, value: 135} + - {offset_seconds: 360, value: 130} + - {offset_seconds: 420, value: 125} + - {offset_seconds: 480, value: 120} + - {offset_seconds: 540, value: 115} + - {offset_seconds: 600, value: 110} + - {offset_seconds: 660, value: 105} + - {offset_seconds: 720, value: 100} + - {offset_seconds: 780, value: 95} + - {offset_seconds: 840, value: 90} + - {offset_seconds: 900, value: 85} + - {offset_seconds: 960, value: 80} + - {offset_seconds: 1020, value: 75} + - {offset_seconds: 1080, value: 70} + - {offset_seconds: 1140, value: 65} + - {offset_seconds: 1200, value: 60} + - {offset_seconds: 1260, value: 55} + - metric: data + labels: + job: backend + instance: i-2 + samples: + - {offset_seconds: 0, value: 150} + - {offset_seconds: 60, value: 145} + - {offset_seconds: 120, value: 140} + - {offset_seconds: 180, value: 135} + - {offset_seconds: 240, value: 130} + - {offset_seconds: 300, value: 125} + - {offset_seconds: 360, value: 120} + - {offset_seconds: 420, value: 115} + - {offset_seconds: 480, value: 110} + - {offset_seconds: 540, value: 105} + - {offset_seconds: 600, value: 100} + - {offset_seconds: 660, value: 95} + - {offset_seconds: 720, value: 90} + - {offset_seconds: 780, value: 85} + - {offset_seconds: 840, value: 80} + - {offset_seconds: 900, value: 75} + - {offset_seconds: 960, value: 70} + - {offset_seconds: 1020, value: 65} + - {offset_seconds: 1080, value: 60} + - {offset_seconds: 1140, value: 55} + - {offset_seconds: 1200, value: 50} + - {offset_seconds: 1260, value: 45} + - metric: data + labels: + job: backend + instance: i-3 + samples: + - {offset_seconds: 0, value: 140} + - {offset_seconds: 60, value: 135} + - {offset_seconds: 120, value: 130} + - {offset_seconds: 180, value: 125} + - {offset_seconds: 240, value: 120} + - {offset_seconds: 300, value: 115} + - {offset_seconds: 360, value: 110} + - {offset_seconds: 420, value: 105} + - {offset_seconds: 480, value: 100} + - {offset_seconds: 540, value: 95} + - {offset_seconds: 600, value: 90} + - {offset_seconds: 660, value: 85} + - {offset_seconds: 720, value: 80} + - {offset_seconds: 780, value: 75} + - {offset_seconds: 840, value: 70} + - {offset_seconds: 900, value: 65} + - {offset_seconds: 960, value: 60} + - {offset_seconds: 1020, value: 55} + - {offset_seconds: 1080, value: 50} + - {offset_seconds: 1140, value: 45} + - {offset_seconds: 1200, value: 40} + - {offset_seconds: 1260, value: 35} + - metric: data + labels: + job: backend + instance: i-4 + samples: + - {offset_seconds: 0, value: 130} + - {offset_seconds: 60, value: 125} + - {offset_seconds: 120, value: 120} + - {offset_seconds: 180, value: 115} + - {offset_seconds: 240, value: 110} + - {offset_seconds: 300, value: 105} + - {offset_seconds: 360, value: 100} + - {offset_seconds: 420, value: 95} + - {offset_seconds: 480, value: 90} + - {offset_seconds: 540, value: 85} + - {offset_seconds: 600, value: 80} + - {offset_seconds: 660, value: 75} + - {offset_seconds: 720, value: 70} + - {offset_seconds: 780, value: 65} + - {offset_seconds: 840, value: 60} + - {offset_seconds: 900, value: 55} + - {offset_seconds: 960, value: 50} + - {offset_seconds: 1020, value: 45} + - {offset_seconds: 1080, value: 40} + - {offset_seconds: 1140, value: 35} + - {offset_seconds: 1200, value: 30} + - {offset_seconds: 1260, value: 25} + - metric: data + labels: + job: worker + instance: i-1 + samples: + - {offset_seconds: 0, value: 15} + - {offset_seconds: 60, value: 55} + - {offset_seconds: 120, value: 95} + - {offset_seconds: 180, value: 135} + - {offset_seconds: 240, value: 175} + - {offset_seconds: 300, value: 215} + - {offset_seconds: 360, value: 255} + - {offset_seconds: 420, value: 295} + - {offset_seconds: 480, value: 335} + - {offset_seconds: 540, value: 375} + - {offset_seconds: 600, value: 415} + - {offset_seconds: 660, value: 455} + - {offset_seconds: 720, value: 495} + - {offset_seconds: 780, value: 535} + - {offset_seconds: 840, value: 575} + - {offset_seconds: 900, value: 615} + - {offset_seconds: 960, value: 655} + - {offset_seconds: 1020, value: 695} + - {offset_seconds: 1080, value: 735} + - {offset_seconds: 1140, value: 775} + - {offset_seconds: 1200, value: 815} + - {offset_seconds: 1260, value: 855} + - metric: data + labels: + job: worker + instance: i-2 + samples: + - {offset_seconds: 0, value: 25} + - {offset_seconds: 60, value: 65} + - {offset_seconds: 120, value: 105} + - {offset_seconds: 180, value: 145} + - {offset_seconds: 240, value: 185} + - {offset_seconds: 300, value: 225} + - {offset_seconds: 360, value: 265} + - {offset_seconds: 420, value: 305} + - {offset_seconds: 480, value: 345} + - {offset_seconds: 540, value: 385} + - {offset_seconds: 600, value: 425} + - {offset_seconds: 660, value: 465} + - {offset_seconds: 720, value: 505} + - {offset_seconds: 780, value: 545} + - {offset_seconds: 840, value: 585} + - {offset_seconds: 900, value: 625} + - {offset_seconds: 960, value: 665} + - {offset_seconds: 1020, value: 705} + - {offset_seconds: 1080, value: 745} + - {offset_seconds: 1140, value: 785} + - {offset_seconds: 1200, value: 825} + - {offset_seconds: 1260, value: 865} + - metric: data + labels: + job: worker + instance: i-3 + samples: + - {offset_seconds: 0, value: 35} + - {offset_seconds: 60, value: 75} + - {offset_seconds: 120, value: 115} + - {offset_seconds: 180, value: 155} + - {offset_seconds: 240, value: 195} + - {offset_seconds: 300, value: 235} + - {offset_seconds: 360, value: 275} + - {offset_seconds: 420, value: 315} + - {offset_seconds: 480, value: 355} + - {offset_seconds: 540, value: 395} + - {offset_seconds: 600, value: 435} + - {offset_seconds: 660, value: 475} + - {offset_seconds: 720, value: 515} + - {offset_seconds: 780, value: 555} + - {offset_seconds: 840, value: 595} + - {offset_seconds: 900, value: 635} + - {offset_seconds: 960, value: 675} + - {offset_seconds: 1020, value: 715} + - {offset_seconds: 1080, value: 755} + - {offset_seconds: 1140, value: 795} + - {offset_seconds: 1200, value: 835} + - {offset_seconds: 1260, value: 875} + - metric: data + labels: + job: worker + instance: i-4 + samples: + - {offset_seconds: 0, value: 45} + - {offset_seconds: 60, value: 85} + - {offset_seconds: 120, value: 125} + - {offset_seconds: 180, value: 165} + - {offset_seconds: 240, value: 205} + - {offset_seconds: 300, value: 245} + - {offset_seconds: 360, value: 285} + - {offset_seconds: 420, value: 325} + - {offset_seconds: 480, value: 365} + - {offset_seconds: 540, value: 405} + - {offset_seconds: 600, value: 445} + - {offset_seconds: 660, value: 485} + - {offset_seconds: 720, value: 525} + - {offset_seconds: 780, value: 565} + - {offset_seconds: 840, value: 605} + - {offset_seconds: 900, value: 645} + - {offset_seconds: 960, value: 685} + - {offset_seconds: 1020, value: 725} + - {offset_seconds: 1080, value: 765} + - {offset_seconds: 1140, value: 805} + - {offset_seconds: 1200, value: 845} + - {offset_seconds: 1260, value: 885} diff --git a/promql-compliance/runner/checked_in_fixture_validation_test.go b/promql-compliance/runner/checked_in_fixture_validation_test.go index e4a08898..94f2be31 100644 --- a/promql-compliance/runner/checked_in_fixture_validation_test.go +++ b/promql-compliance/runner/checked_in_fixture_validation_test.go @@ -66,6 +66,7 @@ func TestCheckedInNativeDagSuites(t *testing.T) { parser := promqlparser.NewParser(promqlparser.Options{}) for _, path := range []string{ "../suites/native-dag-aggregations.yaml", + "../suites/nested-aggregations.yaml", } { suite, err := LoadSuiteFile(path) if err != nil { diff --git a/promql-compliance/suites/nested-aggregations.yaml b/promql-compliance/suites/nested-aggregations.yaml new file mode 100644 index 00000000..736ad2f0 --- /dev/null +++ b/promql-compliance/suites/nested-aggregations.yaml @@ -0,0 +1,373 @@ +name: nested-aggregations +comparison_defaults: + value_tolerance: + relative: 0.01 + absolute: 0.000001 + +# Prometheus is the reference for each instant and range result. Every query +# starts from a planner-supported anchor and leaves one or two aggregation +# stages for ASAPQuery's query-time pipeline. +queries: + - name: one-stage-sum + expr: "sum by (job) (sum by (job, instance) (data))" + instant_offsets_seconds: &evaluation_offsets [300, 600, 900, 1200] + range: &evaluation_range + start_offset_seconds: 300 + end_offset_seconds: 1200 + step_seconds: 60 + - name: one-stage-count + expr: "count by (job) (sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: one-stage-avg + expr: "avg by (job) (sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: one-stage-min + expr: "min by (job) (sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: one-stage-max + expr: "max by (job) (sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: one-stage-topk-3 + expr: "topk by (job) (3, sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: one-stage-quantile-0-5 + expr: "quantile by (job) (0.5, sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: one-stage-quantile-0-75 + expr: "quantile by (job) (0.75, sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: one-stage-quantile-0-9 + expr: "quantile by (job) (0.9, sum by (job, instance) (data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-sum + expr: "sum by (job) (sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-count + expr: "count by (job) (sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-avg + expr: "avg by (job) (sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-min + expr: "min by (job) (sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-max + expr: "max by (job) (sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-topk-3 + expr: "topk by (job) (3, sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-quantile-0-5 + expr: "quantile by (job) (0.5, sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-quantile-0-75 + expr: "quantile by (job) (0.75, sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-sum-then-quantile-0-9 + expr: "quantile by (job) (0.9, sum by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-sum + expr: "sum by (job) (count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-count + expr: "count by (job) (count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-avg + expr: "avg by (job) (count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-min + expr: "min by (job) (count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-max + expr: "max by (job) (count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-topk-3 + expr: "topk by (job) (3, count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-quantile-0-5 + expr: "quantile by (job) (0.5, count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-quantile-0-75 + expr: "quantile by (job) (0.75, count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-count-then-quantile-0-9 + expr: "quantile by (job) (0.9, count by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-sum + expr: "sum by (job) (avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-count + expr: "count by (job) (avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-avg + expr: "avg by (job) (avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-min + expr: "min by (job) (avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-max + expr: "max by (job) (avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-topk-3 + expr: "topk by (job) (3, avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-quantile-0-5 + expr: "quantile by (job) (0.5, avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-quantile-0-75 + expr: "quantile by (job) (0.75, avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-avg-then-quantile-0-9 + expr: "quantile by (job) (0.9, avg by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-sum + expr: "sum by (job) (min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-count + expr: "count by (job) (min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-avg + expr: "avg by (job) (min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-min + expr: "min by (job) (min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-max + expr: "max by (job) (min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-topk-3 + expr: "topk by (job) (3, min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-quantile-0-5 + expr: "quantile by (job) (0.5, min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-quantile-0-75 + expr: "quantile by (job) (0.75, min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-min-then-quantile-0-9 + expr: "quantile by (job) (0.9, min by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-sum + expr: "sum by (job) (max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-count + expr: "count by (job) (max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-avg + expr: "avg by (job) (max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-min + expr: "min by (job) (max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-max + expr: "max by (job) (max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-topk-3 + expr: "topk by (job) (3, max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-quantile-0-5 + expr: "quantile by (job) (0.5, max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-quantile-0-75 + expr: "quantile by (job) (0.75, max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-max-then-quantile-0-9 + expr: "quantile by (job) (0.9, max by (job) (sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-sum + expr: "sum by (job) (topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-count + expr: "count by (job) (topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-avg + expr: "avg by (job) (topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-min + expr: "min by (job) (topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-max + expr: "max by (job) (topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-topk-3 + expr: "topk by (job) (3, topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-quantile-0-5 + expr: "quantile by (job) (0.5, topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-quantile-0-75 + expr: "quantile by (job) (0.75, topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-topk-3-then-quantile-0-9 + expr: "quantile by (job) (0.9, topk by (job) (3, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-sum + expr: "sum by (job) (quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-count + expr: "count by (job) (quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-avg + expr: "avg by (job) (quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-min + expr: "min by (job) (quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-max + expr: "max by (job) (quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-topk-3 + expr: "topk by (job) (3, quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-quantile-0-5 + expr: "quantile by (job) (0.5, quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-quantile-0-75 + expr: "quantile by (job) (0.75, quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-5-then-quantile-0-9 + expr: "quantile by (job) (0.9, quantile by (job) (0.5, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-sum + expr: "sum by (job) (quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-count + expr: "count by (job) (quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-avg + expr: "avg by (job) (quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-min + expr: "min by (job) (quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-max + expr: "max by (job) (quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-topk-3 + expr: "topk by (job) (3, quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-quantile-0-5 + expr: "quantile by (job) (0.5, quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-quantile-0-75 + expr: "quantile by (job) (0.75, quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-75-then-quantile-0-9 + expr: "quantile by (job) (0.9, quantile by (job) (0.75, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-sum + expr: "sum by (job) (quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-count + expr: "count by (job) (quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-avg + expr: "avg by (job) (quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-min + expr: "min by (job) (quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-max + expr: "max by (job) (quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-topk-3 + expr: "topk by (job) (3, quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-quantile-0-5 + expr: "quantile by (job) (0.5, quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-quantile-0-75 + expr: "quantile by (job) (0.75, quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range + - name: two-stage-quantile-0-9-then-quantile-0-9 + expr: "quantile by (job) (0.9, quantile by (job) (0.9, sum by (job, instance) (data)))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range From f7c3b876c54c1c07007a46616ba6e502a70624b1 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 28 Sep 2026 22:10:17 -0400 Subject: [PATCH 11/22] test(query-engine): cover query-time aggregation grouping --- asap-query-engine/src/tests/mod.rs | 1 + .../query_time_aggregation_public_tests.rs | 255 ++++++++++++++++++ 2 files changed, 256 insertions(+) create mode 100644 asap-query-engine/src/tests/query_time_aggregation_public_tests.rs diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 262119e0..a7f3a5f5 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -10,6 +10,7 @@ pub mod native_pipeline_merge_tests; pub mod native_range_query_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; +pub mod query_time_aggregation_public_tests; pub mod range_query_arithmetic_tests; pub mod sql_pattern_matching_tests; pub mod stage_e_instant_range_equivalence_tests; diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs new file mode 100644 index 00000000..a372f0de --- /dev/null +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -0,0 +1,255 @@ +//! Public PromQL coverage for query-time aggregation pipelines. + +#[cfg(test)] +mod tests { + use crate::data_model::{ + AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, PromQLSchema, + QueryConfig, SchemaConfig, + }; + use crate::engines::QueryResult; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::tests::test_utilities::engine_factories::create_engine_single_pop; + use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeAggregationParameter, + QueryTimeGrouping, QueryTimeGroupingMode, + }; + use promql_utilities::data_model::KeyByLabelNames; + + const METRIC: &str = "query_time_aggregation_metric"; + const ANCHOR: &str = "sum by (instance, job, region) (query_time_aggregation_metric)"; + + fn engine() -> crate::SimpleEngine { + create_engine_single_pop( + METRIC, + AggregationType::Sum, + vec!["instance", "job", "region"], + vec![ + ( + Some(vec!["a".to_string(), "api".to_string(), "east".to_string()]), + Box::new(SumAccumulator::with_sum(1.0)), + ), + ( + Some(vec!["b".to_string(), "api".to_string(), "east".to_string()]), + Box::new(SumAccumulator::with_sum(3.0)), + ), + ( + Some(vec![ + "a".to_string(), + "worker".to_string(), + "west".to_string(), + ]), + Box::new(SumAccumulator::with_sum(5.0)), + ), + ( + Some(vec![ + "b".to_string(), + "worker".to_string(), + "west".to_string(), + ]), + Box::new(SumAccumulator::with_sum(7.0)), + ), + ], + ANCHOR, + ) + } + + fn stage( + operator: QueryTimeAggregationOperator, + grouping: QueryTimeGroupingMode, + labels: &[&str], + parameter: Option, + ) -> QueryTimeAggregation { + QueryTimeAggregation { + operator, + grouping: QueryTimeGrouping { + mode: grouping, + labels: labels.iter().map(|label| (*label).to_string()).collect(), + }, + parameter, + } + } + + fn execute( + engine: &crate::SimpleEngine, + query: String, + stage: QueryTimeAggregation, + ) -> Vec<(Vec, f64)> { + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + METRIC.to_string(), + KeyByLabelNames::new(vec![ + "instance".to_string(), + "job".to_string(), + "region".to_string(), + ]), + )), + query_configs: vec![QueryConfig::with_plan( + query.clone(), + ANCHOR.to_string(), + vec![stage], + ) + .add_aggregation(AggregationReference::new(1, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + let (_, result) = engine + .handle_query_promql(query, 1_000.0) + .expect("configured nested query should execute locally"); + let QueryResult::Vector(vector) = result else { + panic!("instant query should return a vector"); + }; + vector + .values + .into_iter() + .map(|value| (value.labels.labels, value.value)) + .collect() + } + + #[test] + fn every_query_time_operator_respects_global_by_and_without_grouping() { + let engine = engine(); + let cases = [ + ( + "sum", + QueryTimeAggregationOperator::Sum, + None, + vec![(vec![], 16.0)], + vec![(vec!["api"], 4.0), (vec!["worker"], 12.0)], + ), + ( + "count", + QueryTimeAggregationOperator::Count, + None, + vec![(vec![], 4.0)], + vec![(vec!["api"], 2.0), (vec!["worker"], 2.0)], + ), + ( + "avg", + QueryTimeAggregationOperator::Avg, + None, + vec![(vec![], 4.0)], + vec![(vec!["api"], 2.0), (vec!["worker"], 6.0)], + ), + ( + "min", + QueryTimeAggregationOperator::Min, + None, + vec![(vec![], 1.0)], + vec![(vec!["api"], 1.0), (vec!["worker"], 5.0)], + ), + ( + "max", + QueryTimeAggregationOperator::Max, + None, + vec![(vec![], 7.0)], + vec![(vec!["api"], 3.0), (vec!["worker"], 7.0)], + ), + ( + "quantile", + QueryTimeAggregationOperator::Quantile, + Some(QueryTimeAggregationParameter::Float(0.5)), + vec![(vec![], 4.0)], + vec![(vec!["api"], 2.0), (vec!["worker"], 6.0)], + ), + ( + "topk", + QueryTimeAggregationOperator::Topk, + Some(QueryTimeAggregationParameter::Integer(3)), + vec![ + (vec!["a", "worker", "west"], 5.0), + (vec!["b", "api", "east"], 3.0), + (vec!["b", "worker", "west"], 7.0), + ], + vec![ + (vec!["a", "api", "east"], 1.0), + (vec!["a", "worker", "west"], 5.0), + (vec!["b", "api", "east"], 3.0), + (vec!["b", "worker", "west"], 7.0), + ], + ), + ]; + + for (name, operator, parameter, global_expected, grouped_expected) in cases { + let parameter_for_global = parameter.clone(); + let global_query = match name { + "quantile" => format!("quantile(0.5, {ANCHOR})"), + "topk" => format!("topk(3, {ANCHOR})"), + _ => format!("{name}({ANCHOR})"), + }; + assert_eq!( + execute( + &engine, + global_query, + stage( + operator.clone(), + QueryTimeGroupingMode::All, + &[], + parameter_for_global + ), + ), + expected(global_expected), + "global {name}" + ); + + let parameter_for_by = parameter.clone(); + let by_query = match name { + "quantile" => format!("quantile by (job) (0.5, {ANCHOR})"), + "topk" => format!("topk by (job) (3, {ANCHOR})"), + _ => format!("{name} by (job) ({ANCHOR})"), + }; + assert_eq!( + execute( + &engine, + by_query, + stage( + operator.clone(), + QueryTimeGroupingMode::By, + &["job"], + parameter_for_by, + ), + ), + expected(grouped_expected.clone()), + "by(job) {name}" + ); + + let without_query = match name { + "quantile" => format!("quantile without (instance) (0.5, {ANCHOR})"), + "topk" => format!("topk without (instance) (3, {ANCHOR})"), + _ => format!("{name} without (instance) ({ANCHOR})"), + }; + let without_expected = if name == "topk" { + grouped_expected + } else { + grouped_expected + .into_iter() + .map(|(labels, value)| { + ( + vec![labels[0], if labels[0] == "api" { "east" } else { "west" }], + value, + ) + }) + .collect() + }; + assert_eq!( + execute( + &engine, + without_query, + stage( + operator, + QueryTimeGroupingMode::Without, + &["instance"], + parameter, + ), + ), + expected(without_expected), + "without(instance) {name}" + ); + } + } + + fn expected(values: Vec<(Vec<&str>, f64)>) -> Vec<(Vec, f64)> { + values + .into_iter() + .map(|(labels, value)| (labels.into_iter().map(str::to_string).collect(), value)) + .collect() + } +} From 0f74902d5974e325b332e115eda4067f897ffe22 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 28 Sep 2026 22:29:18 -0400 Subject: [PATCH 12/22] fix(query-config): reject zero topk limits --- .../rs/asap_types/src/query_config.rs | 59 ++++++++++++++++++- 1 file changed, 58 insertions(+), 1 deletion(-) diff --git a/asap-common/dependencies/rs/asap_types/src/query_config.rs b/asap-common/dependencies/rs/asap_types/src/query_config.rs index eef9140f..f71b1d85 100644 --- a/asap-common/dependencies/rs/asap_types/src/query_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/query_config.rs @@ -78,10 +78,16 @@ impl QueryTimeAggregation { } match (&self.operator, &self.parameter) { + ( + QueryTimeAggregationOperator::Topk, + Some(QueryTimeAggregationParameter::Integer(k)), + ) if *k > 0 => {} ( QueryTimeAggregationOperator::Topk, Some(QueryTimeAggregationParameter::Integer(_)), - ) => {} + ) => { + return Err("topk requires a positive integer parameter".to_string()); + } (QueryTimeAggregationOperator::Topk, _) => { return Err("topk requires an integer parameter".to_string()); } @@ -138,3 +144,54 @@ impl QueryConfig { self } } + +#[cfg(test)] +mod tests { + use super::*; + + fn topk(k: u64) -> QueryTimeAggregation { + QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Topk, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Integer(k)), + } + } + + #[test] + fn topk_requires_a_positive_k() { + assert!(topk(1).validate().is_ok()); + assert_eq!( + topk(0).validate(), + Err("topk requires a positive integer parameter".into()) + ); + } + + #[test] + fn quantile_accepts_endpoints_and_rejects_out_of_range_values() { + for phi in [0.0, 1.0] { + assert!(QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Quantile, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Float(phi)), + } + .validate() + .is_ok()); + } + assert!(QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Quantile, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Float(1.01)), + } + .validate() + .is_err()); + } +} From 1313532d9b4fcdc9ca15368a828c12b053356683 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 28 Sep 2026 22:34:52 -0400 Subject: [PATCH 13/22] test(query-engine): cover query-time aggregation boundaries --- .../query_time_aggregation_public_tests.rs | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs index a372f0de..b54bb6c0 100644 --- a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -246,6 +246,49 @@ mod tests { } } + #[test] + fn topk_and_quantile_boundary_parameters_execute_through_the_public_handler() { + let engine = engine(); + let topk = |k| QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Topk, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Integer(k)), + }; + let quantile = |phi| QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Quantile, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Float(phi)), + }; + + assert_eq!( + execute(&engine, format!("topk(1, {ANCHOR})"), topk(1)), + expected(vec![(vec!["b", "worker", "west"], 7.0)]) + ); + assert_eq!( + execute(&engine, format!("topk(10, {ANCHOR})"), topk(10)), + expected(vec![ + (vec!["a", "api", "east"], 1.0), + (vec!["a", "worker", "west"], 5.0), + (vec!["b", "api", "east"], 3.0), + (vec!["b", "worker", "west"], 7.0), + ]) + ); + assert_eq!( + execute(&engine, format!("quantile(0, {ANCHOR})"), quantile(0.0)), + expected(vec![(vec![], 1.0)]) + ); + assert_eq!( + execute(&engine, format!("quantile(1, {ANCHOR})"), quantile(1.0)), + expected(vec![(vec![], 7.0)]) + ); + } + fn expected(values: Vec<(Vec<&str>, f64)>) -> Vec<(Vec, f64)> { values .into_iter() From ff280a07ec39d05f256db087e6c27644076ce953 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 29 Sep 2026 10:59:02 -0400 Subject: [PATCH 14/22] test(query-engine): cover query-time topk ties --- .../query_time_aggregation_public_tests.rs | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs index b54bb6c0..c9be7ef6 100644 --- a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -289,6 +289,37 @@ mod tests { ); } + #[test] + fn query_time_topk_breaks_ties_by_full_label_set() { + let engine = create_engine_single_pop( + METRIC, + AggregationType::Sum, + vec!["instance", "job", "region"], + vec![ + ( + Some(vec!["a".to_string(), "api".to_string(), "east".to_string()]), + Box::new(SumAccumulator::with_sum(7.0)), + ), + ( + Some(vec!["b".to_string(), "api".to_string(), "east".to_string()]), + Box::new(SumAccumulator::with_sum(7.0)), + ), + ], + ANCHOR, + ); + let topk = stage( + QueryTimeAggregationOperator::Topk, + QueryTimeGroupingMode::All, + &[], + Some(QueryTimeAggregationParameter::Integer(1)), + ); + + assert_eq!( + execute(&engine, format!("topk(1, {ANCHOR})"), topk), + expected(vec![(vec!["a", "api", "east"], 7.0)]) + ); + } + fn expected(values: Vec<(Vec<&str>, f64)>) -> Vec<(Vec, f64)> { values .into_iter() From bbc7f0a639586cd8cefaef400433a4a0fa8c2fbd Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 3 Oct 2026 23:01:00 -0400 Subject: [PATCH 15/22] feat(query-engine): execute nested aggregation DAG nodes --- asap-query-engine/src/engines/query_plan.rs | 73 +++++++++++++- .../src/engines/query_time_aggregation.rs | 4 +- .../src/engines/simple_engine/mod.rs | 96 ++++++++++++++----- .../src/engines/simple_engine/promql.rs | 36 +++---- .../query_time_aggregation_public_tests.rs | 7 +- .../src/tests/structural_matching_tests.rs | 7 +- .../tests/e2e_precompute_equivalence.rs | 21 ++-- 7 files changed, 193 insertions(+), 51 deletions(-) diff --git a/asap-query-engine/src/engines/query_plan.rs b/asap-query-engine/src/engines/query_plan.rs index 629d1464..c2f291db 100644 --- a/asap-query-engine/src/engines/query_plan.rs +++ b/asap-query-engine/src/engines/query_plan.rs @@ -1,7 +1,9 @@ //! Request-specific native query DAGs. +use crate::engines::query_time_aggregation::output_labels_for_aggregation; use crate::engines::simple_engine::{RangeQueryExecutionContext, StoreQueryParams}; use asap_types::enums::WindowType; +use asap_types::query_config::QueryTimeAggregation; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::Statistic; use tracing::debug; @@ -36,6 +38,12 @@ pub(crate) enum QueryPlanNode { input: NodeId, statistic: Statistic, query_kwargs: std::collections::HashMap, + output_labels: KeyByLabelNames, + }, + AggregateVector { + input: NodeId, + aggregation: QueryTimeAggregation, + input_labels: KeyByLabelNames, }, LimitTopK { input: NodeId, @@ -102,6 +110,7 @@ impl QueryPlan { pub(crate) fn compile_range( context: &RangeQueryExecutionContext, options: PlanOptions, + query_time_aggregations: &[QueryTimeAggregation], ) -> Result { let mut nodes = Vec::new(); let values_read = Self::push_read( @@ -143,8 +152,21 @@ impl QueryPlan { input: resolved, statistic: context.base.metadata.statistic_to_compute, query_kwargs: context.base.metadata.query_kwargs.clone(), + output_labels: context.base.metadata.query_output_labels.clone(), }, ); + let mut labels = context.base.metadata.query_output_labels.clone(); + for aggregation in query_time_aggregations { + root = Self::push( + &mut nodes, + QueryPlanNode::AggregateVector { + input: root, + aggregation: aggregation.clone(), + input_labels: labels.clone(), + }, + ); + labels = output_labels_for_aggregation(&labels, aggregation)?; + } if options.limit_topk && context.base.metadata.statistic_to_compute == Statistic::Topk { let k = context .base @@ -302,7 +324,7 @@ impl QueryPlan { values.0, keys.map(|id| format!("n{}", id.0)).unwrap_or_else(|| "self".to_string()) ), - QueryPlanNode::Estimate { input, statistic, query_kwargs } => { + 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) @@ -310,6 +332,9 @@ impl QueryPlan { QueryPlanNode::LimitTopK { input, k, .. } => { format!("n{index} LimitTopK(n{}, k={k})", input.0) } + QueryPlanNode::AggregateVector { input, aggregation, .. } => { + format!("n{index} AggregateVector(n{}, {:?})", input.0, aggregation) + } QueryPlanNode::Format { input, include_metric_name, metric } => format!( "n{index} Format(n{}, include_metric_name={include_metric_name}) metric={metric}", input.0 ), @@ -328,6 +353,7 @@ impl QueryPlanNode { Self::ComposeWindows { .. } => "ComposeWindows", Self::ResolveKeys { .. } => "ResolveKeys", Self::Estimate { .. } => "Estimate", + Self::AggregateVector { .. } => "AggregateVector", Self::LimitTopK { .. } => "LimitTopK", Self::Format { .. } => "Format", } @@ -338,6 +364,7 @@ impl QueryPlanNode { Self::StoreRead { .. } => Vec::new(), Self::ComposeWindows { input, .. } | Self::Estimate { input, .. } + | Self::AggregateVector { input, .. } | Self::LimitTopK { input, .. } | Self::Format { input, .. } => vec![*input], Self::ResolveKeys { values, keys } => { @@ -355,6 +382,10 @@ mod tests { use super::*; use crate::data_model::AggregationIdInfo; use crate::engines::simple_engine::{QueryExecutionContext, QueryMetadata, StoreQueryPlan}; + use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeGrouping, + QueryTimeGroupingMode, + }; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::AggregationType; use std::cell::RefCell; @@ -423,6 +454,7 @@ mod tests { limit_topk: false, format_output: false, }, + &[], ) .unwrap() .explain(); @@ -432,6 +464,41 @@ mod tests { assert!(explanation.ends_with("root: n5")); } + #[test] + fn nested_aggregation_pipeline_is_visible_as_ordered_plan_nodes() { + let explanation = QueryPlan::compile_range( + &context(), + PlanOptions { + limit_topk: false, + format_output: true, + }, + &[ + QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Sum, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: None, + }, + QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Max, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: None, + }, + ], + ) + .unwrap() + .explain(); + + assert!(explanation.contains("n4 AggregateVector(n3, QueryTimeAggregation { operator: Sum")); + assert!(explanation.contains("n5 AggregateVector(n4, QueryTimeAggregation { operator: Max")); + assert!(explanation.contains("n6 Format(n5")); + } + #[test] fn range_plan_keeps_every_output_timestamp() { let mut context = context(); @@ -443,6 +510,7 @@ mod tests { limit_topk: false, format_output: false, }, + &[], ) .unwrap() .explain(); @@ -467,6 +535,7 @@ mod tests { limit_topk: true, format_output: true, }, + &[], ) .unwrap() .explain(); @@ -487,6 +556,7 @@ mod tests { limit_topk: true, format_output: false, }, + &[], ) .expect_err("topk plan without k must fail loudly"); @@ -500,6 +570,7 @@ mod tests { input: NodeId(1), statistic: Statistic::Sum, query_kwargs: HashMap::new(), + output_labels: KeyByLabelNames::empty(), }], root: NodeId(0), }; diff --git a/asap-query-engine/src/engines/query_time_aggregation.rs b/asap-query-engine/src/engines/query_time_aggregation.rs index 59ab3ebd..d02e981b 100644 --- a/asap-query-engine/src/engines/query_time_aggregation.rs +++ b/asap-query-engine/src/engines/query_time_aggregation.rs @@ -9,7 +9,7 @@ use std::collections::BTreeMap; const METRIC_NAME_LABEL: &str = "__name__"; -fn output_labels( +pub(crate) fn output_labels_for_aggregation( input_labels: &KeyByLabelNames, aggregation: &QueryTimeAggregation, ) -> Result { @@ -130,7 +130,7 @@ fn apply_stage( input: Vec, aggregation: &QueryTimeAggregation, ) -> Result<(KeyByLabelNames, Vec), String> { - let output_labels = output_labels(input_labels, aggregation)?; + let output_labels = output_labels_for_aggregation(input_labels, aggregation)?; let grouping_labels = match aggregation.operator { QueryTimeAggregationOperator::Topk => match aggregation.grouping.mode { QueryTimeGroupingMode::All => KeyByLabelNames::empty(), diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 2f42cc86..43e624fd 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -208,7 +208,15 @@ enum NativePlanOutput { Read(TimestampedBucketsMap), Composed(ComposedRangeRead), Resolved(ResolvedRangeReads), - Results(Vec), + Results { + labels: KeyByLabelNames, + values: Vec, + }, +} + +struct RangePipelineOutput { + labels: KeyByLabelNames, + values: Vec, } struct NativePlanRuntime<'a> { @@ -290,22 +298,43 @@ impl QueryPlanRuntime for NativePlanRuntime<'_> { "ResolveKeys received incompatible inputs".into(), )), }, - QueryPlanNode::Estimate { .. } => match inputs { + QueryPlanNode::Estimate { output_labels, .. } => match inputs { [NativePlanOutput::Resolved(reads)] => self .engine .estimate_range_query(self.context, reads.clone()) - .map(NativePlanOutput::Results), + .map(|values| NativePlanOutput::Results { + labels: output_labels.clone(), + values, + }), _ => Err(QueryExecutionError::Native( "Estimate expected resolved reads".into(), )), }, + QueryPlanNode::AggregateVector { + aggregation, + input_labels, + .. + } => match inputs { + [NativePlanOutput::Results { values, .. }] => { + crate::engines::query_time_aggregation::apply_range_pipeline( + input_labels.clone(), + values.clone(), + std::slice::from_ref(aggregation), + ) + .map(|(labels, values)| NativePlanOutput::Results { labels, values }) + .map_err(QueryExecutionError::Native) + } + _ => Err(QueryExecutionError::Native( + "AggregateVector expected estimates".into(), + )), + }, QueryPlanNode::LimitTopK { k, grouping_labels, .. } => match inputs { - [NativePlanOutput::Results(results)] => self + [NativePlanOutput::Results { labels, values }] => self .engine .limit_range_topk( - results, + values, k, &SimpleEngine::topk_row_label_order( &self.context.base.metadata, @@ -315,7 +344,10 @@ impl QueryPlanRuntime for NativePlanRuntime<'_> { grouping_labels, ) .map_err(QueryExecutionError::Native) - .map(NativePlanOutput::Results), + .map(|values| NativePlanOutput::Results { + labels: labels.clone(), + values, + }), _ => Err(QueryExecutionError::Native( "LimitTopK expected estimates".into(), )), @@ -325,10 +357,12 @@ impl QueryPlanRuntime for NativePlanRuntime<'_> { metric, .. } => match inputs { - [NativePlanOutput::Results(results)] => Ok(NativePlanOutput::Results( - self.engine - .format_range_results(results, *include_metric_name, metric), - )), + [NativePlanOutput::Results { labels, values }] => Ok(NativePlanOutput::Results { + labels: labels.clone(), + values: self + .engine + .format_range_results(values, *include_metric_name, metric), + }), _ => Err(QueryExecutionError::Native( "result node expected estimates".into(), )), @@ -1625,11 +1659,14 @@ impl SimpleEngine { )) })?; - let range_results = self.execute_observed_range_query_pipeline( - &range_context, - enable_topk_limiting, - enable_topk_formatting, - )?; + let range_results = self + .execute_observed_range_query_pipeline( + &range_context, + enable_topk_limiting, + enable_topk_formatting, + &[], + )? + .values; let mut results: Vec = range_results .into_iter() @@ -2395,18 +2432,29 @@ impl SimpleEngine { context: &RangeQueryExecutionContext, enable_topk_limiting: bool, enable_topk_formatting: bool, - ) -> Result, QueryExecutionError> { + query_time_aggregations: &[asap_types::query_config::QueryTimeAggregation], + ) -> Result { Self::reject_off_grid_sliding_counter_query(context)?; #[cfg(feature = "native_query_legacy_test_support")] if matches!( self.native_range_execution_mode, NativeRangeExecutionMode::Legacy ) { - return self.execute_legacy_range_query_pipeline( - context, - enable_topk_limiting, - enable_topk_formatting, - ); + if !query_time_aggregations.is_empty() { + return Err(QueryExecutionError::Native( + "Legacy range execution does not support query-time aggregations".to_string(), + )); + } + return self + .execute_legacy_range_query_pipeline( + context, + enable_topk_limiting, + enable_topk_formatting, + ) + .map(|values| RangePipelineOutput { + labels: context.base.metadata.query_output_labels.clone(), + values, + }); } #[cfg(feature = "native_query_legacy_test_support")] let plan = if matches!( @@ -2421,6 +2469,7 @@ impl SimpleEngine { limit_topk: enable_topk_limiting, format_output: enable_topk_formatting, }, + query_time_aggregations, ) .map_err(QueryExecutionError::Native)? }; @@ -2431,6 +2480,7 @@ impl SimpleEngine { limit_topk: enable_topk_limiting, format_output: enable_topk_formatting, }, + query_time_aggregations, ) .map_err(QueryExecutionError::Native)?; debug!(plan = %plan.explain(), "Compiled native query plan"); @@ -2443,7 +2493,9 @@ impl SimpleEngine { QueryPlanExecutionError::InvalidPlan(reason) => QueryExecutionError::Native(reason), QueryPlanExecutionError::Node { source, .. } => source, })? { - NativePlanOutput::Results(results) => Ok(results), + NativePlanOutput::Results { labels, values } => { + Ok(RangePipelineOutput { labels, values }) + } _ => Err(QueryExecutionError::Native( "Query plan root did not produce results".to_string(), )), diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 604fb020..74e27a4b 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -10,7 +10,7 @@ use super::{ }; use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig}; use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVectorElement}; -use crate::engines::query_time_aggregation::{apply_instant_pipeline, apply_range_pipeline}; +use crate::engines::query_time_aggregation::apply_instant_pipeline; use asap_types::query_requirements::build_query_requirements_promql; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; @@ -766,7 +766,8 @@ impl SimpleEngine { // unformatted intermediate label representation until after the // arithmetic operation. let Some(results) = Self::map_local_execution_outcome( - self.execute_observed_range_query_pipeline(&ctx, true, false), + self.execute_observed_range_query_pipeline(&ctx, true, false, &[]) + .map(|output| output.values), )? else { return Ok(None); @@ -809,13 +810,15 @@ impl SimpleEngine { } // Binary arms need Topk limiting, but not final presentation formatting. let Some(lhs_results) = Self::map_local_execution_outcome( - self.execute_observed_range_query_pipeline(&lhs_ctx, true, false), + self.execute_observed_range_query_pipeline(&lhs_ctx, true, false, &[]) + .map(|output| output.values), )? else { return Ok(None); }; let Some(rhs_results) = Self::map_local_execution_outcome( - self.execute_observed_range_query_pipeline(&rhs_ctx, true, false), + self.execute_observed_range_query_pipeline(&rhs_ctx, true, false, &[]) + .map(|output| output.values), )? else { return Ok(None); @@ -1139,8 +1142,11 @@ impl SimpleEngine { ) else { return Ok(None); }; - let (anchor_labels, anchor_result) = - self.execute_context_result(context, false, false)?; + let Some((anchor_labels, anchor_result)) = + self.execute_context_result(context, false, false)? + else { + return Ok(None); + }; let QueryResult::Vector(anchor_values) = anchor_result else { return Ok(None); }; @@ -1445,16 +1451,13 @@ impl SimpleEngine { else { return Ok(None); }; - let anchor_results = self - .execute_observed_range_query_pipeline(&context, false, false) - .map_err(QueryExecutionError::Native)?; - let (labels, results) = apply_range_pipeline( - context.base.metadata.query_output_labels, - anchor_results, + let output = self.execute_observed_range_query_pipeline( + &context, + false, + false, &config.query_time_aggregations, - ) - .map_err(QueryExecutionError::Native)?; - return Ok(Some((labels, QueryResult::matrix(results)))); + )?; + return Ok(Some((output.labels, QueryResult::matrix(output.values)))); } } @@ -1468,7 +1471,8 @@ impl SimpleEngine { // instant's handle_query_promql -- both flags are no-ops unless this // query's statistic is Topk. let Some(results): Option> = Self::map_local_execution_outcome( - self.execute_observed_range_query_pipeline(&context, true, true), + self.execute_observed_range_query_pipeline(&context, true, true, &[]) + .map(|output| output.values), )? else { return Ok(None); diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs index c9be7ef6..60c405b2 100644 --- a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -91,9 +91,12 @@ mod tests { .add_aggregation(AggregationReference::new(1, None))], cleanup_policy: CleanupPolicy::NoCleanup, }); - let (_, result) = engine + let Some((_, result)) = engine .handle_query_promql(query, 1_000.0) - .expect("configured nested query should execute locally"); + .expect("configured nested query should execute locally") + else { + panic!("configured nested query should execute locally"); + }; let QueryResult::Vector(vector) = result else { panic!("instant query should return a vector"); }; diff --git a/asap-query-engine/src/tests/structural_matching_tests.rs b/asap-query-engine/src/tests/structural_matching_tests.rs index b9e9f8aa..10ca99eb 100644 --- a/asap-query-engine/src/tests/structural_matching_tests.rs +++ b/asap-query-engine/src/tests/structural_matching_tests.rs @@ -84,9 +84,12 @@ mod tests { cleanup_policy: CleanupPolicy::NoCleanup, }); - let (labels, result) = engine + let Some((labels, result)) = engine .handle_query_promql(query.to_string(), 1_000.0) - .expect("the exact configured query should use its planned anchor"); + .expect("the exact configured query should use its planned anchor") + else { + panic!("the exact configured query should use its planned anchor"); + }; assert_eq!(labels.labels, vec!["job"]); let QueryResult::Vector(vector) = result else { diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index f04feba1..5035933f 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -434,9 +434,12 @@ async fn e2e_nested_topk_executes_after_its_planned_sum_anchor() { }], ) .await; - let (_, result) = engine + let Some((_, result)) = engine .handle_query_promql(query, 1.0) - .expect("nested query should execute through the planned anchor"); + .expect("nested query should execute through the planned anchor") + else { + panic!("nested query should execute through the planned anchor"); + }; let QueryResult::Vector(vector) = result else { panic!("instant query should return a vector"); }; @@ -570,9 +573,12 @@ async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { .add_aggregation(AggregationReference::new(18, None))], cleanup_policy: CleanupPolicy::NoCleanup, }); - let (_, instant) = engine + let Some((_, instant)) = engine .handle_query_promql(query.clone(), 1.0) - .unwrap_or_else(|| panic!("instant {name}")); + .unwrap_or_else(|error| panic!("instant {name}: {error}")) + else { + panic!("instant {name} returned no local result"); + }; let QueryResult::Vector(instant) = instant else { panic!("instant {name} should return a vector"); }; @@ -587,9 +593,12 @@ async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { .collect(); assert_eq!(instant_values, expected, "instant {name}"); - let (_, range) = engine + let Some((_, range)) = engine .handle_range_query_promql(query, 1.0, 2.0, 1.0) - .unwrap_or_else(|| panic!("range {name}")); + .unwrap_or_else(|error| panic!("range {name}: {error}")) + else { + panic!("range {name} returned no local result"); + }; let QueryResult::Matrix(range) = range else { panic!("range {name} should return a matrix"); }; From b1e5cc9fe60021a7a687130aa572b1ef6166b6d2 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 3 Oct 2026 23:17:10 -0400 Subject: [PATCH 16/22] fix(query-engine): preserve missing query-time labels --- .../src/engines/query_time_aggregation.rs | 29 +++++++++---------- .../query_time_aggregation_public_tests.rs | 19 ++++++++++++ 2 files changed, 33 insertions(+), 15 deletions(-) diff --git a/asap-query-engine/src/engines/query_time_aggregation.rs b/asap-query-engine/src/engines/query_time_aggregation.rs index d02e981b..a33f146e 100644 --- a/asap-query-engine/src/engines/query_time_aggregation.rs +++ b/asap-query-engine/src/engines/query_time_aggregation.rs @@ -29,20 +29,13 @@ pub(crate) fn output_labels_for_aggregation( .cloned() .collect(), }; - for label in &labels { - if !input_labels.labels.contains(label) { - return Err(format!( - "query-time aggregation references unknown label '{label}'" - )); - } - } Ok(KeyByLabelNames::new(labels)) } fn label_indices( input_labels: &KeyByLabelNames, output_labels: &KeyByLabelNames, -) -> Result, String> { +) -> Vec> { output_labels .labels .iter() @@ -51,19 +44,25 @@ fn label_indices( .labels .iter() .position(|candidate| candidate == label) - .ok_or_else(|| format!("query-time aggregation references unknown label '{label}'")) }) .collect() } -fn select_labels(labels: &KeyByLabelValues, indices: &[usize]) -> Result { +fn select_labels( + labels: &KeyByLabelValues, + indices: &[Option], +) -> Result { indices .iter() .map(|index| { - labels - .get(*index) - .cloned() - .ok_or_else(|| "result labels do not match the configured label schema".to_string()) + index.map_or_else( + || Ok(String::new()), + |index| { + labels.get(index).cloned().ok_or_else(|| { + "result labels do not match the configured label schema".to_string() + }) + }, + ) }) .collect::, _>>() .map(KeyByLabelValues::new_with_labels) @@ -146,7 +145,7 @@ fn apply_stage( }, _ => output_labels.clone(), }; - let indices = label_indices(input_labels, &grouping_labels)?; + let indices = label_indices(input_labels, &grouping_labels); let mut groups: BTreeMap, Vec> = BTreeMap::new(); for element in input { if !element.value.is_finite() { diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs index 60c405b2..86570f10 100644 --- a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -292,6 +292,25 @@ mod tests { ); } + #[test] + fn grouping_by_a_missing_label_uses_prometheus_empty_label_value() { + let engine = engine(); + + assert_eq!( + execute( + &engine, + format!("sum by (missing) ({ANCHOR})"), + stage( + QueryTimeAggregationOperator::Sum, + QueryTimeGroupingMode::By, + &["missing"], + None, + ), + ), + expected(vec![(vec![""], 16.0)]) + ); + } + #[test] fn query_time_topk_breaks_ties_by_full_label_set() { let engine = create_engine_single_pop( From cd5f31e09e87186ee5abc1de18a742185d6b1bdb Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 3 Oct 2026 23:22:33 -0400 Subject: [PATCH 17/22] fix(planner): preserve empty without aggregations --- asap-planner-rs/src/planner/promql.rs | 31 +++++++++++++++---- .../query_time_aggregation_public_tests.rs | 24 ++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/asap-planner-rs/src/planner/promql.rs b/asap-planner-rs/src/planner/promql.rs index ee5b92a3..5fbe0eb7 100644 --- a/asap-planner-rs/src/planner/promql.rs +++ b/asap-planner-rs/src/planner/promql.rs @@ -59,12 +59,10 @@ fn aggregation_stage( labels: labels.labels.clone(), } } - Some(promql_parser::parser::LabelModifier::Exclude(labels)) if !labels.is_empty() => { - QueryTimeGrouping { - mode: QueryTimeGroupingMode::Without, - labels: labels.labels.clone(), - } - } + Some(promql_parser::parser::LabelModifier::Exclude(labels)) => QueryTimeGrouping { + mode: QueryTimeGroupingMode::Without, + labels: labels.labels.clone(), + }, _ => QueryTimeGrouping { mode: QueryTimeGroupingMode::All, labels: Vec::new(), @@ -440,3 +438,24 @@ impl SingleQueryProcessor { Ok((configs, cleanup_param)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_without_modifier_is_preserved_as_without() { + let expression = promql_parser::parser::parse("sum without () (requests_total)") + .expect("query should parse"); + let promql_parser::parser::Expr::Aggregate(aggregate) = expression else { + panic!("query should parse as an aggregation"); + }; + + let stage = aggregation_stage(&aggregate).expect("aggregation should be supported"); + assert!(matches!( + stage.grouping.mode, + QueryTimeGroupingMode::Without + )); + assert!(stage.grouping.labels.is_empty()); + } +} diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs index 86570f10..5cf62b80 100644 --- a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -311,6 +311,30 @@ mod tests { ); } + #[test] + fn empty_without_modifier_preserves_every_input_label() { + let engine = engine(); + + assert_eq!( + execute( + &engine, + format!("sum without () ({ANCHOR})"), + stage( + QueryTimeAggregationOperator::Sum, + QueryTimeGroupingMode::Without, + &[], + None, + ), + ), + expected(vec![ + (vec!["a", "api", "east"], 1.0), + (vec!["a", "worker", "west"], 5.0), + (vec!["b", "api", "east"], 3.0), + (vec!["b", "worker", "west"], 7.0), + ]) + ); + } + #[test] fn query_time_topk_breaks_ties_by_full_label_set() { let engine = create_engine_single_pop( From 35a424fe95c1c4d86ae73b7db6c8e678e8456306 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sat, 3 Oct 2026 23:24:07 -0400 Subject: [PATCH 18/22] fix(query-engine): apply nested topk anchors first --- asap-query-engine/src/engines/query_plan.rs | 24 ++-- .../src/engines/simple_engine/promql.rs | 131 +++++++++++++++++- 2 files changed, 136 insertions(+), 19 deletions(-) diff --git a/asap-query-engine/src/engines/query_plan.rs b/asap-query-engine/src/engines/query_plan.rs index c2f291db..2a951607 100644 --- a/asap-query-engine/src/engines/query_plan.rs +++ b/asap-query-engine/src/engines/query_plan.rs @@ -155,18 +155,6 @@ impl QueryPlan { output_labels: context.base.metadata.query_output_labels.clone(), }, ); - let mut labels = context.base.metadata.query_output_labels.clone(); - for aggregation in query_time_aggregations { - root = Self::push( - &mut nodes, - QueryPlanNode::AggregateVector { - input: root, - aggregation: aggregation.clone(), - input_labels: labels.clone(), - }, - ); - labels = output_labels_for_aggregation(&labels, aggregation)?; - } if options.limit_topk && context.base.metadata.statistic_to_compute == Statistic::Topk { let k = context .base @@ -186,6 +174,18 @@ impl QueryPlan { }, ); } + let mut labels = context.base.metadata.query_output_labels.clone(); + for aggregation in query_time_aggregations { + root = Self::push( + &mut nodes, + QueryPlanNode::AggregateVector { + input: root, + aggregation: aggregation.clone(), + input_labels: labels.clone(), + }, + ); + labels = output_labels_for_aggregation(&labels, aggregation)?; + } if options.format_output { root = Self::push( &mut nodes, diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 74e27a4b..9c52c36b 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -1143,7 +1143,7 @@ impl SimpleEngine { return Ok(None); }; let Some((anchor_labels, anchor_result)) = - self.execute_context_result(context, false, false)? + self.execute_context_result(context, true, false)? else { return Ok(None); }; @@ -1451,12 +1451,16 @@ impl SimpleEngine { else { return Ok(None); }; - let output = self.execute_observed_range_query_pipeline( - &context, - false, - false, - &config.query_time_aggregations, - )?; + let Some(output) = + Self::map_local_execution_outcome(self.execute_observed_range_query_pipeline( + &context, + true, + false, + &config.query_time_aggregations, + ))? + else { + return Ok(None); + }; return Ok(Some((output.labels, QueryResult::matrix(output.values)))); } } @@ -1560,6 +1564,10 @@ mod topk_pipeline_tests { use crate::stores::simple_map_store::SimpleMapStore; use crate::stores::Store; use crate::utils::http::convert_query_result_to_prometheus; + use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeGrouping, + QueryTimeGroupingMode, + }; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::Statistic; use std::collections::{HashMap, HashSet}; @@ -1832,6 +1840,115 @@ mod topk_pipeline_tests { } } + #[test] + fn nested_sum_applies_topk_anchor_before_its_outer_aggregation() { + let (engine, store) = build_topk_engine(); + let nested_query = "sum(topk(3, transfer_events))"; + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + METRIC.to_string(), + KeyByLabelNames::new(vec!["srcip".to_string()]), + )), + query_configs: vec![ + QueryConfig::with_plan( + nested_query.to_string(), + TOPK_QUERY.replace("10", "3"), + vec![QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Sum, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: None, + }], + ) + .add_aggregation(AggregationReference::new(AGG_ID, None)), + QueryConfig::new("topk(3, transfer_events)".to_string()) + .add_aggregation(AggregationReference::new(AGG_ID, None)), + ], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + + let context = engine + .build_query_execution_context_promql( + "topk(3, transfer_events)".to_string(), + QUERY_TIME, + ) + .expect("topk anchor should build a context"); + let window = &context.store_plan.values_query; + let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); + for i in 1..=15u64 { + sketch.inner.update(&format!("10.0.0.{i}"), (i * 10) as f64); + } + store + .insert_precomputed_output( + PrecomputedOutput::new(window.start_timestamp, window.end_timestamp, None, AGG_ID), + Box::new(sketch), + ) + .expect("insert should succeed"); + + let (_, result) = engine + .handle_query_promql(nested_query.to_string(), QUERY_TIME) + .expect("nested query should not fail") + .expect("nested query should execute locally"); + let QueryResult::Vector(vector) = result else { + panic!("nested instant query should return a vector"); + }; + assert_eq!(vector.values.len(), 1); + assert_eq!(vector.values[0].value, 420.0); + + let (_, result) = engine + .handle_range_query_promql(nested_query.to_string(), QUERY_TIME - 1.0, QUERY_TIME, 1.0) + .expect("nested range query should not fail") + .expect("nested range query should execute locally"); + let QueryResult::Matrix(matrix) = result else { + panic!("nested range query should return a matrix"); + }; + let end_sample = matrix + .values + .iter() + .flat_map(|element| &element.samples) + .find(|sample| sample.timestamp == (QUERY_TIME * 1_000.0) as u64) + .expect("range result should contain the end timestamp"); + assert_eq!(end_sample.value, 420.0); + } + + #[test] + fn nested_range_no_local_data_remains_a_prometheus_fallback() { + let (engine, _store) = build_topk_engine(); + let nested_query = "sum(topk(3, transfer_events))"; + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + METRIC.to_string(), + KeyByLabelNames::new(vec!["srcip".to_string()]), + )), + query_configs: vec![QueryConfig::with_plan( + nested_query.to_string(), + "topk(3, transfer_events)".to_string(), + vec![QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Sum, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: None, + }], + ) + .add_aggregation(AggregationReference::new(AGG_ID, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + + assert!(matches!( + engine.handle_range_query_promql( + nested_query.to_string(), + QUERY_TIME - 1.0, + QUERY_TIME, + 1.0, + ), + Ok(None) + )); + } + /// A topk leaf wrapped in an arithmetic binary expr (`topk(10, ...) + 0`) /// must still truncate to the top 10, while arithmetic output drops the /// metric name. Binary-arm evaluation must not apply standalone Topk From 072ef1e63dd7c76eef8bc0f5cbdc8147d0451f94 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 4 Oct 2026 10:46:19 -0400 Subject: [PATCH 19/22] fix(query-engine): validate nested aggregation plans --- .../rs/asap_types/src/query_config.rs | 20 ++++- asap-planner-rs/src/planner/promql.rs | 17 +++- .../src/engines/query_time_aggregation.rs | 1 - .../src/engines/simple_engine/promql.rs | 17 ++-- .../query_time_aggregation_public_tests.rs | 83 +++++++++++++++++-- .../tests/e2e_precompute_equivalence.rs | 13 ++- 6 files changed, 128 insertions(+), 23 deletions(-) diff --git a/asap-common/dependencies/rs/asap_types/src/query_config.rs b/asap-common/dependencies/rs/asap_types/src/query_config.rs index f71b1d85..644b32b4 100644 --- a/asap-common/dependencies/rs/asap_types/src/query_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/query_config.rs @@ -56,10 +56,8 @@ impl QueryTimeAggregation { QueryTimeGroupingMode::All if !self.grouping.labels.is_empty() => { return Err("all grouping cannot name labels".to_string()); } - QueryTimeGroupingMode::By | QueryTimeGroupingMode::Without - if self.grouping.labels.is_empty() => - { - return Err("by and without grouping must name at least one label".to_string()); + QueryTimeGroupingMode::By if self.grouping.labels.is_empty() => { + return Err("by grouping must name at least one label".to_string()); } _ => {} } @@ -194,4 +192,18 @@ mod tests { .validate() .is_err()); } + + #[test] + fn without_empty_labels_is_a_valid_promql_grouping() { + assert!(QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Sum, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::Without, + labels: Vec::new(), + }, + parameter: None, + } + .validate() + .is_ok()); + } } diff --git a/asap-planner-rs/src/planner/promql.rs b/asap-planner-rs/src/planner/promql.rs index 5fbe0eb7..744cbbb9 100644 --- a/asap-planner-rs/src/planner/promql.rs +++ b/asap-planner-rs/src/planner/promql.rs @@ -75,7 +75,7 @@ fn aggregation_stage( else { return None; }; - if number.val < 0.0 || number.val.fract() != 0.0 { + if number.val <= 0.0 || number.val.fract() != 0.0 { return None; } Some(QueryTimeAggregationParameter::Integer(number.val as u64)) @@ -85,6 +85,9 @@ fn aggregation_stage( else { return None; }; + if !number.val.is_finite() || !(0.0..=1.0).contains(&number.val) { + return None; + } Some(QueryTimeAggregationParameter::Float(number.val)) } _ => None, @@ -458,4 +461,16 @@ mod tests { )); assert!(stage.grouping.labels.is_empty()); } + + #[test] + fn invalid_query_time_parameters_are_not_planned() { + for query in ["topk(0, requests_total)", "quantile(1.1, requests_total)"] { + let expression = promql_parser::parser::parse(query).expect("query should parse"); + let promql_parser::parser::Expr::Aggregate(aggregate) = expression else { + panic!("query should parse as an aggregation"); + }; + + assert!(aggregation_stage(&aggregate).is_none(), "{query}"); + } + } } diff --git a/asap-query-engine/src/engines/query_time_aggregation.rs b/asap-query-engine/src/engines/query_time_aggregation.rs index a33f146e..e11dcda1 100644 --- a/asap-query-engine/src/engines/query_time_aggregation.rs +++ b/asap-query-engine/src/engines/query_time_aggregation.rs @@ -168,7 +168,6 @@ fn apply_stage( group.truncate(k); results.extend(group); } - results.sort_by(|left, right| left.labels.labels.cmp(&right.labels.labels)); return Ok((output_labels, results)); } diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 9c52c36b..4bce34e9 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -437,6 +437,9 @@ impl SimpleEngine { Expr::Paren(paren) => self.resolve_arm_leaf_context(&paren.expr, time), other => { let config = self.find_query_config_promql_structural(other)?; + if !config.query_time_aggregations.is_empty() { + return None; + } let ctx = self.build_query_execution_context_from_ast( other, &config, @@ -1126,12 +1129,9 @@ impl SimpleEngine { let anchor_ast = match promql_parser::parser::parse(&config.planned_subquery) { Ok(ast) => ast, Err(error) => { - warn!( - query = %query, - planned_subquery = %config.planned_subquery, + return Err(QueryExecutionError::Native(format!( "configured query-time aggregation anchor does not parse: {error}" - ); - return Ok(None); + ))); } }; let Some(context) = self.build_query_execution_context_from_ast( @@ -1431,12 +1431,9 @@ impl SimpleEngine { let anchor_ast = match promql_parser::parser::parse(&config.planned_subquery) { Ok(ast) => ast, Err(error) => { - warn!( - query = %query, - planned_subquery = %config.planned_subquery, + return Err(QueryExecutionError::Native(format!( "configured query-time aggregation anchor does not parse: {error}" - ); - return Ok(None); + ))); } }; let Some(anchor_context) = self.build_query_execution_context_from_ast( diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs index 5cf62b80..b2d6ee1f 100644 --- a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -158,15 +158,15 @@ mod tests { QueryTimeAggregationOperator::Topk, Some(QueryTimeAggregationParameter::Integer(3)), vec![ + (vec!["b", "worker", "west"], 7.0), (vec!["a", "worker", "west"], 5.0), (vec!["b", "api", "east"], 3.0), - (vec!["b", "worker", "west"], 7.0), ], vec![ - (vec!["a", "api", "east"], 1.0), - (vec!["a", "worker", "west"], 5.0), (vec!["b", "api", "east"], 3.0), + (vec!["a", "api", "east"], 1.0), (vec!["b", "worker", "west"], 7.0), + (vec!["a", "worker", "west"], 5.0), ], ), ]; @@ -276,10 +276,10 @@ mod tests { assert_eq!( execute(&engine, format!("topk(10, {ANCHOR})"), topk(10)), expected(vec![ - (vec!["a", "api", "east"], 1.0), + (vec!["b", "worker", "west"], 7.0), (vec!["a", "worker", "west"], 5.0), (vec!["b", "api", "east"], 3.0), - (vec!["b", "worker", "west"], 7.0), + (vec!["a", "api", "east"], 1.0), ]) ); assert_eq!( @@ -335,6 +335,79 @@ mod tests { ); } + #[test] + fn malformed_planned_subquery_fails_loudly_for_instant_and_range() { + let engine = engine(); + let query = "sum(query_time_aggregation_metric)".to_string(); + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + METRIC.to_string(), + KeyByLabelNames::new(vec![ + "instance".to_string(), + "job".to_string(), + "region".to_string(), + ]), + )), + query_configs: vec![QueryConfig::with_plan( + query.clone(), + "sum(".to_string(), + vec![stage( + QueryTimeAggregationOperator::Sum, + QueryTimeGroupingMode::All, + &[], + None, + )], + ) + .add_aggregation(AggregationReference::new(1, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + + for result in [ + engine.handle_query_promql(query.clone(), 1_000.0), + engine.handle_range_query_promql(query.clone(), 999.0, 1_000.0, 1.0), + ] { + assert!(matches!( + result, + Err(crate::QueryExecutionError::Native(message)) + if message.contains("configured query-time aggregation anchor does not parse") + )); + } + } + + #[test] + fn binary_expression_with_a_query_time_pipeline_falls_back_until_complete_dag_support() { + let engine = engine(); + let arm = format!("sum({ANCHOR})"); + let query = format!("{arm} + 1"); + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + METRIC.to_string(), + KeyByLabelNames::new(vec![ + "instance".to_string(), + "job".to_string(), + "region".to_string(), + ]), + )), + query_configs: vec![QueryConfig::with_plan( + arm, + ANCHOR.to_string(), + vec![stage( + QueryTimeAggregationOperator::Sum, + QueryTimeGroupingMode::All, + &[], + None, + )], + ) + .add_aggregation(AggregationReference::new(1, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + + assert!(matches!( + engine.handle_query_promql(query, 1_000.0), + Ok(None) + )); + } + #[test] fn query_time_topk_breaks_ties_by_full_label_set() { let engine = create_engine_single_pop( diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 5035933f..720f5aa9 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -539,7 +539,7 @@ async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { "topk", QueryTimeAggregationOperator::Topk, Some(QueryTimeAggregationParameter::Integer(3)), - vec![(vec!["api"], 5.0), (vec!["worker"], 9.0)], + vec![(vec!["worker"], 9.0), (vec!["api"], 5.0)], ), ]; @@ -593,6 +593,15 @@ async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { .collect(); assert_eq!(instant_values, expected, "instant {name}"); + let range_expected = if name == "topk" { + vec![ + (vec!["api".to_string()], 5.0), + (vec!["worker".to_string()], 9.0), + ] + } else { + expected.clone() + }; + let Some((_, range)) = engine .handle_range_query_promql(query, 1.0, 2.0, 1.0) .unwrap_or_else(|error| panic!("range {name}: {error}")) @@ -613,7 +622,7 @@ async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { (value.labels.labels, sample.value) }) .collect(); - assert_eq!(range_values, expected, "range {name}"); + assert_eq!(range_values, range_expected, "range {name}"); } } From b0a14f9e1e8f67e5288bd9cc5ce89041a5bc43d5 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 4 Oct 2026 13:20:51 -0400 Subject: [PATCH 20/22] fix(query-engine): preserve nested topk labels --- asap-query-engine/src/engines/query_plan.rs | 15 +- .../src/engines/query_time_aggregation.rs | 21 ++ .../src/engines/simple_engine/promql.rs | 23 ++- .../query_time_aggregation_public_tests.rs | 83 +++++++- .../tests/e2e_precompute_equivalence.rs | 194 ++++++++++++++++++ 5 files changed, 322 insertions(+), 14 deletions(-) diff --git a/asap-query-engine/src/engines/query_plan.rs b/asap-query-engine/src/engines/query_plan.rs index 2a951607..2384b75e 100644 --- a/asap-query-engine/src/engines/query_plan.rs +++ b/asap-query-engine/src/engines/query_plan.rs @@ -174,6 +174,19 @@ impl QueryPlan { }, ); } + let materialize_metric_name = !query_time_aggregations.is_empty() + && context.base.metadata.statistic_to_compute == Statistic::Topk + && context.base.metadata.keep_metric_name; + if materialize_metric_name { + root = Self::push( + &mut nodes, + QueryPlanNode::Format { + input: root, + include_metric_name: true, + metric: context.base.metric.clone(), + }, + ); + } let mut labels = context.base.metadata.query_output_labels.clone(); for aggregation in query_time_aggregations { root = Self::push( @@ -186,7 +199,7 @@ impl QueryPlan { ); labels = output_labels_for_aggregation(&labels, aggregation)?; } - if options.format_output { + if options.format_output && !materialize_metric_name { root = Self::push( &mut nodes, QueryPlanNode::Format { diff --git a/asap-query-engine/src/engines/query_time_aggregation.rs b/asap-query-engine/src/engines/query_time_aggregation.rs index e11dcda1..48495926 100644 --- a/asap-query-engine/src/engines/query_time_aggregation.rs +++ b/asap-query-engine/src/engines/query_time_aggregation.rs @@ -32,6 +32,27 @@ pub(crate) fn output_labels_for_aggregation( Ok(KeyByLabelNames::new(labels)) } +pub(crate) fn pipeline_supports_labels( + input_labels: &KeyByLabelNames, + pipeline: &[QueryTimeAggregation], +) -> Result { + let mut labels = input_labels.clone(); + for aggregation in pipeline { + if !matches!(aggregation.operator, QueryTimeAggregationOperator::Topk) + && matches!(aggregation.grouping.mode, QueryTimeGroupingMode::By) + && aggregation + .grouping + .labels + .iter() + .any(|label| !labels.labels.contains(label)) + { + return Ok(false); + } + labels = output_labels_for_aggregation(&labels, aggregation)?; + } + Ok(true) +} + fn label_indices( input_labels: &KeyByLabelNames, output_labels: &KeyByLabelNames, diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 4bce34e9..0ecc764d 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -10,7 +10,7 @@ use super::{ }; use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig}; use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVectorElement}; -use crate::engines::query_time_aggregation::apply_instant_pipeline; +use crate::engines::query_time_aggregation::{apply_instant_pipeline, pipeline_supports_labels}; use asap_types::query_requirements::build_query_requirements_promql; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; @@ -1142,14 +1142,25 @@ impl SimpleEngine { ) else { return Ok(None); }; + let anchor_metric = context.metric.clone(); let Some((anchor_labels, anchor_result)) = self.execute_context_result(context, true, false)? else { return Ok(None); }; - let QueryResult::Vector(anchor_values) = anchor_result else { + if !pipeline_supports_labels(&anchor_labels, &config.query_time_aggregations) + .map_err(QueryExecutionError::Native)? + { + return Ok(None); + } + let QueryResult::Vector(mut anchor_values) = anchor_result else { return Ok(None); }; + if anchor_labels.labels.first().map(String::as_str) == Some(METRIC_NAME_LABEL) { + for value in &mut anchor_values.values { + Self::prepend_metric_name(&anchor_metric, &mut value.labels); + } + } let (labels, values) = apply_instant_pipeline( anchor_labels, anchor_values.values, @@ -1448,6 +1459,14 @@ impl SimpleEngine { else { return Ok(None); }; + if !pipeline_supports_labels( + &context.base.metadata.query_output_labels, + &config.query_time_aggregations, + ) + .map_err(QueryExecutionError::Native)? + { + return Ok(None); + } let Some(output) = Self::map_local_execution_outcome(self.execute_observed_range_query_pipeline( &context, diff --git a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs index b2d6ee1f..5035f1ba 100644 --- a/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs +++ b/asap-query-engine/src/tests/query_time_aggregation_public_tests.rs @@ -293,22 +293,36 @@ mod tests { } #[test] - fn grouping_by_a_missing_label_uses_prometheus_empty_label_value() { + fn grouping_by_a_missing_label_falls_back_to_prometheus() { let engine = engine(); - - assert_eq!( - execute( - &engine, - format!("sum by (missing) ({ANCHOR})"), - stage( + let query = format!("sum by (missing) ({ANCHOR})"); + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + METRIC.to_string(), + KeyByLabelNames::new(vec![ + "instance".to_string(), + "job".to_string(), + "region".to_string(), + ]), + )), + query_configs: vec![QueryConfig::with_plan( + query.clone(), + ANCHOR.to_string(), + vec![stage( QueryTimeAggregationOperator::Sum, QueryTimeGroupingMode::By, &["missing"], None, - ), - ), - expected(vec![(vec![""], 16.0)]) - ); + )], + ) + .add_aggregation(AggregationReference::new(1, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + + assert!(matches!( + engine.handle_query_promql(query, 1_000.0), + Ok(None) + )); } #[test] @@ -408,6 +422,53 @@ mod tests { )); } + #[test] + fn pipeline_falls_back_when_a_later_by_stage_uses_a_removed_label() { + let engine = engine(); + let query = format!("max by (instance) (sum by (job) ({ANCHOR}))"); + engine.update_inference_config(InferenceConfig { + schema: SchemaConfig::PromQL(PromQLSchema::new().add_metric( + METRIC.to_string(), + KeyByLabelNames::new(vec![ + "instance".to_string(), + "job".to_string(), + "region".to_string(), + ]), + )), + query_configs: vec![QueryConfig::with_plan( + query.clone(), + ANCHOR.to_string(), + vec![ + stage( + QueryTimeAggregationOperator::Sum, + QueryTimeGroupingMode::By, + &["job"], + None, + ), + stage( + QueryTimeAggregationOperator::Max, + QueryTimeGroupingMode::By, + &["instance"], + None, + ), + ], + ) + .add_aggregation(AggregationReference::new(1, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }); + + // Prometheus omits an absent grouping label; this native representation + // cannot distinguish it from an explicitly empty value. + assert!(matches!( + engine.handle_query_promql(query.clone(), 1_000.0), + Ok(None) + )); + assert!(matches!( + engine.handle_range_query_promql(query, 999.0, 1_000.0, 1.0), + Ok(None) + )); + } + #[test] fn query_time_topk_breaks_ties_by_full_label_set() { let engine = create_engine_single_pop( diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 720f5aa9..59641c57 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -454,6 +454,200 @@ async fn e2e_nested_topk_executes_after_its_planned_sum_anchor() { assert_eq!(vector.values[0].value, 9.0); } +#[tokio::test] +async fn e2e_nested_aggregation_preserves_plain_topk_labels() { + use asap_types::query_config::{ + QueryTimeAggregationOperator, QueryTimeGrouping, QueryTimeGroupingMode, + }; + + let metric = "nested_plain_topk_labels"; + let anchor = format!("topk(3, {metric})"); + let query = format!("sum by (srcip) ({anchor})"); + let mut config = make_agg_config_full( + 19, + metric, + AggregationType::CountMinSketchWithHeap, + "count", + 1_000, + 0, + vec![], + vec!["srcip"], + ); + config.parameters.insert("depth".to_string(), json!(3_u64)); + config + .parameters + .insert("width".to_string(), json!(128_u64)); + config + .parameters + .insert("heapsize".to_string(), json!(16_u64)); + let samples = [ + ("10.0.0.1", 5), + ("10.0.0.2", 4), + ("10.0.0.3", 3), + ("10.0.0.4", 2), + ] + .into_iter() + .flat_map(|(srcip, count)| { + std::iter::repeat_with(move || make_timeseries(metric, vec![("srcip", srcip)], 1_000, 1.0)) + .take(count) + }) + .chain(std::iter::once(make_timeseries(metric, vec![], 3_000, 0.0))) + .collect(); + let scenario = NativeDagScenario { + port: 19425, + metric, + query: &query, + aggregation_configs: vec![config], + schema_labels: vec!["srcip".to_string()], + samples, + evaluation_time_seconds: 1.0, + base_interval_ms: 1_000, + }; + let (engine, query) = scenario + .build_engine_with_plan( + &anchor, + vec![QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Sum, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::By, + labels: vec!["srcip".to_string()], + }, + parameter: None, + }], + ) + .await; + + let expected = vec![ + (vec!["10.0.0.1".to_string()], 5.0), + (vec!["10.0.0.2".to_string()], 4.0), + (vec!["10.0.0.3".to_string()], 3.0), + ]; + let (_, instant) = engine + .handle_query_promql(query.clone(), 1.0) + .expect("instant query should execute locally") + .expect("instant query should return a result"); + let QueryResult::Vector(instant) = instant else { + panic!("instant query should return a vector"); + }; + assert_eq!( + instant + .values + .into_iter() + .map(|value| (value.labels.labels, value.value)) + .collect::>(), + expected + ); + + let (_, range) = engine + .handle_range_query_promql(query, 1.0, 2.0, 1.0) + .expect("range query should execute locally") + .expect("range query should return a result"); + let QueryResult::Matrix(range) = range else { + panic!("range query should return a matrix"); + }; + assert_eq!( + range + .values + .into_iter() + .map(|value| { + ( + value.labels.labels, + value.samples.last().expect("range sample").value, + ) + }) + .collect::>(), + expected + ); +} + +#[tokio::test] +async fn e2e_nested_topk_preserves_plain_topk_metric_name() { + use asap_types::query_config::{ + QueryTimeAggregationOperator, QueryTimeAggregationParameter, QueryTimeGrouping, + QueryTimeGroupingMode, + }; + + let metric = "nested_plain_topk_metric_name"; + let anchor = format!("topk(3, {metric})"); + let query = format!("topk(1, {anchor})"); + let mut config = make_agg_config_full( + 20, + metric, + AggregationType::CountMinSketchWithHeap, + "count", + 1_000, + 0, + vec![], + vec!["srcip"], + ); + config.parameters.insert("depth".to_string(), json!(3_u64)); + config + .parameters + .insert("width".to_string(), json!(128_u64)); + config + .parameters + .insert("heapsize".to_string(), json!(16_u64)); + let samples = [("10.0.0.1", 5), ("10.0.0.2", 4), ("10.0.0.3", 3)] + .into_iter() + .flat_map(|(srcip, count)| { + std::iter::repeat_with(move || { + make_timeseries(metric, vec![("srcip", srcip)], 1_000, 1.0) + }) + .take(count) + }) + .chain(std::iter::once(make_timeseries(metric, vec![], 3_000, 0.0))) + .collect(); + let scenario = NativeDagScenario { + port: 19426, + metric, + query: &query, + aggregation_configs: vec![config], + schema_labels: vec!["srcip".to_string()], + samples, + evaluation_time_seconds: 1.0, + base_interval_ms: 1_000, + }; + let (engine, query) = scenario + .build_engine_with_plan( + &anchor, + vec![QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Topk, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::All, + labels: Vec::new(), + }, + parameter: Some(QueryTimeAggregationParameter::Integer(1)), + }], + ) + .await; + + let expected_labels = vec![metric.to_string(), "10.0.0.1".to_string()]; + let (_, instant) = engine + .handle_query_promql(query.clone(), 1.0) + .expect("instant query should execute locally") + .expect("instant query should return a result"); + let QueryResult::Vector(instant) = instant else { + panic!("instant query should return a vector"); + }; + assert_eq!(instant.values.len(), 1); + assert_eq!(instant.values[0].labels.labels, expected_labels); + assert_eq!(instant.values[0].value, 5.0); + + let (_, range) = engine + .handle_range_query_promql(query, 1.0, 2.0, 1.0) + .expect("range query should execute locally") + .expect("range query should return a result"); + let QueryResult::Matrix(range) = range else { + panic!("range query should return a matrix"); + }; + assert_eq!(range.values.len(), 1); + assert_eq!(range.values[0].labels.labels, expected_labels); + assert_eq!( + range.values[0].samples.last().expect("range sample").value, + 5.0 + ); +} + #[tokio::test] async fn e2e_nested_aggregation_operator_matrix_executes_instant_and_range() { use asap_types::query_config::{ From 353d82868a76ab0042ffb9bf557d68333fd6891e Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 4 Oct 2026 18:43:03 -0400 Subject: [PATCH 21/22] fix(query-engine): keep grouped topk buckets contiguous --- .../src/engines/simple_engine/mod.rs | 146 +++++++++++++++++- 1 file changed, 143 insertions(+), 3 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 43e624fd..9299035f 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1517,6 +1517,53 @@ impl SimpleEngine { .then_with(|| a_labels.cmp(b_labels)) } + fn sort_instant_topk_results( + results: &mut [InstantVectorElement], + output_labels: &KeyByLabelNames, + grouping_labels: &KeyByLabelNames, + ) -> Result<(), QueryExecutionError> { + let grouping_positions = grouping_labels + .labels + .iter() + .map(|label| { + output_labels + .labels + .iter() + .position(|candidate| candidate == label) + .ok_or_else(|| { + QueryExecutionError::Native(format!( + "Topk grouping label '{label}' is absent from output" + )) + }) + }) + .collect::, _>>()?; + for result in results.iter() { + if grouping_positions + .iter() + .any(|position| result.labels.labels.get(*position).is_none()) + { + return Err(QueryExecutionError::Native( + "Topk result labels do not match the configured output schema".to_string(), + )); + } + } + results.sort_by(|left, right| { + for position in &grouping_positions { + let order = left.labels.labels[*position].cmp(&right.labels.labels[*position]); + if !order.is_eq() { + return order; + } + } + Self::cmp_topk_value_desc( + left.value, + &left.labels.labels, + right.value, + &right.labels.labels, + ) + }); + Ok(()) + } + /// Returns the required `k` parameter for a Topk query. /// /// PromQL context construction validates this before execution, but the @@ -1690,9 +1737,11 @@ impl SimpleEngine { // restore it. Tie-broken by label for determinism, matching the // range engine's own topk sort (#581 stage E.3). if context.metadata.statistic_to_compute == Statistic::Topk { - results.sort_by(|a, b| { - Self::cmp_topk_value_desc(a.value, &a.labels.labels, b.value, &b.labels.labels) - }); + Self::sort_instant_topk_results( + &mut results, + &context.metadata.query_output_labels, + &context.grouping_labels, + )?; } Ok(results) @@ -3096,6 +3145,9 @@ impl SimpleEngine { #[cfg(test)] mod topk_metadata_tests { use super::{QueryExecutionError, SimpleEngine}; + use crate::data_model::KeyByLabelValues; + use crate::engines::query_result::InstantVectorElement; + use promql_utilities::data_model::KeyByLabelNames; use std::collections::HashMap; #[test] @@ -3126,6 +3178,94 @@ mod topk_metadata_tests { )); assert!(matches!(native_error, Err(QueryExecutionError::Native(_)))); } + + #[test] + fn grouped_topk_keeps_each_bucket_contiguous() { + let output_labels = KeyByLabelNames::new(vec![ + "__name__".to_string(), + "instance".to_string(), + "job".to_string(), + ]); + let grouping_labels = KeyByLabelNames::new(vec!["job".to_string()]); + let mut results = vec![ + InstantVectorElement::new( + KeyByLabelValues::new_with_labels(vec![ + "ordered_data".to_string(), + "b".to_string(), + "backend".to_string(), + ]), + 2.0, + ), + InstantVectorElement::new( + KeyByLabelValues::new_with_labels(vec![ + "ordered_data".to_string(), + "a".to_string(), + "frontend".to_string(), + ]), + 4.0, + ), + InstantVectorElement::new( + KeyByLabelValues::new_with_labels(vec![ + "ordered_data".to_string(), + "a".to_string(), + "backend".to_string(), + ]), + 3.0, + ), + InstantVectorElement::new( + KeyByLabelValues::new_with_labels(vec![ + "ordered_data".to_string(), + "b".to_string(), + "frontend".to_string(), + ]), + 1.0, + ), + ]; + + SimpleEngine::sort_instant_topk_results(&mut results, &output_labels, &grouping_labels) + .unwrap(); + + assert_eq!( + results + .into_iter() + .map(|result| (result.labels.labels, result.value)) + .collect::>(), + vec![ + ( + vec![ + "ordered_data".to_string(), + "a".to_string(), + "backend".to_string(), + ], + 3.0, + ), + ( + vec![ + "ordered_data".to_string(), + "b".to_string(), + "backend".to_string(), + ], + 2.0, + ), + ( + vec![ + "ordered_data".to_string(), + "a".to_string(), + "frontend".to_string(), + ], + 4.0, + ), + ( + vec![ + "ordered_data".to_string(), + "b".to_string(), + "frontend".to_string(), + ], + 1.0, + ), + ] + ); + } } #[cfg(test)] From d450f5c224c1be7d8c0226a9fa47a46f29369628 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 5 Oct 2026 09:47:27 -0400 Subject: [PATCH 22/22] fix(query-engine): align raw topk labels --- .../src/engines/simple_engine/mod.rs | 14 ++- .../tests/e2e_precompute_equivalence.rs | 108 ++++++++++++++++++ .../suites/nested-aggregations.yaml | 6 + 3 files changed, 124 insertions(+), 4 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 9299035f..347d5d1a 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1737,11 +1737,17 @@ impl SimpleEngine { // restore it. Tie-broken by label for determinism, matching the // range engine's own topk sort (#581 stage E.3). if context.metadata.statistic_to_compute == Statistic::Topk { - Self::sort_instant_topk_results( - &mut results, - &context.metadata.query_output_labels, + let raw_labels = Self::topk_row_label_order( + &context.metadata, &context.grouping_labels, - )?; + &context.aggregated_labels, + ); + let result_labels = if enable_topk_formatting && context.metadata.keep_metric_name { + &context.metadata.query_output_labels + } else { + &raw_labels + }; + Self::sort_instant_topk_results(&mut results, result_labels, &context.grouping_labels)?; } Ok(results) diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 59641c57..5d6592a5 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -560,6 +560,114 @@ async fn e2e_nested_aggregation_preserves_plain_topk_labels() { ); } +#[tokio::test] +async fn e2e_nested_grouped_plain_topk_uses_raw_label_positions() { + use asap_types::query_config::{ + QueryTimeAggregationOperator, QueryTimeGrouping, QueryTimeGroupingMode, + }; + + let metric = "nested_grouped_plain_topk"; + let anchor = format!("topk by (job) (3, {metric})"); + let query = format!("sum by (job) ({anchor})"); + let mut config = make_agg_config_full( + 21, + metric, + AggregationType::CountMinSketchWithHeap, + "count", + 1_000, + 0, + vec!["job"], + vec!["instance"], + ); + config.parameters.insert("depth".to_string(), json!(3_u64)); + config + .parameters + .insert("width".to_string(), json!(128_u64)); + config + .parameters + .insert("heapsize".to_string(), json!(16_u64)); + let samples = [ + ("backend", "a", 3), + ("backend", "b", 2), + ("backend", "c", 1), + ("backend", "d", 1), + ] + .into_iter() + .flat_map(|(job, instance, count)| { + std::iter::repeat_with(move || { + make_timeseries( + metric, + vec![("job", job), ("instance", instance)], + 1_000, + 1.0, + ) + }) + .take(count) + }) + .chain(std::iter::once(make_timeseries(metric, vec![], 3_000, 0.0))) + .collect(); + let scenario = NativeDagScenario { + port: 19427, + metric, + query: &query, + aggregation_configs: vec![config], + schema_labels: vec!["instance".to_string(), "job".to_string()], + samples, + evaluation_time_seconds: 1.0, + base_interval_ms: 1_000, + }; + let (engine, query) = scenario + .build_engine_with_plan( + &anchor, + vec![QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Sum, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::By, + labels: vec!["job".to_string()], + }, + parameter: None, + }], + ) + .await; + + let (_, instant) = engine + .handle_query_promql(query.clone(), 1.0) + .expect("instant query should execute locally") + .expect("instant query should return a result"); + let QueryResult::Vector(instant) = instant else { + panic!("instant query should return a vector"); + }; + assert_eq!( + instant + .values + .into_iter() + .map(|value| (value.labels.labels, value.value)) + .collect::>(), + vec![(vec!["backend".to_string()], 6.0)] + ); + + let (_, range) = engine + .handle_range_query_promql(query, 1.0, 2.0, 1.0) + .expect("range query should execute locally") + .expect("range query should return a result"); + let QueryResult::Matrix(range) = range else { + panic!("range query should return a matrix"); + }; + assert_eq!( + range + .values + .into_iter() + .map(|value| { + ( + value.labels.labels, + value.samples.last().expect("range sample").value, + ) + }) + .collect::>(), + vec![(vec!["backend".to_string()], 6.0)] + ); +} + #[tokio::test] async fn e2e_nested_topk_preserves_plain_topk_metric_name() { use asap_types::query_config::{ diff --git a/promql-compliance/suites/nested-aggregations.yaml b/promql-compliance/suites/nested-aggregations.yaml index 736ad2f0..21fd0317 100644 --- a/promql-compliance/suites/nested-aggregations.yaml +++ b/promql-compliance/suites/nested-aggregations.yaml @@ -35,6 +35,12 @@ queries: expr: "topk by (job) (3, sum by (job, instance) (data))" instant_offsets_seconds: *evaluation_offsets range: *evaluation_range + # A bare-selector TopK keeps __name__ until final formatting. The outer + # aggregation makes the native path execute that raw TopK result first. + - name: bare-topk-by-job-then-sum + expr: "sum by (job) (topk by (job) (3, data))" + instant_offsets_seconds: *evaluation_offsets + range: *evaluation_range - name: one-stage-quantile-0-5 expr: "quantile by (job) (0.5, sum by (job, instance) (data))" instant_offsets_seconds: *evaluation_offsets