From d24b6a1e9705c09ebadfb4f3b94e6f5ee7e35f1a Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 22:41:48 -0400 Subject: [PATCH 1/9] feat(promql-utils): derive topk-of-sum output labels from the inner sum Co-Authored-By: Claude Opus 5.5 --- .../rs/asap_types/src/query_requirements.rs | 10 +- .../src/query_logics/parsing.rs | 112 +++++++++++++++++- 2 files changed, 114 insertions(+), 8 deletions(-) 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..82ea0694 100644 --- a/asap-common/dependencies/rs/asap_types/src/query_requirements.rs +++ b/asap-common/dependencies/rs/asap_types/src/query_requirements.rs @@ -44,8 +44,9 @@ pub struct QueryRequirements { /// for `topk by (job) (k, x)`. `None` for a bare `topk(k, x)` (single /// global ranking) 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 + /// (#714). 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, see #508) 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/query_logics/parsing.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs index 606dab93..da0f3d7d 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 @@ -170,12 +170,39 @@ 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) +} + +/// Token for the inner reducer of `topk(k, sum by (...) (f(m[w])))`. +pub const INNER_AGGREGATION_TOKEN: &str = "inner_aggregation"; + +/// 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, @@ -237,7 +264,7 @@ pub fn get_topk_by_labels( AggregationModifierType::By => KeyByLabelNames::new(modifier.labels.clone()), AggregationModifierType::Without => { let without_labels = KeyByLabelNames::new(modifier.labels.clone()); - all_labels.difference(&without_labels) + topk_input_labels(match_result, all_labels).difference(&without_labels) } }) } @@ -296,4 +323,81 @@ 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 { + use crate::ast_matching::{PromQLPattern, PromQLPatternBuilder}; + let pattern = PromQLPattern::new(PromQLPatternBuilder::aggregation( + vec!["topk"], + PromQLPatternBuilder::aggregation( + vec!["sum"], + PromQLPatternBuilder::function( + vec!["sum_over_time"], + vec![PromQLPatternBuilder::matrix_selector( + PromQLPatternBuilder::metric(None, None, None, Some("metric")), + None, + Some("range_vector"), + )], + Some("function"), + Some("function_args"), + ), + None, + None, + None, + Some(INNER_AGGREGATION_TOKEN), + ), + None, + None, + None, + Some("aggregation"), + )); + let result = pattern.matches(&promql_parser::parser::parse(query).unwrap()); + assert!(result.matches, "{query} should match"); + result + } + + #[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"]))); + } + + #[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() + ); + } } From cee0a2e1e00e21dadd6fbfa168e33f3d91830ff6 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 22:48:31 -0400 Subject: [PATCH 2/9] feat(planner): collapse topk(sum by (l) (f(m[w]))) into one CountMinSketchWithHeap For f linear under both sum and topk (sum_over_time, count_over_time), plan one heap keyed by the inner sum's labels instead of an exact sum anchor with a query-time topk. Both the legacy and MILP planners use it. Co-Authored-By: Claude Opus 5.5 --- .../src/query_logics/parsing.rs | 84 ++++++++++++------- asap-planner-rs/src/optimizer/milp.rs | 16 ++++ asap-planner-rs/src/planner/patterns.rs | 33 +++++++- asap-planner-rs/tests/integration.rs | 41 +++++++++ 4 files changed, 140 insertions(+), 34 deletions(-) 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 da0f3d7d..acd0e460 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 @@ -2,7 +2,7 @@ use promql_parser::parser::Expr; use tracing::debug; use crate::ast_matching::promql_pattern::AggregationModifierType; -use crate::ast_matching::PromQLMatchResult; +use crate::ast_matching::{PromQLMatchResult, PromQLPattern, PromQLPatternBuilder}; use crate::data_model::KeyByLabelNames; use crate::query_logics::enums::{AggregationOperator, PromQLFunction, Statistic}; use crate::query_logics::logics::get_is_collapsable; @@ -183,6 +183,53 @@ pub fn get_spatial_aggregation_output_labels( /// Token for the inner reducer of `topk(k, sum by (...) (f(m[w])))`. pub const INNER_AGGREGATION_TOKEN: &str = "inner_aggregation"; +/// Patterns for `topk(k, sum by (l) (f(m[w])))`, one per `f` collapsable +/// under both `sum` and `topk`: such a query collapses into one +/// CountMinSketchWithHeap keyed by `l`. +pub fn build_topk_of_sum_patterns() -> Vec { + [ + PromQLFunction::Rate, + PromQLFunction::Increase, + PromQLFunction::SumOverTime, + PromQLFunction::CountOverTime, + PromQLFunction::AvgOverTime, + PromQLFunction::MinOverTime, + PromQLFunction::MaxOverTime, + ] + .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()], + 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"), + ), + None, + None, + None, + Some(INNER_AGGREGATION_TOKEN), + ), + None, + None, + None, + Some("aggregation"), + )) + }) + .collect() +} + /// Labels of the series topk ranks: the inner reducer's output labels when /// one is present, else all labels. fn topk_input_labels( @@ -329,36 +376,13 @@ mod tests { } fn match_topk_of_sum(query: &str) -> PromQLMatchResult { - use crate::ast_matching::{PromQLPattern, PromQLPatternBuilder}; - let pattern = PromQLPattern::new(PromQLPatternBuilder::aggregation( - vec!["topk"], - PromQLPatternBuilder::aggregation( - vec!["sum"], - PromQLPatternBuilder::function( - vec!["sum_over_time"], - vec![PromQLPatternBuilder::matrix_selector( - PromQLPatternBuilder::metric(None, None, None, Some("metric")), - None, - Some("range_vector"), - )], - Some("function"), - Some("function_args"), - ), - None, - None, - None, - Some(INNER_AGGREGATION_TOKEN), - ), - None, - None, - None, - Some("aggregation"), - )); - let result = pattern.matches(&promql_parser::parser::parse(query).unwrap()); - assert!(result.matches, "{query} should match"); - result + let ast = promql_parser::parser::parse(query).unwrap(); + let result = 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"]); diff --git a/asap-planner-rs/src/optimizer/milp.rs b/asap-planner-rs/src/optimizer/milp.rs index 71d8d83c..18401e25 100644 --- a/asap-planner-rs/src/optimizer/milp.rs +++ b/asap-planner-rs/src/optimizer/milp.rs @@ -664,6 +664,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/planner/patterns.rs b/asap-planner-rs/src/planner/patterns.rs index 6790a1f1..404a4750 100644 --- a/asap-planner-rs/src/planner/patterns.rs +++ b/asap-planner-rs/src/planner/patterns.rs @@ -1,10 +1,12 @@ use promql_utilities::ast_matching::{PromQLPattern, PromQLPatternBuilder}; use promql_utilities::query_logics::enums::{AggregationOperator, PromQLFunction}; use promql_utilities::query_logics::logics::get_is_collapsable; +use promql_utilities::query_logics::parsing::build_topk_of_sum_patterns; /// 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 = @@ -117,6 +119,7 @@ pub fn build_patterns() -> Vec { let mut patterns = vec![ot_quantile, ot_temporal_funcs, os_spatial]; patterns.extend(one_temporal_one_spatial_collapsable); + patterns.extend(build_topk_of_sum_patterns()); patterns } @@ -135,8 +138,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( From 06b883c19ad0c9ef5f0bf9d043da8774a5d89076 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Mon, 12 Oct 2026 22:54:00 -0400 Subject: [PATCH 3/9] feat(query-engine): serve topk(sum by (l) (f(m[w]))) from its planned heap Register the shared topk-of-sum patterns so the engine matches the planner's collapsed CountMinSketchWithHeap and answers with the inner sum's labels. Co-Authored-By: Claude Opus 5.5 --- .../src/engines/simple_engine/mod.rs | 2 + .../tests/e2e_precompute_equivalence.rs | 55 +++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index c0bb98bd..9c1bfd1d 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -561,6 +561,8 @@ impl SimpleEngine { spatial_pattern("generic", &spatial_pattern_blocks), ]; controller_patterns.extend(one_temporal_one_spatial_collapsable_patterns); + controller_patterns + .extend(promql_utilities::query_logics::parsing::build_topk_of_sum_patterns()); Self { store, 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::{ From 252df57105851459443641f6f3789799ea4f5e94 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 13 Oct 2026 08:57:58 -0400 Subject: [PATCH 4/9] fix(promql-utils): bucket topk-of-sum only by labels the inner sum kept `topk by (instance) (k, sum by (job) (...))` ranks series that no longer carry `instance`, so it is one bucket, not a heap per raw `instance` of partial sums. Intersect `by` labels with topk's input labels, and report no bucket labels left as a single ranking. Also correct doc comments that still said topk outputs all labels, and drop issue references from them. Co-Authored-By: Claude Opus 5.5 --- .../rs/asap_types/src/query_requirements.rs | 14 ++--- .../src/query_logics/parsing.rs | 59 +++++++++++++++---- .../src/engines/simple_engine/promql.rs | 7 +-- 3 files changed, 56 insertions(+), 24 deletions(-) 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 82ea0694..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,12 +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, or an inner `sum by (...)`'s labels; - /// unaffected by topk's own modifier) rather than the bucketing labels - /// (#714). + /// unaffected by topk's own modifier) rather than the bucketing labels. pub topk_by_labels: Option, } @@ -98,9 +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, (collapsable, see #508) OneTemporalOneSpatial and - // topk-of-sum 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/query_logics/parsing.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs index acd0e460..27fee521 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 @@ -284,15 +284,17 @@ fn reducer_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, @@ -307,13 +309,24 @@ 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()), + let input_labels = topk_input_labels(match_result, all_labels); + let buckets = match modifier.modifier_type { + // A `by` label the input series don't carry is empty in every + // series, so it doesn't split the buckets. + AggregationModifierType::By => KeyByLabelNames::new( + modifier + .labels + .iter() + .filter(|label| input_labels.labels.contains(label)) + .cloned() + .collect(), + ), AggregationModifierType::Without => { - let without_labels = KeyByLabelNames::new(modifier.labels.clone()); - topk_input_labels(match_result, all_labels).difference(&without_labels) + input_labels.difference(&KeyByLabelNames::new(modifier.labels.clone())) } - }) + }; + // No bucket labels left is one ranking over the whole input. + Some(buckets).filter(|buckets| !buckets.is_empty()) } #[cfg(test)] @@ -414,6 +427,26 @@ mod tests { assert_eq!(get_topk_by_labels(&result, &all), 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"]); 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 From 0cd5a2fbc53544f54c61bbe5fd398150c4dc8455 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 13 Oct 2026 08:59:28 -0400 Subject: [PATCH 5/9] fix(planner): refuse a MILP deployment whose queries need different labels A Raqe doesn't carry a keyed sketch's key labels, so the solver can put topk(sum by (job) ...) and topk(sum by (instance) ...) on one heap deployment; the config then keyed the heap by the first query's labels and answered the other wrongly. Derive each served query's labels and error when they disagree. Also note that the heap is priced from the bucket grouping's shape, not its key count. Co-Authored-By: Claude Opus 5.5 --- asap-planner-rs/src/optimizer/milp.rs | 5 ++ asap-planner-rs/src/optimizer/milp_output.rs | 87 ++++++++++++++++---- 2 files changed, 78 insertions(+), 14 deletions(-) diff --git a/asap-planner-rs/src/optimizer/milp.rs b/asap-planner-rs/src/optimizer/milp.rs index 18401e25..b06813f5 100644 --- a/asap-planner-rs/src/optimizer/milp.rs +++ b/asap-planner-rs/src/optimizer/milp.rs @@ -396,6 +396,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 diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index 8f934931..93cac77b 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -57,6 +57,16 @@ pub enum MilpOutputError { raqe: Vec, deployment: Vec, }, + #[error( + "queries {first:?} and {other:?} share a deployment but need different labels \ + ({first_labels} vs {other_labels})" + )] + DeploymentLabelsDisagree { + first: String, + other: String, + first_labels: String, + other_labels: String, + }, #[error(transparent)] Generator(#[from] ControllerError), } @@ -233,21 +243,25 @@ fn aggregation_config( } }; - let requirements = &items + // A Raqe doesn't carry the labels a keyed sketch is keyed by (a + // topk-of-sum heap's inner `sum by` labels), so the solver may share one + // deployment between queries that need different keys. + let first = 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, - ); + .expect("every planned deployment serves a Raqe"); + let (grouping_labels, aggregated_labels) = subpopulation_labels(first, aggregation_type); + if let Some(other) = items.iter().find(|item| { + subpopulation_labels(item, aggregation_type) + != (grouping_labels.clone(), aggregated_labels.clone()) + }) { + let (other_grouping, other_aggregated) = subpopulation_labels(other, aggregation_type); + return Err(MilpOutputError::DeploymentLabelsDisagree { + first: first.query_strings[0].clone(), + other: other.query_strings[0].clone(), + first_labels: format!("grouping {grouping_labels:?}, aggregated {aggregated_labels:?}"), + other_labels: format!("grouping {other_grouping:?}, aggregated {other_aggregated:?}"), + }); + } Ok(IntermediateAggConfig { aggregation_type, @@ -273,6 +287,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(), @@ -672,6 +707,30 @@ 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. + #[test] + fn topk_of_sum_queries_needing_different_heap_keys_cannot_share_a_deployment() { + 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 err = plan_to_planner_output(&config, &workload, &solution) + .err() + .unwrap(); + assert!( + matches!(&err, MilpOutputError::DeploymentLabelsDisagree { first, other, .. } + if [first, other].contains(&&by_job.to_string()) + && [first, other].contains(&&by_instance.to_string())), + "{err}" + ); + } + #[test] fn query_split_across_deployments_is_an_error() { // The same query string at two cadences is two items; if they land on From 3b378487e8685f7a47055ff712011d61472ce648 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 13 Oct 2026 09:02:10 -0400 Subject: [PATCH 6/9] refactor(promql-utils): share collapsable query patterns between planner and engine The planner and the engine each built the same ONE_TEMPORAL_ONE_SPATIAL patterns from their own copy of the temporal function list, and the topk-of-sum builder sat in the parsing module as a third copy. Move both builders into ast_matching::collapsable_patterns, so the two sides can't drift apart on which shapes they accept. Co-Authored-By: Claude Opus 5.5 --- .../src/ast_matching/collapsable_patterns.rs | 106 ++++++++++++++++++ .../promql_utilities/src/ast_matching/mod.rs | 1 + .../src/query_logics/parsing.rs | 55 +-------- asap-planner-rs/src/planner/patterns.rs | 57 +--------- .../src/engines/simple_engine/mod.rs | 67 +---------- 5 files changed, 119 insertions(+), 167 deletions(-) create mode 100644 asap-common/dependencies/rs/promql_utilities/src/ast_matching/collapsable_patterns.rs 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 27fee521..dd74fd9e 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,8 +1,9 @@ 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, PromQLPattern, PromQLPatternBuilder}; +use crate::ast_matching::PromQLMatchResult; use crate::data_model::KeyByLabelNames; use crate::query_logics::enums::{AggregationOperator, PromQLFunction, Statistic}; use crate::query_logics::logics::get_is_collapsable; @@ -180,56 +181,6 @@ pub fn get_spatial_aggregation_output_labels( reducer_output_labels(aggregation_token, all_labels) } -/// Token for the inner reducer of `topk(k, sum by (...) (f(m[w])))`. -pub const INNER_AGGREGATION_TOKEN: &str = "inner_aggregation"; - -/// Patterns for `topk(k, sum by (l) (f(m[w])))`, one per `f` collapsable -/// under both `sum` and `topk`: such a query collapses into one -/// CountMinSketchWithHeap keyed by `l`. -pub fn build_topk_of_sum_patterns() -> Vec { - [ - PromQLFunction::Rate, - PromQLFunction::Increase, - PromQLFunction::SumOverTime, - PromQLFunction::CountOverTime, - PromQLFunction::AvgOverTime, - PromQLFunction::MinOverTime, - PromQLFunction::MaxOverTime, - ] - .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()], - 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"), - ), - None, - None, - None, - Some(INNER_AGGREGATION_TOKEN), - ), - None, - None, - None, - Some("aggregation"), - )) - }) - .collect() -} - /// Labels of the series topk ranks: the inner reducer's output labels when /// one is present, else all labels. fn topk_input_labels( @@ -390,7 +341,7 @@ mod tests { fn match_topk_of_sum(query: &str) -> PromQLMatchResult { let ast = promql_parser::parser::parse(query).unwrap(); - let result = build_topk_of_sum_patterns() + let result = crate::ast_matching::collapsable_patterns::build_topk_of_sum_patterns() .iter() .map(|pattern| pattern.matches(&ast)) .find(|result| result.matches); diff --git a/asap-planner-rs/src/planner/patterns.rs b/asap-planner-rs/src/planner/patterns.rs index 404a4750..285753ba 100644 --- a/asap-planner-rs/src/planner/patterns.rs +++ b/asap-planner-rs/src/planner/patterns.rs @@ -1,7 +1,8 @@ +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; -use promql_utilities::query_logics::parsing::build_topk_of_sum_patterns; /// Build all patterns in priority order: ONLY_TEMPORAL (2), ONLY_SPATIAL (1), /// ONE_TEMPORAL_ONE_SPATIAL (one per collapsable (function, op) pair), @@ -67,58 +68,8 @@ 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 } diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 9c1bfd1d..dc8e6e85 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, }; @@ -496,73 +498,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(promql_utilities::query_logics::parsing::build_topk_of_sum_patterns()); + controller_patterns.extend(build_one_temporal_one_spatial_patterns()); + controller_patterns.extend(build_topk_of_sum_patterns()); Self { store, From bc5fdfc53c47358ae5117d32afe1f0be6a499134 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 13 Oct 2026 09:06:03 -0400 Subject: [PATCH 7/9] fix(planner): pass MilpConfig in the shared-heap-deployment test plan_to_planner_output now takes a MilpConfig on main. Co-Authored-By: Claude Opus 5.5 --- asap-planner-rs/src/optimizer/milp_output.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index f5e4ee82..edb0aaf4 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -861,7 +861,7 @@ mod tests { } solution.deployments.truncate(1); - let err = plan_to_planner_output(&config, &workload, &solution) + let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) .err() .unwrap(); assert!( From 3ab6e4ad56b54747a8273af591f4f1acc1f77d0d Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 13 Oct 2026 09:17:57 -0400 Subject: [PATCH 8/9] fix(promql-utils): keep plain topk's by labels when the schema lacks them Intersecting by labels with topk's input labels turned plain topk by (job) on a metric missing from the schema (no known labels) into a global top k. Only a topk-of-sum's inner sum removes labels, so intersect only there. Co-Authored-By: Claude Opus 5.5 --- .../src/query_logics/parsing.rs | 52 +++++++++++++++---- 1 file changed, 43 insertions(+), 9 deletions(-) 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 dd74fd9e..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 @@ -260,21 +260,31 @@ pub fn get_topk_by_labels( } let modifier = aggregation_token.modifier.as_ref()?; - let input_labels = topk_input_labels(match_result, all_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 { - // A `by` label the input series don't carry is empty in every - // series, so it doesn't split the buckets. + // 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 + modifier_labels .labels - .iter() + .into_iter() .filter(|label| input_labels.labels.contains(label)) - .cloned() .collect(), ), - AggregationModifierType::Without => { - input_labels.difference(&KeyByLabelNames::new(modifier.labels.clone())) - } + 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()) @@ -378,6 +388,30 @@ mod tests { 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. From ec2fda31110d8483df547bb81747ac2f5dc0b906 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 13 Oct 2026 09:18:29 -0400 Subject: [PATCH 9/9] fix(planner): give each heap key set its own aggregation instead of failing the plan A MILP deployment shared by topk-of-sum queries keyed by different labels used to fail the whole plan. Build one aggregation per (deployment, labels a served query needs), so each query reads a correctly keyed heap. The plan still prices only one sketch for such a deployment. Co-Authored-By: Claude Opus 5.5 --- asap-planner-rs/src/optimizer/milp_output.rs | 100 ++++++++----------- 1 file changed, 39 insertions(+), 61 deletions(-) diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index edb0aaf4..453cf4cf 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; @@ -56,16 +56,6 @@ pub enum MilpOutputError { raqe: Vec, deployment: Vec, }, - #[error( - "queries {first:?} and {other:?} share a deployment but need different labels \ - ({first_labels} vs {other_labels})" - )] - DeploymentLabelsDisagree { - first: String, - other: String, - first_labels: String, - other_labels: String, - }, #[error(transparent)] Generator(#[from] ControllerError), } @@ -121,16 +111,6 @@ pub fn plan_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(); solution.deployments.len()]; - let mut seen = HashSet::new(); - for (raqe, planned) in solution.raqes.iter().enumerate() { - if seen.insert((planned.deployment, workload.raqe_items[raqe])) { - served[planned.deployment].push(item_of(raqe)); - } - } - // A SQL aggregation names its table and value column, and its metric is // the table. let sql_columns: HashMap = config @@ -139,10 +119,21 @@ pub fn plan_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 (planned, items) in solution.deployments.iter().zip(&served) { - let mut aggregation = aggregation_config(&planned.deployment, items)?; + let mut item_keys: HashMap<(usize, usize), String> = HashMap::new(); + for (raqe, planned) in solution.raqes.iter().enumerate() { + let item = workload.raqe_items[raqe]; + if item_keys.contains_key(&(planned.deployment, item)) { + continue; + } + let deployment = &solution.deployments[planned.deployment].deployment; + let mut aggregation = aggregation_config(deployment, &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()); @@ -150,13 +141,14 @@ pub fn plan_to_planner_output( } let key = aggregation.identifying_key(); aggregations.entry(key.clone()).or_insert(aggregation); - deployment_keys.push(key); + item_keys.insert((planned.deployment, item), key); } let mut queries: IndexMap = IndexMap::new(); for (raqe, planned) in solution.raqes.iter().enumerate() { - let key = &deployment_keys[planned.deployment]; - for query in &item_of(raqe).query_strings { + let item = workload.raqe_items[raqe]; + let key = &item_keys[&(planned.deployment, item)]; + for query in &workload.items[item].query_strings { match queries.entry(query.clone()) { Entry::Vacant(entry) => { entry.insert(QueryPlanEntry::fully_planned( @@ -199,12 +191,12 @@ pub fn plan_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| { @@ -270,25 +262,7 @@ fn aggregation_config( } }; - // A Raqe doesn't carry the labels a keyed sketch is keyed by (a - // topk-of-sum heap's inner `sum by` labels), so the solver may share one - // deployment between queries that need different keys. - let first = items - .first() - .expect("every planned deployment serves a Raqe"); - let (grouping_labels, aggregated_labels) = subpopulation_labels(first, aggregation_type); - if let Some(other) = items.iter().find(|item| { - subpopulation_labels(item, aggregation_type) - != (grouping_labels.clone(), aggregated_labels.clone()) - }) { - let (other_grouping, other_aggregated) = subpopulation_labels(other, aggregation_type); - return Err(MilpOutputError::DeploymentLabelsDisagree { - first: first.query_strings[0].clone(), - other: other.query_strings[0].clone(), - first_labels: format!("grouping {grouping_labels:?}, aggregated {aggregated_labels:?}"), - other_labels: format!("grouping {other_grouping:?}, aggregated {other_aggregated:?}"), - }); - } + let (grouping_labels, aggregated_labels) = subpopulation_labels(item, aggregation_type); Ok(IntermediateAggConfig { aggregation_type, @@ -849,9 +823,10 @@ 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. + /// 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_cannot_share_a_deployment() { + 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) = @@ -861,15 +836,18 @@ mod tests { } solution.deployments.truncate(1); - let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) - .err() - .unwrap(); - assert!( - matches!(&err, MilpOutputError::DeploymentLabelsDisagree { first, other, .. } - if [first, other].contains(&&by_job.to_string()) - && [first, other].contains(&&by_instance.to_string())), - "{err}" - ); + 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]