fix(planner,query-engine): recognize topk over sum_over_time/count_over_time

The planner and query-engine's PromQL pattern matchers only recognized
topk over a raw vector selector. Their spatial-over-temporal collapsable
pattern lists covered sum/count/avg/quantile/min/max but omitted topk,
so `topk(k, sum_over_time(...))` and `topk(k, count_over_time(...))`
were silently dropped from inference_config.yaml even though the
underlying temporal aggregation was planned (#699).

- get_is_collapsable now treats (Topk, SumOverTime) and (Topk,
  CountOverTime) as collapsable; both pattern-generator lists (planner
  and query-engine each carry their own copy) include Topk.
- get_statistics_to_compute special-cases Topk: the outer op always
  wins over the inner function's statistic (unlike sum/min/max, which
  absorb into the same accumulator), still checked by the existing
  collapsability assert.
- topk_count_events is now derived per query via the new
  promql_topk_count_events helper: sum_over_time (or a raw vector) is
  value-weighted, count_over_time is count-weighted. Previously this was
  hardcoded to value-weighted everywhere PromQL topk requirements were
  built.

Covers all 24 topk-over-temporal cases in the aggregations differential
suite.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01CD4LRkJNNDb1TQ17XeQjmy

Author: milindsrivastava1997

asap-common/dependencies/rs/asap_types/src/query_requirements.rs

@@ -1,6 +1,7 @@
 use promql_utilities::ast_matching::PromQLMatchResult;
 use promql_utilities::data_model::KeyByLabelNames;
 use promql_utilities::query_logics::enums::Statistic;
+use promql_utilities::query_logics::logics::promql_topk_count_events;
 use promql_utilities::query_logics::parsing::{
     get_metric_and_spatial_filter, get_spatial_aggregation_output_labels, get_statistics_to_compute,
 };
@@ -32,8 +33,9 @@ pub struct QueryRequirements {
     ///   * `Some(true)`  → COUNT semantics (`count_events: true`, weight 1/event),
     ///   * `Some(false)` → SUM semantics (`count_events: false`, weight = value).
     ///
-    /// PromQL top-k is value-weighted, so it uses `Some(false)`. `None` is only
-    /// valid for non-top-k requirements.
+    /// PromQL top-k over a raw vector or `sum_over_time` is value-weighted
+    /// (`Some(false)`); over `count_over_time` it is count-weighted
+    /// (`Some(true)`). `None` is only valid for non-top-k requirements.
     pub topk_count_events: Option<bool>,
 }
 
@@ -91,7 +93,7 @@ pub fn build_query_requirements_promql(
         all_labels
     };
 
-    let topk_count_events = statistics.contains(&Statistic::Topk).then_some(false);
+    let topk_count_events = promql_topk_count_events(match_result);
 
     Some(QueryRequirements {
         metric,

asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs

@@ -107,10 +107,45 @@ pub fn get_is_collapsable(
         ),
         AggregationOperator::Min => temporal_aggregation == PromQLFunction::MinOverTime,
         AggregationOperator::Max => temporal_aggregation == PromQLFunction::MaxOverTime,
+        // topk ranks by whatever value the temporal function produces per
+        // series: sum_over_time's summed value or count_over_time's event
+        // count. See `Statistic::Topk`'s handling in `get_statistics_to_compute`
+        // and `promql_topk_count_events`, which derive from the same pairing.
+        AggregationOperator::Topk => matches!(
+            temporal_aggregation,
+            PromQLFunction::SumOverTime | PromQLFunction::CountOverTime
+        ),
         _ => false,
     }
 }
 
+/// For a topk match result, the `count_events` weighting the wrapped temporal
+/// function implies:
+///   * no temporal function (raw vector topk), or `sum_over_time` → `Some(false)`
+///     (value-weighted: topk ranks by the value itself),
+///   * `count_over_time` → `Some(true)` (count-weighted: topk ranks by the
+///     per-series event count).
+///
+/// Returns `None` if the match result is not a topk aggregation.
+pub fn promql_topk_count_events(
+    match_result: &crate::ast_matching::PromQLMatchResult,
+) -> Option<bool> {
+    if match_result
+        .get_aggregation_op()?
+        .parse::<AggregationOperator>()
+        != Ok(AggregationOperator::Topk)
+    {
+        return None;
+    }
+    match match_result
+        .get_function_name()
+        .and_then(|name| name.parse::<PromQLFunction>().ok())
+    {
+        Some(PromQLFunction::CountOverTime) => Some(true),
+        _ => Some(false),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

asap-common/dependencies/rs/promql_utilities/src/query_logics/parsing.rs

@@ -84,11 +84,16 @@ pub fn get_metric_and_spatial_filter(match_result: &PromQLMatchResult) -> (Strin
 ///   reachable via a pattern already narrowed to a collapsable `(function,
 ///   op)` pair (see `get_is_collapsable`, and #508's pattern-narrowing fix
 ///   that makes non-collapsable combinations fail to match at all) — asserted
-///   below rather than silently trusted. The statistic still comes from the
-///   *function*, never the outer op: e.g. `count_over_time` + `sum` needs a
-///   `Count` accumulator, not a `Sum` one — summing per-series counts gives
-///   the group's total count, so the outer op only describes how per-series
-///   results combine, never which statistic must be precomputed.
+///   below rather than silently trusted. Which side supplies the statistic
+///   then depends on the outer op:
+///     - `topk`: the statistic is always `Topk` — a `topk(k, sum_over_time(x))`
+///       still needs a heavy-hitter sketch, not a `Sum` accumulator, so the
+///       outer op is never dropped (see #699).
+///     - every other collapsable op: the statistic comes from the *function*,
+///       never the outer op — e.g. `count_over_time` + `sum` needs a `Count`
+///       accumulator, not a `Sum` one, since summing per-series counts gives
+///       the group's total count. Here the outer op only describes how
+///       per-series results combine, never which statistic must be precomputed.
 ///
 /// Returns a typed error if the matched statistic/function name is not
 /// recognized, so callers can decide whether to skip or fail the query.
@@ -107,20 +112,23 @@ pub fn get_statistics_to_compute(
     };
 
     let statistic_to_compute: Option<String> = if has_function && has_aggregation {
+        let aggregation_op = match_result
+            .get_aggregation_op()
+            .and_then(|o| o.parse::<AggregationOperator>().ok());
         debug_assert!(
             match_result
                 .get_function_name()
                 .and_then(|f| f.parse::<PromQLFunction>().ok())
-                .zip(
-                    match_result
-                        .get_aggregation_op()
-                        .and_then(|o| o.parse::<AggregationOperator>().ok())
-                )
+                .zip(aggregation_op)
                 .is_some_and(|(f, o)| get_is_collapsable(f, o)),
             "a match with both function and aggregation tokens must be collapsable \
              (patterns are narrowed to only collapsable pairs, see #508)"
         );
-        function_statistic(match_result)
+        if aggregation_op == Some(AggregationOperator::Topk) {
+            Some(AggregationOperator::Topk.as_str().to_string())
+        } else {
+            function_statistic(match_result)
+        }
     } else if has_function {
         function_statistic(match_result)
     } else if has_aggregation {

asap-planner-rs/src/planner/patterns.rs

@@ -90,6 +90,7 @@ pub fn build_patterns() -> Vec<PromQLPattern> {
             AggregationOperator::Quantile,
             AggregationOperator::Min,
             AggregationOperator::Max,
+            AggregationOperator::Topk,
         ]
         .into_iter()
         .filter_map(move |op| {
@@ -131,10 +132,11 @@ mod tests {
     }
 
     #[test]
-    fn exactly_four_collapsable_one_temporal_one_spatial_patterns() {
-        // 2 ONLY_TEMPORAL + 1 ONLY_SPATIAL + 4 collapsable ONE_TEMPORAL_ONE_SPATIAL
-        // (sum+sum_over_time, sum+count_over_time, min+min_over_time, max+max_over_time).
-        assert_eq!(build_patterns().len(), 7);
+    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);
     }
 
     #[test]
@@ -143,6 +145,10 @@ mod tests {
         assert!(matches_some_pattern("sum(count_over_time(x[5m]))"));
         assert!(matches_some_pattern("min(min_over_time(x[5m]))"));
         assert!(matches_some_pattern("max(max_over_time(x[5m]))"));
+        assert!(matches_some_pattern("topk(1, sum_over_time(x[5m]))"));
+        assert!(matches_some_pattern(
+            "topk by (job) (3, count_over_time(x[5m]))"
+        ));
     }
 
     #[test]
@@ -163,5 +169,13 @@ mod tests {
             !matches_some_pattern("min(max_over_time(x[5m]))"),
             "min+max_over_time is not collapsable"
         );
+        assert!(
+            !matches_some_pattern("topk(1, rate(x[5m]))"),
+            "topk+rate is not collapsable"
+        );
+        assert!(
+            !matches_some_pattern("topk(1, avg_over_time(x[5m]))"),
+            "topk+avg_over_time is not collapsable"
+        );
     }
 }

asap-planner-rs/src/planner/sketch.rs

@@ -1,6 +1,7 @@
 use crate::config::input::SketchParameterOverrides;
 use promql_utilities::ast_matching::PromQLMatchResult;
 use promql_utilities::query_logics::enums::AggregationType;
+use promql_utilities::query_logics::logics::promql_topk_count_events;
 use std::collections::HashMap;
 
 // Default sketch parameters
@@ -163,11 +164,18 @@ pub fn build_sketch_parameters_from_promql(
     } else {
         None
     };
+    let topk_count_events = if aggregation_type == AggregationType::CountMinSketchWithHeap {
+        Some(promql_topk_count_events(match_result).ok_or_else(|| {
+            "topk query missing required aggregation match to derive count_events".to_string()
+        })?)
+    } else {
+        None
+    };
     build_sketch_parameters(
         aggregation_type,
         aggregation_sub_type,
         topk_k,
-        Some(false),
+        topk_count_events,
         sketch_params,
     )
 }

asap-planner-rs/tests/integration.rs

@@ -463,6 +463,62 @@ fn topk_produces_count_min_sketch_with_heap() {
     );
 }
 
+#[test]
+fn topk_over_sum_over_time_produces_value_weighted_heap() {
+    // https://github.com/ProjectASAP/asap-internal/issues/699 — topk wrapping
+    // a temporal aggregation must still be planned, not silently omitted.
+    let c = Controller::from_yaml_with_schema(
+        r#"
+query_groups:
+  - id: 1
+    queries:
+      - "topk(1, sum_over_time(http_requests_total[5m]))"
+    repetition_delay_ms: 60000
+    controller_options:
+      accuracy_sla: 0.99
+      latency_sla: 1.0
+"#,
+        http_requests_schema(),
+        arroyo_opts(),
+    )
+    .unwrap();
+    let out = c.generate().unwrap();
+    assert_eq!(out.inference_query_count(), 1);
+    assert!(out.has_aggregation_type("CountMinSketchWithHeap"));
+    assert_eq!(
+        out.aggregation_parameter("CountMinSketchWithHeap", "count_events"),
+        Some(serde_yaml::Value::Bool(false)),
+        "topk ranks the summed value, not the observation count"
+    );
+}
+
+#[test]
+fn topk_over_count_over_time_produces_count_weighted_heap() {
+    let c = Controller::from_yaml_with_schema(
+        r#"
+query_groups:
+  - id: 1
+    queries:
+      - "topk by (job) (3, count_over_time(http_requests_total[5m]))"
+    repetition_delay_ms: 60000
+    controller_options:
+      accuracy_sla: 0.99
+      latency_sla: 1.0
+"#,
+        http_requests_schema(),
+        arroyo_opts(),
+    )
+    .unwrap();
+    let out = c.generate().unwrap();
+    assert_eq!(out.inference_query_count(), 1);
+    assert!(out.has_aggregation_type("CountMinSketchWithHeap"));
+    assert_eq!(
+        out.aggregation_parameter("CountMinSketchWithHeap", "count_events"),
+        Some(serde_yaml::Value::Bool(true)),
+        "topk over count_over_time ranks the observation count"
+    );
+}
+
 #[test]
 fn heap_parameters_require_explicit_count_events_weighting() {
     let error = asap_planner::planner::sketch::build_sketch_parameters(

asap-query-engine/src/engines/simple_engine/mod.rs

@@ -301,6 +301,7 @@ impl SimpleEngine {
                 AggregationOperator::Quantile,
                 AggregationOperator::Min,
                 AggregationOperator::Max,
+                AggregationOperator::Topk,
             ]
             .into_iter()
             .filter_map(move |op| {

asap-query-engine/src/engines/simple_engine/promql.rs

@@ -1453,6 +1453,47 @@ mod topk_pipeline_tests {
         assert_eq!(requirements.topk_count_events, Some(false));
     }
 
+    #[test]
+    fn topk_over_sum_over_time_matches_and_is_value_weighted() {
+        // https://github.com/ProjectASAP/asap-internal/issues/699
+        let (engine, _store) = build_topk_engine();
+        let query = "topk(1, sum_over_time(transfer_events[5m]))";
+        let ast = promql_parser::parser::parse(query).unwrap();
+        let match_result = engine
+            .find_matching_controller_pattern(&ast, query)
+            .expect("topk over sum_over_time should match a controller pattern");
+        let schema = PromQLSchema::new().add_metric(
+            METRIC.to_string(),
+            KeyByLabelNames::new(vec!["srcip".to_string()]),
+        );
+        let requirements =
+            asap_types::build_query_requirements_promql(query, &match_result, &schema, 1000)
+                .expect("topk over sum_over_time requirements should build");
+
+        assert_eq!(requirements.statistics, vec![Statistic::Topk]);
+        assert_eq!(requirements.topk_count_events, Some(false));
+    }
+
+    #[test]
+    fn topk_over_count_over_time_matches_and_is_count_weighted() {
+        let (engine, _store) = build_topk_engine();
+        let query = "topk by (job) (3, count_over_time(transfer_events[5m]))";
+        let ast = promql_parser::parser::parse(query).unwrap();
+        let match_result = engine
+            .find_matching_controller_pattern(&ast, query)
+            .expect("topk over count_over_time should match a controller pattern");
+        let schema = PromQLSchema::new().add_metric(
+            METRIC.to_string(),
+            KeyByLabelNames::new(vec!["srcip".to_string()]),
+        );
+        let requirements =
+            asap_types::build_query_requirements_promql(query, &match_result, &schema, 1000)
+                .expect("topk over count_over_time requirements should build");
+
+        assert_eq!(requirements.statistics, vec![Statistic::Topk]);
+        assert_eq!(requirements.topk_count_events, Some(true));
+    }
+
     fn build_topk_engine() -> (SimpleEngine, Arc<SimpleMapStore>) {
         let promql_schema = PromQLSchema::new().add_metric(
             METRIC.to_string(),