diff --git a/asap-common/dependencies/rs/asap_types/src/query_requirements.rs b/asap-common/dependencies/rs/asap_types/src/query_requirements.rs index 57e2214a..21dce81e 100644 --- a/asap-common/dependencies/rs/asap_types/src/query_requirements.rs +++ b/asap-common/dependencies/rs/asap_types/src/query_requirements.rs @@ -41,11 +41,12 @@ pub struct QueryRequirements { /// For `Statistic::Topk` requirements with an explicit `by`/`without` /// modifier, the labels used to bucket the input for independent /// per-bucket ranking (see `get_topk_by_labels`) -- e.g. `Some(["job"])` - /// for `topk by (job) (k, x)`. `None` for a bare `topk(k, x)` (single - /// global ranking) and for every non-topk requirement. Deliberately + /// for `topk by (job) (k, x)`. `None` when topk ranks its whole input as + /// one bucket (e.g. a bare `topk(k, x)`) and for every non-topk + /// requirement. Deliberately /// separate from `grouping_labels`, which for topk always reports the - /// *output* label set (all labels, unaffected by any modifier) rather - /// than the bucketing labels (#714). + /// *output* label set (all labels, or an inner `sum by (...)`'s labels; + /// unaffected by topk's own modifier) rather than the bucketing labels. pub topk_by_labels: Option, } @@ -97,8 +98,9 @@ pub fn build_query_requirements_promql( let topk_by_labels = get_topk_by_labels(match_result, &all_labels); let grouping_labels = if has_aggregation { - // OnlySpatial and (collapsable, see #508) OneTemporalOneSpatial encode - // their output labels in the AST's `by (...)` / `without (...)` clause. + // OnlySpatial, collapsable OneTemporalOneSpatial and topk-of-sum + // encode their output labels in the AST's `by (...)` / `without (...)` + // clause. get_spatial_aggregation_output_labels(match_result, &all_labels) } else { // OnlyTemporal preserves all labels. diff --git a/asap-common/dependencies/rs/promql_utilities/src/ast_matching/collapsable_patterns.rs b/asap-common/dependencies/rs/promql_utilities/src/ast_matching/collapsable_patterns.rs new file mode 100644 index 00000000..e496a99d --- /dev/null +++ b/asap-common/dependencies/rs/promql_utilities/src/ast_matching/collapsable_patterns.rs @@ -0,0 +1,106 @@ +//! Patterns for spatial aggregations over a temporal function that collapse +//! into one precomputed statistic. The planner and the query engine both +//! build their pattern lists from these, so they accept the same shapes. + +use std::collections::HashMap; + +use serde_json::Value; + +use crate::ast_matching::{PromQLPattern, PromQLPatternBuilder}; +use crate::query_logics::enums::{AggregationOperator, PromQLFunction}; +use crate::query_logics::logics::get_is_collapsable; + +/// Token for the inner reducer of `topk(k, sum by (...) (f(m[w])))`. +pub const INNER_AGGREGATION_TOKEN: &str = "inner_aggregation"; + +const TEMPORAL_FUNCTIONS: [PromQLFunction; 8] = [ + PromQLFunction::Rate, + PromQLFunction::Increase, + PromQLFunction::SumOverTime, + PromQLFunction::CountOverTime, + PromQLFunction::AvgOverTime, + PromQLFunction::MinOverTime, + PromQLFunction::MaxOverTime, + PromQLFunction::QuantileOverTime, +]; + +const SPATIAL_OPERATORS: [AggregationOperator; 7] = [ + AggregationOperator::Sum, + AggregationOperator::Count, + AggregationOperator::Avg, + AggregationOperator::Quantile, + AggregationOperator::Min, + AggregationOperator::Max, + AggregationOperator::Topk, +]; + +/// `func(metric[w])`, binding `metric`, `range_vector`, `function` and +/// `function_args`. +fn temporal_function(func: PromQLFunction) -> Option> { + PromQLPatternBuilder::function( + vec![func.as_str()], + vec![PromQLPatternBuilder::matrix_selector( + PromQLPatternBuilder::metric(None, None, None, Some("metric")), + None, + Some("range_vector"), + )], + Some("function"), + Some("function_args"), + ) +} + +/// `op(func(metric[w]))`: one narrow pattern per collapsable (function, op) +/// pair (see `get_is_collapsable`). `sum(min_over_time(x[5m]))` can't be +/// served by one precomputed statistic the way `sum(sum_over_time(x[5m]))` +/// can, and a broad any-op-over-any-function pattern would match it and +/// silently drop the outer aggregation instead of rejecting the query. +pub fn build_one_temporal_one_spatial_patterns() -> Vec { + TEMPORAL_FUNCTIONS + .into_iter() + .flat_map(|func| { + SPATIAL_OPERATORS + .into_iter() + .filter(move |&op| get_is_collapsable(func, op)) + .map(move |op| { + PromQLPattern::new(PromQLPatternBuilder::aggregation( + vec![op.as_str()], + temporal_function(func), + None, + None, + None, + Some("aggregation"), + )) + }) + }) + .collect() +} + +/// `topk(k, sum by (l) (func(metric[w])))`, one per `func` collapsable under +/// both `sum` and `topk`: such a query collapses into one +/// CountMinSketchWithHeap keyed by `l`. +pub fn build_topk_of_sum_patterns() -> Vec { + TEMPORAL_FUNCTIONS + .into_iter() + .filter(|&func| { + get_is_collapsable(func, AggregationOperator::Sum) + && get_is_collapsable(func, AggregationOperator::Topk) + }) + .map(|func| { + PromQLPattern::new(PromQLPatternBuilder::aggregation( + vec![AggregationOperator::Topk.as_str()], + PromQLPatternBuilder::aggregation( + vec![AggregationOperator::Sum.as_str()], + temporal_function(func), + None, + None, + None, + Some(INNER_AGGREGATION_TOKEN), + ), + None, + None, + None, + Some("aggregation"), + )) + }) + .collect() +} diff --git a/asap-common/dependencies/rs/promql_utilities/src/ast_matching/mod.rs b/asap-common/dependencies/rs/promql_utilities/src/ast_matching/mod.rs index bcc58df9..30ea4c73 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/ast_matching/mod.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/ast_matching/mod.rs @@ -1,3 +1,4 @@ +pub mod collapsable_patterns; pub mod promql_pattern; pub mod promql_pattern_builder; //pub mod promql_pattern_factory; diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs index 606dab93..e522ff27 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs @@ -1,6 +1,7 @@ use promql_parser::parser::Expr; use tracing::debug; +use crate::ast_matching::collapsable_patterns::INNER_AGGREGATION_TOKEN; use crate::ast_matching::promql_pattern::AggregationModifierType; use crate::ast_matching::PromQLMatchResult; use crate::data_model::KeyByLabelNames; @@ -170,12 +171,36 @@ pub fn get_spatial_aggregation_output_labels( .and_then(|token| token.aggregation.as_ref()) .expect("aggregation token missing"); - // Patching: When the query is topk, we should always return all labels + // topk is a selector: it keeps its input's labels, which are all labels + // unless an inner `sum by (...)` reduced them first. if aggregation_token.op.parse::() == Ok(AggregationOperator::Topk) { - debug!("Aggregation operation is 'topk', returning all labels"); - return all_labels.clone(); + debug!("Aggregation operation is 'topk', returning its input's labels"); + return topk_input_labels(match_result, all_labels); } + reducer_output_labels(aggregation_token, all_labels) +} + +/// Labels of the series topk ranks: the inner reducer's output labels when +/// one is present, else all labels. +fn topk_input_labels( + match_result: &PromQLMatchResult, + all_labels: &KeyByLabelNames, +) -> KeyByLabelNames { + match match_result + .tokens + .get(INNER_AGGREGATION_TOKEN) + .and_then(|token| token.aggregation.as_ref()) + { + Some(inner) => reducer_output_labels(inner, all_labels), + None => all_labels.clone(), + } +} + +fn reducer_output_labels( + aggregation_token: &crate::AggregationToken, + all_labels: &KeyByLabelNames, +) -> KeyByLabelNames { // Fixing issue https://github.com/ProjectASAP/asap-internal/issues/24 let modifier: &crate::AggregationModifier = match aggregation_token.modifier.as_ref() { Some(m) => m, @@ -210,15 +235,17 @@ pub fn get_spatial_aggregation_output_labels( /// Per PromQL semantics, `topk`/`bottomk` are selectors, not reducers: `by`/ /// `without` only buckets the input series for independent per-bucket /// ranking, it never changes the output's label set (unlike `sum by (...)` -/// etc, which collapse the output down to the by-clause labels) -- that's -/// exactly why `get_spatial_aggregation_output_labels` always returns all -/// labels for `topk` regardless of any modifier. This function recovers the -/// modifier's labels separately, for callers that need the bucketing -/// information itself (#714: `topk by (job) (k, x)` needs one independent -/// top-k ranking per `job` value, not one ranking over the whole input). +/// etc, which collapse the output down to the by-clause labels). So +/// `get_spatial_aggregation_output_labels` returns topk's input labels +/// regardless of any modifier: all labels, or the inner `sum`'s output +/// labels in `topk(k, sum by (...) (...))`. This function recovers the +/// bucket labels separately, for callers that need them: `topk by (job) +/// (k, x)` ranks each `job` value independently. Bucket labels are always +/// within topk's input labels. /// -/// Returns `None` when there's no modifier (a bare `topk(k, x)` ranks across -/// the whole input in a single bucket) or when the aggregation isn't `topk`. +/// Returns `None` when the input forms a single bucket: no modifier (a bare +/// `topk(k, x)`), no bucket labels left within the input, or when the +/// aggregation isn't `topk`. pub fn get_topk_by_labels( match_result: &PromQLMatchResult, all_labels: &KeyByLabelNames, @@ -233,13 +260,34 @@ pub fn get_topk_by_labels( } let modifier = aggregation_token.modifier.as_ref()?; - Some(match modifier.modifier_type { - AggregationModifierType::By => KeyByLabelNames::new(modifier.labels.clone()), - AggregationModifierType::Without => { - let without_labels = KeyByLabelNames::new(modifier.labels.clone()); - all_labels.difference(&without_labels) - } - }) + let modifier_labels = KeyByLabelNames::new(modifier.labels.clone()); + let Some(inner) = match_result + .tokens + .get(INNER_AGGREGATION_TOKEN) + .and_then(|token| token.aggregation.as_ref()) + else { + // Plain topk: `all_labels` may be empty for a metric missing from + // the schema, so `by` labels are taken as written. + return Some(match modifier.modifier_type { + AggregationModifierType::By => modifier_labels, + AggregationModifierType::Without => all_labels.difference(&modifier_labels), + }); + }; + let input_labels = reducer_output_labels(inner, all_labels); + let buckets = match modifier.modifier_type { + // The inner sum removed every label outside its output, so such a + // `by` label is empty in every input series and doesn't split them. + AggregationModifierType::By => KeyByLabelNames::new( + modifier_labels + .labels + .into_iter() + .filter(|label| input_labels.labels.contains(label)) + .collect(), + ), + AggregationModifierType::Without => input_labels.difference(&modifier_labels), + }; + // No bucket labels left is one ranking over the whole input. + Some(buckets).filter(|buckets| !buckets.is_empty()) } #[cfg(test)] @@ -296,4 +344,102 @@ mod tests { assert_eq!(err, StatisticExtractionError::MissingStatistic); } + + fn labels(names: &[&str]) -> KeyByLabelNames { + KeyByLabelNames::new(names.iter().map(|n| n.to_string()).collect()) + } + + fn match_topk_of_sum(query: &str) -> PromQLMatchResult { + let ast = promql_parser::parser::parse(query).unwrap(); + let result = crate::ast_matching::collapsable_patterns::build_topk_of_sum_patterns() + .iter() + .map(|pattern| pattern.matches(&ast)) + .find(|result| result.matches); + result.unwrap_or_else(|| panic!("{query} should match")) + } + #[test] + fn topk_of_sum_outputs_the_inner_sum_labels() { + let all = labels(&["instance", "job", "label_0"]); + let result = match_topk_of_sum("topk(32, sum by (label_0) (sum_over_time(data[15m])))"); + + assert_eq!( + get_spatial_aggregation_output_labels(&result, &all), + labels(&["label_0"]) + ); + assert_eq!(get_topk_by_labels(&result, &all), None); + assert_eq!( + get_statistics_to_compute(&result), + Ok(vec![Statistic::Topk]) + ); + } + + #[test] + fn topk_of_sum_without_buckets_within_the_inner_sum_labels() { + let all = labels(&["instance", "job", "label_0"]); + let result = match_topk_of_sum( + "topk without (label_0) (3, sum by (job, label_0) (sum_over_time(data[5m])))", + ); + + assert_eq!( + get_spatial_aggregation_output_labels(&result, &all), + labels(&["job", "label_0"]) + ); + // Only `job` is left of the inner sum's labels, not `instance`. + assert_eq!(get_topk_by_labels(&result, &all), Some(labels(&["job"]))); + } + + /// A metric missing from the schema has no known labels; plain topk's + /// `by` labels must still bucket, not collapse into a global ranking. + #[test] + fn plain_topk_by_keeps_its_labels_for_a_metric_missing_from_the_schema() { + let ast = promql_parser::parser::parse("topk by (job) (5, x)").unwrap(); + let result = crate::ast_matching::PromQLPattern::new( + crate::ast_matching::PromQLPatternBuilder::aggregation( + vec!["topk"], + crate::ast_matching::PromQLPatternBuilder::metric(None, None, None, Some("metric")), + None, + None, + None, + Some("aggregation"), + ), + ) + .matches(&ast); + assert!(result.matches); + + assert_eq!( + get_topk_by_labels(&result, &KeyByLabelNames::empty()), + Some(labels(&["job"])) + ); + } + + /// The inner sum already removed `instance`, so every input series has + /// the same (empty) `instance` and there is one bucket: the global top + /// jobs, not a heap per raw `instance` of partial sums. + #[test] + fn topk_of_sum_by_a_label_the_inner_sum_removed_is_one_bucket() { + let all = labels(&["instance", "job"]); + let result = + match_topk_of_sum("topk by (instance) (3, sum by (job) (sum_over_time(m[5m])))"); + + assert_eq!(get_topk_by_labels(&result, &all), None); + + let partly_kept = match_topk_of_sum( + "topk by (instance, job) (3, sum by (job, label_0) (sum_over_time(m[5m])))", + ); + assert_eq!( + get_topk_by_labels(&partly_kept, &labels(&["instance", "job", "label_0"])), + Some(labels(&["job"])) + ); + } + + #[test] + fn topk_of_sum_without_inner_modifier_collapses_to_no_labels() { + let all = labels(&["instance", "job"]); + let result = match_topk_of_sum("topk(1, sum(sum_over_time(data[5m])))"); + + assert_eq!( + get_spatial_aggregation_output_labels(&result, &all), + KeyByLabelNames::empty() + ); + } } diff --git a/asap-planner-rs/src/optimizer/milp.rs b/asap-planner-rs/src/optimizer/milp.rs index 655a92bd..07165935 100644 --- a/asap-planner-rs/src/optimizer/milp.rs +++ b/asap-planner-rs/src/optimizer/milp.rs @@ -486,6 +486,11 @@ fn item_to_raqe(item: &OptimizerItem) -> Result { // TopK keeps one heap per `topk by` bucket; its `grouping_labels` is the // output label set (every label), which would cost one heap per series. + // ponytail: a topk-of-sum heap holds one key per inner `sum by` group, but + // the Raqe has no field for those labels, so it is priced from the + // bucket grouping's shape. Key count is bounded by the series count, but + // accuracy curves measured per series may not match summed values; needs + // key labels in sketch-bench's Raqe. let grouping_labels: LabelSet = match capability { Capability::TopKByValue | Capability::TopKByCount => req .topk_by_labels @@ -754,6 +759,22 @@ metrics: assert_eq!(r.grouping_labels, labels(&["job"])); } + #[test] + fn topk_of_sum_is_one_global_heap_keyed_by_the_inner_labels() { + let w = workload(&group( + "topk(5, sum by (job) (sum_over_time(http_requests_total[5m])))", + 0.9, + )); + let r = &w.raqes[0]; + assert_eq!(r.capability, Capability::TopKByValue); + assert_eq!(r.grouping_labels, LabelSet::new()); + assert_eq!(r.topk_k, Some(5)); + assert_eq!( + w.items[0].requirements.grouping_labels, + promql_utilities::data_model::KeyByLabelNames::new(vec!["job".into()]) + ); + } + #[test] fn bare_topk_is_one_global_heap() { let w = workload(&group("topk(5, http_requests_total)", 0.9)); diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index e45e2733..1b282d6b 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -1,6 +1,6 @@ //! Turns a solved MILP plan into the planner's streaming and inference YAML. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use asap_types::enums::{CleanupPolicy, WindowType}; use indexmap::map::Entry; @@ -133,16 +133,6 @@ pub fn deployments_to_planner_output( } } - let item_of = |raqe: usize| &workload.items[workload.raqe_items[raqe]]; - // Each item once per deployment, however many occurrences it has. - let mut served: Vec> = vec![Vec::new(); deployments.len()]; - let mut seen = HashSet::new(); - for (raqe, &index) in served_by.iter().enumerate() { - if seen.insert((index, workload.raqe_items[raqe])) { - served[index].push(item_of(raqe)); - } - } - // A SQL aggregation names its table and value column, and its metric is // the table. let sql_columns: HashMap = config @@ -151,10 +141,20 @@ pub fn deployments_to_planner_output( .map(|(metric, table, column)| (metric, (table.name.as_str(), column))) .collect(); + // One aggregation per (deployment, labels its queries need). A Raqe + // doesn't carry a keyed sketch's key labels (a topk-of-sum heap's inner + // `sum by` labels), so the solver may give one deployment to queries + // needing different keys; those get one aggregation each. + // ponytail: the plan then prices one sketch where several are deployed; + // key labels in the Raqe would let the solver see them. let mut aggregations: IndexMap = IndexMap::new(); - let mut deployment_keys = Vec::with_capacity(served.len()); - for (deployment, items) in deployments.iter().zip(&served) { - let mut aggregation = aggregation_config(deployment, items)?; + let mut item_keys: HashMap<(usize, usize), String> = HashMap::new(); + for (raqe, &index) in served_by.iter().enumerate() { + let item = workload.raqe_items[raqe]; + if item_keys.contains_key(&(index, item)) { + continue; + } + let mut aggregation = aggregation_config(&deployments[index], &workload.items[item])?; if let Some(&(table, column)) = sql_columns.get(&aggregation.metric) { aggregation.metric = table.to_string(); aggregation.table_name = Some(table.to_string()); @@ -162,13 +162,14 @@ pub fn deployments_to_planner_output( } let key = aggregation.identifying_key(); aggregations.entry(key.clone()).or_insert(aggregation); - deployment_keys.push(key); + item_keys.insert((index, item), key); } let mut queries: IndexMap = IndexMap::new(); for (raqe, &index) in served_by.iter().enumerate() { - let key = &deployment_keys[index]; - for query in &item_of(raqe).query_strings { + let item = workload.raqe_items[raqe]; + let key = &item_keys[&(index, item)]; + for query in &workload.items[item].query_strings { match queries.entry(query.clone()) { Entry::Vacant(entry) => { entry.insert(QueryPlanEntry::fully_planned( @@ -211,12 +212,12 @@ pub fn deployments_to_planner_output( })) } -/// `items` are the optimizer items of the Raqes `deployment` serves; they -/// share its capability and grouping (`plan_to_planner_output` refuses -/// roll-ups), so the first one decides the labels. +/// The aggregation `deployment` runs for `item`, one of the optimizer items +/// it serves. Items share the deployment's capability and grouping +/// (`plan_to_planner_output` refuses roll-ups); `item` decides the labels. fn aggregation_config( deployment: &Deployment, - items: &[&OptimizerItem], + item: &OptimizerItem, ) -> Result { let variant = deployment.config.sketch.as_str(); let integer_param = |name: &'static str| { @@ -282,21 +283,7 @@ fn aggregation_config( } }; - let requirements = &items - .first() - .expect("every planned deployment serves a Raqe") - .requirements; - let mut grouping_labels = KeyByLabelNames::empty(); - let mut aggregated_labels = KeyByLabelNames::empty(); - set_subpopulation_labels( - requirements.statistics[0], - aggregation_type, - &requirements.grouping_labels, - requirements.topk_by_labels.as_ref(), - &mut KeyByLabelNames::empty(), - &mut grouping_labels, - &mut aggregated_labels, - ); + let (grouping_labels, aggregated_labels) = subpopulation_labels(item, aggregation_type); Ok(IntermediateAggConfig { aggregation_type, @@ -322,6 +309,27 @@ fn aggregation_config( }) } +/// The (grouping, aggregated) labels `item`'s query needs from an aggregation +/// of `aggregation_type`. +fn subpopulation_labels( + item: &OptimizerItem, + aggregation_type: AggregationType, +) -> (KeyByLabelNames, KeyByLabelNames) { + let requirements = &item.requirements; + let mut grouping_labels = KeyByLabelNames::empty(); + let mut aggregated_labels = KeyByLabelNames::empty(); + set_subpopulation_labels( + requirements.statistics[0], + aggregation_type, + &requirements.grouping_labels, + requirements.topk_by_labels.as_ref(), + &mut KeyByLabelNames::empty(), + &mut grouping_labels, + &mut aggregated_labels, + ); + (grouping_labels, aggregated_labels) +} + fn missing_param(variant: &str, param: &'static str) -> MilpOutputError { MilpOutputError::MissingParam { variant: variant.to_string(), @@ -835,6 +843,34 @@ mod tests { ); } + /// Both are global top-k Raqes on one metric, so the solver may give them + /// one deployment, but their heaps are keyed by different labels: each + /// query gets its own heap. + #[test] + fn topk_of_sum_queries_needing_different_heap_keys_get_separate_heaps() { + let by_job = "topk(5, sum by (job) (sum_over_time(http_requests_total[1m])))"; + let by_instance = "topk(5, sum by (instance) (sum_over_time(http_requests_total[1m])))"; + let (config, workload, mut solution) = + plan(&(group(by_job, 0.9) + &group(by_instance, 0.9))); + for raqe in &mut solution.raqes { + raqe.deployment = 0; + } + solution.deployments.truncate(1); + + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); + + assert_eq!(output.streaming_aggregation_count(), 2); + for (query, key) in [(by_job, "job"), (by_instance, "instance")] { + let aggregation = aggregation_for(&output, query); + assert_eq!( + aggregation.aggregated_labels, + KeyByLabelNames::new(vec![key.to_string()]), + "{query}" + ); + } + } + #[test] fn query_split_across_deployments_is_an_error() { // The same query string at two cadences is two items; if they land on diff --git a/asap-planner-rs/src/planner/patterns.rs b/asap-planner-rs/src/planner/patterns.rs index 6790a1f1..285753ba 100644 --- a/asap-planner-rs/src/planner/patterns.rs +++ b/asap-planner-rs/src/planner/patterns.rs @@ -1,10 +1,13 @@ +use promql_utilities::ast_matching::collapsable_patterns::{ + build_one_temporal_one_spatial_patterns, build_topk_of_sum_patterns, +}; use promql_utilities::ast_matching::{PromQLPattern, PromQLPatternBuilder}; use promql_utilities::query_logics::enums::{AggregationOperator, PromQLFunction}; -use promql_utilities::query_logics::logics::get_is_collapsable; /// Build all patterns in priority order: ONLY_TEMPORAL (2), ONLY_SPATIAL (1), -/// ONE_TEMPORAL_ONE_SPATIAL (one per collapsable (function, op) pair). Tried -/// in order until one matches. +/// ONE_TEMPORAL_ONE_SPATIAL (one per collapsable (function, op) pair), +/// TOPK_OF_SUM (one per function collapsable under both `sum` and `topk`). +/// Tried in order until one matches. pub fn build_patterns() -> Vec { let metric_pattern = || PromQLPatternBuilder::metric(None, None, None, Some("metric")); let range_vector_pattern = @@ -65,58 +68,9 @@ pub fn build_patterns() -> Vec { Some("aggregation"), )); - // ONE_TEMPORAL_ONE_SPATIAL: one narrow pattern per collapsable (function, op) - // pair (see `get_is_collapsable`) — e.g. `sum(min_over_time(x[5m]))` cannot be - // served by a single precomputed statistic the way `sum(sum_over_time(x[5m]))` - // can. A broad any-op-wrapping-any-function pattern would structurally match - // non-collapsable combinations too, silently dropping the outer aggregation - // instead of rejecting the query (see #508). - let one_temporal_one_spatial_collapsable: Vec = [ - PromQLFunction::Rate, - PromQLFunction::Increase, - PromQLFunction::SumOverTime, - PromQLFunction::CountOverTime, - PromQLFunction::AvgOverTime, - PromQLFunction::MinOverTime, - PromQLFunction::MaxOverTime, - PromQLFunction::QuantileOverTime, - ] - .into_iter() - .flat_map(|func| { - [ - AggregationOperator::Sum, - AggregationOperator::Count, - AggregationOperator::Avg, - AggregationOperator::Quantile, - AggregationOperator::Min, - AggregationOperator::Max, - AggregationOperator::Topk, - ] - .into_iter() - .filter_map(move |op| { - if !get_is_collapsable(func, op) { - return None; - } - let pattern = PromQLPatternBuilder::aggregation( - vec![op.as_str()], - PromQLPatternBuilder::function( - vec![func.as_str()], - vec![range_vector_pattern()], - Some("function"), - Some("function_args"), - ), - None, - None, - None, - Some("aggregation"), - ); - Some(PromQLPattern::new(pattern)) - }) - }) - .collect(); - let mut patterns = vec![ot_quantile, ot_temporal_funcs, os_spatial]; - patterns.extend(one_temporal_one_spatial_collapsable); + patterns.extend(build_one_temporal_one_spatial_patterns()); + patterns.extend(build_topk_of_sum_patterns()); patterns } @@ -135,8 +89,30 @@ mod tests { fn exactly_six_collapsable_one_temporal_one_spatial_patterns() { // 2 ONLY_TEMPORAL + 1 ONLY_SPATIAL + 6 collapsable ONE_TEMPORAL_ONE_SPATIAL // (sum+sum_over_time, sum+count_over_time, min+min_over_time, - // max+max_over_time, topk+sum_over_time, topk+count_over_time). - assert_eq!(build_patterns().len(), 9); + // max+max_over_time, topk+sum_over_time, topk+count_over_time) + // + 2 TOPK_OF_SUM (sum_over_time, count_over_time). + assert_eq!(build_patterns().len(), 11); + } + + #[test] + fn topk_of_sum_matches_only_linear_functions() { + assert!(matches_some_pattern( + "topk(32, sum by (label_0) (sum_over_time(data[15m])))" + )); + assert!(matches_some_pattern( + "topk(3, sum by (job) (count_over_time(x[5m])))" + )); + // The counter shapes need per-series deltas before grouping. + assert!(!matches_some_pattern( + "topk(32, sum by (label_0) (rate(data_total[15m])))" + )); + assert!(!matches_some_pattern( + "sum by (job) (rate(data_total[15m]))" + )); + // A CMS sums, so only `sum` can sit between topk and the function. + assert!(!matches_some_pattern( + "topk(3, max by (job) (sum_over_time(x[5m])))" + )); } #[test] diff --git a/asap-planner-rs/tests/integration.rs b/asap-planner-rs/tests/integration.rs index f193a1db..02da9b14 100644 --- a/asap-planner-rs/tests/integration.rs +++ b/asap-planner-rs/tests/integration.rs @@ -746,6 +746,47 @@ query_groups: ); } +#[test] +fn topk_of_sum_collapses_into_one_heap_keyed_by_the_inner_labels() { + let query = "topk(3, sum by (job) (sum_over_time(http_requests_total[5m])))"; + let c = Controller::from_yaml_with_schema( + &format!( + r#" +query_groups: + - id: 1 + queries: + - "{query}" + repetition_delay_ms: 60000 +"# + ), + http_requests_schema(), + default_opts(), + ) + .unwrap(); + let out = c.generate().unwrap(); + + assert_eq!(out.streaming_aggregation_count(), 1); + assert!(out.has_aggregation_type_and_sub_type("CountMinSketchWithHeap", "sum")); + assert_eq!( + out.aggregation_labels("CountMinSketchWithHeap", "aggregated"), + vec!["job".to_string()], + "the heap ranks per-job sums, not per-series values" + ); + assert!(out + .aggregation_labels("CountMinSketchWithHeap", "grouping") + .is_empty()); + + let inference: serde_yaml::Value = + serde_yaml::from_str(&out.to_inference_yaml_string().unwrap()).unwrap(); + let planned_query = &inference["queries"][0]; + assert_eq!(planned_query["planned_subquery"].as_str(), Some(query)); + assert_eq!( + planned_query["query_time_aggregations"].as_sequence(), + Some(&vec![]), + "the heap answers topk itself; no query-time stage runs" + ); +} + #[test] fn heap_parameters_require_explicit_count_events_weighting() { let error = asap_planner::planner::sketch::build_sketch_parameters( diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 2fcc3894..b66fec41 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -29,9 +29,11 @@ use crate::AggregateCore; use asap_types::capability_matching::key_agg_compatible_with_value; use asap_types::enums::WindowType; +use promql_utilities::ast_matching::collapsable_patterns::{ + build_one_temporal_one_spatial_patterns, build_topk_of_sum_patterns, +}; use promql_utilities::ast_matching::{PromQLPattern, PromQLPatternBuilder}; use promql_utilities::data_model::KeyByLabelNames; -use promql_utilities::get_is_collapsable; use promql_utilities::query_logics::enums::{ AggregationOperator, AggregationType, PromQLFunction, Statistic, }; @@ -507,71 +509,14 @@ impl SimpleEngine { PromQLPattern::new(blocks[pattern_type].clone()) } - // A spatial aggregation wrapping a temporal function only collapses into a - // single equivalent statistic for specific (function, op) pairs (see - // `get_is_collapsable`) — e.g. `sum(min_over_time(x[5m]))` cannot be served - // by a single precomputed statistic the way `sum(sum_over_time(x[5m]))` can. - // Building one narrow pattern per collapsable pair (rather than a broad - // any-op-wrapping-any-function pattern) means a non-collapsable combination - // never structurally matches at all, instead of matching and then silently - // dropping the outer aggregation (see #508). - let one_temporal_one_spatial_collapsable_patterns: Vec = [ - PromQLFunction::Rate, - PromQLFunction::Increase, - PromQLFunction::SumOverTime, - PromQLFunction::CountOverTime, - PromQLFunction::AvgOverTime, - PromQLFunction::MinOverTime, - PromQLFunction::MaxOverTime, - PromQLFunction::QuantileOverTime, - ] - .into_iter() - .flat_map(|func| { - [ - AggregationOperator::Sum, - AggregationOperator::Count, - AggregationOperator::Avg, - AggregationOperator::Quantile, - AggregationOperator::Min, - AggregationOperator::Max, - AggregationOperator::Topk, - ] - .into_iter() - .filter_map(move |op| { - if !get_is_collapsable(func, op) { - return None; - } - let range_vector = PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - ); - let function_pattern = PromQLPatternBuilder::function( - vec![func.as_str()], - vec![range_vector], - Some("function"), - Some("function_args"), - ); - let pattern = PromQLPatternBuilder::aggregation( - vec![op.as_str()], - function_pattern, - None, - None, - None, - Some("aggregation"), - ); - Some(PromQLPattern::new(pattern)) - }) - }) - .collect(); - // Create controller patterns: tried in order until one matches. let mut controller_patterns = vec![ temporal_pattern("quantile", &temporal_pattern_blocks), temporal_pattern("generic", &temporal_pattern_blocks), spatial_pattern("generic", &spatial_pattern_blocks), ]; - controller_patterns.extend(one_temporal_one_spatial_collapsable_patterns); + controller_patterns.extend(build_one_temporal_one_spatial_patterns()); + controller_patterns.extend(build_topk_of_sum_patterns()); Self { store, diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 0ecc764d..22ddc621 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -188,10 +188,9 @@ impl SimpleEngine { .map(|s| s.to_string()) } - /// Extracts topk k parameter from PromQL match result. Topk is only ever - /// produced by a spatial `topk` aggregation (excluded from every - /// collapsable spatial-of-temporal pattern), so no shape check is needed - /// before looking at the aggregation token. + /// Extracts topk k parameter from PromQL match result. Every pattern that + /// yields Topk binds its outer `topk` to the `aggregation` token (an inner + /// `sum` binds to `inner_aggregation`), so k is always that token's param. fn extract_topk_param(&self, match_result: &PromQLMatchResult) -> Result { match_result .tokens diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index ee1be666..2a7ec9a1 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -454,6 +454,61 @@ async fn e2e_nested_topk_executes_after_its_planned_sum_anchor() { assert_eq!(vector.values[0].value, 9.0); } +#[tokio::test] +async fn e2e_topk_of_sum_ranks_groups_from_one_planned_heap() { + let metric = "topk_of_sum_requests"; + let query = "topk(1, sum by (job) (sum_over_time(topk_of_sum_requests[1s])))"; + let (streaming_config, inference_config) = plan_promql_query( + metric, + vec!["instance".to_string(), "job".to_string()], + query, + 1_000, + ); + let heaps: Vec<_> = streaming_config + .get_all_aggregation_configs() + .values() + .map(|config| (config.aggregation_type, config.aggregated_labels.clone())) + .collect(); + assert_eq!( + heaps, + vec![( + AggregationType::CountMinSketchWithHeap, + KeyByLabelNames::new(vec!["job".to_string()]) + )] + ); + + // The top series is worker's `c` (4), but the top job is api (2 + 3). + let 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, + 4.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, + ), + ]; + let engine = + build_engine_from_configs(19429, streaming_config, inference_config, samples, 1_000).await; + + let Some((_, QueryResult::Vector(vector))) = engine + .handle_query_promql(query.to_string(), 1.0) + .expect("topk of sum should execute natively") + else { + panic!("topk of sum should return an instant vector"); + }; + assert_eq!(vector.values.len(), 1, "result: {:?}", vector.values); + assert_eq!(vector.values[0].labels.labels, vec!["api"]); + assert_eq!(vector.values[0].value, 5.0); +} + #[tokio::test] async fn e2e_nested_aggregation_preserves_plain_topk_labels() { use asap_types::query_config::{