refactor(planner, query-engine): remove QueryPatternType for PromQL path (#518)

* Restricted OneTemporalOneSpatial patterns to only collapsable ones

* made get_statistics_to_compute independent of QueryPatternType

* made build_query_requirements_promql independent of QueryPatternType

* Removed remaining QueryPatternType from asap-query-engine

* Removed QueryPatternType from asap-planner

* removed enum

* added check

* fix

* fix clippy

Author: milindsrivastava1997

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

@@ -1,6 +1,6 @@
 use promql_utilities::ast_matching::PromQLMatchResult;
 use promql_utilities::data_model::KeyByLabelNames;
-use promql_utilities::query_logics::enums::{QueryPatternType, Statistic};
+use promql_utilities::query_logics::enums::Statistic;
 use promql_utilities::query_logics::parsing::{
     get_metric_and_spatial_filter, get_spatial_aggregation_output_labels, get_statistics_to_compute,
 };
@@ -46,13 +46,12 @@ pub struct QueryRequirements {
 pub fn build_query_requirements_promql(
     query: &str,
     match_result: &PromQLMatchResult,
-    pattern_type: QueryPatternType,
     metric_schema: &PromQLSchema,
     data_ingestion_interval_ms: u64,
 ) -> Option<QueryRequirements> {
     let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result);
 
-    let statistics = get_statistics_to_compute(pattern_type, match_result)
+    let statistics = get_statistics_to_compute(match_result)
         .map_err(|err| {
             warn!(
                 query = %query,
@@ -63,28 +62,33 @@ pub fn build_query_requirements_promql(
         })
         .ok()?;
 
-    let data_range_ms = match pattern_type {
-        QueryPatternType::OnlySpatial => data_ingestion_interval_ms,
+    let has_temporal_function = match_result.tokens.contains_key("function");
+    let has_aggregation = match_result.tokens.contains_key("aggregation");
+
+    let data_range_ms = if has_temporal_function {
         // promql-parser supports a literal `ms` duration suffix (e.g. `[500ms]`),
         // so .num_seconds() would truncate sub-second ranges to 0.
-        _ => match_result
+        match_result
             .get_range_duration()
-            .map(|d| d.num_milliseconds() as u64)?,
+            .map(|d| d.num_milliseconds() as u64)?
+    } else {
+        // OnlySpatial (no temporal component): the query has no range of its
+        // own, so its data range is exactly one scrape interval.
+        data_ingestion_interval_ms
     };
 
     let all_labels = metric_schema
         .get_labels(&metric)
         .cloned()
         .unwrap_or_else(KeyByLabelNames::empty);
 
-    let grouping_labels = match pattern_type {
+    let grouping_labels = if has_aggregation {
+        // OnlySpatial and (collapsable, see #508) OneTemporalOneSpatial 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.
-        QueryPatternType::OnlyTemporal => all_labels,
-        // OnlySpatial and OneTemporalOneSpatial encode their output labels in
-        // the AST's `by (...)` / `without (...)` clause.
-        QueryPatternType::OnlySpatial | QueryPatternType::OneTemporalOneSpatial => {
-            get_spatial_aggregation_output_labels(match_result, &all_labels)
-        }
+        all_labels
     };
 
     Some(QueryRequirements {

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

@@ -3,24 +3,6 @@ use std::fmt;
 use std::str::FromStr;
 use tracing::debug;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
-pub enum QueryPatternType {
-    OnlyTemporal,
-    OnlySpatial,
-    OneTemporalOneSpatial,
-}
-
-impl std::fmt::Display for QueryPatternType {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        debug!("Formatting QueryPatternType: {:?}", self);
-        match self {
-            QueryPatternType::OnlyTemporal => write!(f, "only_temporal"),
-            QueryPatternType::OnlySpatial => write!(f, "only_spatial"),
-            QueryPatternType::OneTemporalOneSpatial => write!(f, "one_temporal_one_spatial"),
-        }
-    }
-}
-
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub enum QueryTreatmentType {
     Exact,
@@ -69,6 +51,22 @@ impl std::fmt::Display for Statistic {
 
 #[allow(clippy::should_implement_trait)]
 impl Statistic {
+    /// Returns `true` for statistics whose result requires approximate
+    /// pre-aggregation, regardless of whether they were reached via a
+    /// temporal function (`PromQLFunction::is_approximate`) or a spatial
+    /// aggregation operator (`AggregationOperator::is_approximate`) — for
+    /// every `Statistic` reachable from either origin, the two origins agree.
+    pub fn is_approximate(self) -> bool {
+        matches!(
+            self,
+            Statistic::Count
+                | Statistic::Sum
+                | Statistic::Cardinality
+                | Statistic::Quantile
+                | Statistic::Topk
+        )
+    }
+
     pub fn from_str(s: &str) -> Option<Self> {
         debug!("Parsing Statistic from string: {}", s);
         match s.to_lowercase().as_str() {

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

@@ -4,21 +4,22 @@ use tracing::debug;
 use crate::ast_matching::promql_pattern::AggregationModifierType;
 use crate::ast_matching::PromQLMatchResult;
 use crate::data_model::KeyByLabelNames;
-use crate::query_logics::enums::{AggregationOperator, QueryPatternType, Statistic};
+use crate::query_logics::enums::{AggregationOperator, PromQLFunction, Statistic};
+use crate::query_logics::logics::get_is_collapsable;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum StatisticExtractionError {
-    MissingStatistic { pattern_type: QueryPatternType },
+    MissingStatistic,
     UnsupportedStatistic { statistic: String },
 }
 
 impl std::fmt::Display for StatisticExtractionError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::MissingStatistic { pattern_type } => {
+            Self::MissingStatistic => {
                 write!(
                     f,
-                    "No statistic found for query pattern type {pattern_type:?}"
+                    "No temporal function or aggregation operation found in match result"
                 )
             }
             Self::UnsupportedStatistic { statistic } => {
@@ -72,28 +73,66 @@ pub fn get_metric_and_spatial_filter(match_result: &PromQLMatchResult) -> (Strin
     (metric_name, spatial_filter)
 }
 
-/// Get statistics to compute based on pattern type and tokens.
+/// Get statistics to compute from a matched query's tokens.
+///
+/// Explicitly handles the three reachable shapes:
+/// - Only a temporal function (`"function"` token, no `"aggregation"`): the
+///   statistic comes from the function name.
+/// - Only a spatial aggregation (`"aggregation"` token, no `"function"`): the
+///   statistic comes from the aggregation operator.
+/// - Both (a spatial aggregation wrapping a temporal function): only ever
+///   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.
+///
 /// Returns a typed error if the matched statistic/function name is not
 /// recognized, so callers can decide whether to skip or fail the query.
 pub fn get_statistics_to_compute(
-    pattern_type: QueryPatternType,
     match_result: &PromQLMatchResult,
 ) -> Result<Vec<Statistic>, StatisticExtractionError> {
-    debug!("Computing statistics for pattern type {:?}", pattern_type);
-    let statistic_to_compute: Option<String> = match pattern_type {
-        QueryPatternType::OnlyTemporal | QueryPatternType::OneTemporalOneSpatial => {
-            match_result.get_function_name().map(|function_name| {
-                let name = function_name.to_lowercase();
-                name.split('_').next().unwrap_or(&name).to_string()
-            })
-        }
-        QueryPatternType::OnlySpatial => match_result
+    let has_function = match_result.tokens.contains_key("function");
+    let has_aggregation = match_result.tokens.contains_key("aggregation");
+    debug!("Computing statistics (has_function={has_function}, has_aggregation={has_aggregation})");
+
+    let function_statistic = |match_result: &PromQLMatchResult| {
+        match_result.get_function_name().map(|function_name| {
+            let name = function_name.to_lowercase();
+            name.split('_').next().unwrap_or(&name).to_string()
+        })
+    };
+
+    let statistic_to_compute: Option<String> = if has_function && has_aggregation {
+        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())
+                )
+                .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)
+    } else if has_function {
+        function_statistic(match_result)
+    } else if has_aggregation {
+        match_result
             .get_aggregation_op()
-            .map(|agg| agg.to_lowercase()),
+            .map(|agg| agg.to_lowercase())
+    } else {
+        None
     };
 
     let Some(statistic_to_compute) = statistic_to_compute else {
-        return Err(StatisticExtractionError::MissingStatistic { pattern_type });
+        return Err(StatisticExtractionError::MissingStatistic);
     };
 
     debug!("Found statistic to compute: {}", statistic_to_compute);
@@ -194,9 +233,7 @@ mod tests {
 
     #[test]
     fn unsupported_matched_temporal_statistic_returns_typed_error() {
-        let err =
-            get_statistics_to_compute(QueryPatternType::OnlyTemporal, &temporal_match("stddev"))
-                .unwrap_err();
+        let err = get_statistics_to_compute(&temporal_match("stddev")).unwrap_err();
 
         assert_eq!(
             err,
@@ -208,17 +245,9 @@ mod tests {
 
     #[test]
     fn missing_statistic_returns_typed_error() {
-        let err = get_statistics_to_compute(
-            QueryPatternType::OnlyTemporal,
-            &PromQLMatchResult::with_tokens(HashMap::new()),
-        )
-        .unwrap_err();
+        let err =
+            get_statistics_to_compute(&PromQLMatchResult::with_tokens(HashMap::new())).unwrap_err();
 
-        assert_eq!(
-            err,
-            StatisticExtractionError::MissingStatistic {
-                pattern_type: QueryPatternType::OnlyTemporal
-            }
-        );
+        assert_eq!(err, StatisticExtractionError::MissingStatistic);
     }
 }

asap-planner-rs/src/optimizer/aqe_extractor.rs

@@ -170,10 +170,10 @@ fn extract_requirements(
     let ast = promql_parser::parser::parse(query).ok()?;
     let patterns = build_patterns();
 
-    let (pattern_type, match_result) = patterns.iter().find_map(|(pt, pat)| {
+    let match_result = patterns.iter().find_map(|pat| {
         let r = pat.matches(&ast);
         if r.matches {
-            Some((*pt, r))
+            Some(r)
         } else {
             None
         }
@@ -182,7 +182,6 @@ fn extract_requirements(
     build_query_requirements_promql(
         query,
         &match_result,
-        pattern_type,
         metric_schema,
         data_ingestion_interval_ms,
     )

asap-planner-rs/src/optimizer/greedy.rs

@@ -126,7 +126,7 @@ mod tests {
         );
 
         let mut seen_ids: StdHashMap<u64, ()> = StdHashMap::new();
-        for (id, _) in solution.deployed_configs.iter() {
+        for id in solution.deployed_configs.keys() {
             assert!(
                 seen_ids.insert(*id, ()).is_none(),
                 "duplicate aggregation_id"