refactor(planner): dedup logic for extracting PromQL QueryRequirements (#430)

Author: milindsrivastava1997

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

@@ -1,5 +1,13 @@
+use promql_utilities::ast_matching::PromQLMatchResult;
 use promql_utilities::data_model::KeyByLabelNames;
-use promql_utilities::query_logics::enums::Statistic;
+use promql_utilities::query_logics::enums::{QueryPatternType, Statistic};
+use promql_utilities::query_logics::parsing::{
+    get_metric_and_spatial_filter, get_spatial_aggregation_output_labels, get_statistics_to_compute,
+};
+use tracing::warn;
+
+use crate::promql_schema::PromQLSchema;
+use crate::utils::normalize_spatial_filter;
 
 /// What a query needs in order to be answered by a stored aggregation.
 #[derive(Debug, Clone)]
@@ -28,3 +36,60 @@ pub struct QueryRequirements {
     /// constrain the sketch weighting); matching ignores it when `None`.
     pub topk_count_events: Option<bool>,
 }
+
+/// Build `QueryRequirements` from an already-pattern-matched PromQL query.
+///
+/// Shared by the query engine's capability-matching fallback and the planner's
+/// AQE extractor, which independently parse and pattern-match a query before
+/// calling this. `query` is used only for diagnostics on the unsupported-stats
+/// warning.
+pub fn build_query_requirements_promql(
+    query: &str,
+    match_result: &PromQLMatchResult,
+    pattern_type: QueryPatternType,
+    metric_schema: &PromQLSchema,
+) -> Option<QueryRequirements> {
+    let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result);
+
+    let statistics = get_statistics_to_compute(pattern_type, match_result)
+        .map_err(|err| {
+            warn!(
+                query = %query,
+                error = %err,
+                "skipping matched query with unsupported statistics"
+            );
+            err
+        })
+        .ok()?;
+
+    let data_range_ms = match pattern_type {
+        QueryPatternType::OnlySpatial => None,
+        _ => match_result
+            .get_range_duration()
+            .map(|d| d.num_seconds() as u64 * 1000),
+    };
+
+    let all_labels = metric_schema
+        .get_labels(&metric)
+        .cloned()
+        .unwrap_or_else(KeyByLabelNames::empty);
+
+    let grouping_labels = match pattern_type {
+        // 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)
+        }
+    };
+
+    Some(QueryRequirements {
+        metric,
+        statistics,
+        data_range_ms,
+        grouping_labels,
+        spatial_filter_normalized: normalize_spatial_filter(&spatial_filter),
+        topk_count_events: None,
+    })
+}

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

@@ -1,13 +1,9 @@
 use std::collections::HashMap;
 
-use asap_types::query_requirements::QueryRequirements;
-use asap_types::utils::normalize_spatial_filter;
+use asap_types::query_requirements::{build_query_requirements_promql, QueryRequirements};
 use asap_types::PromQLSchema;
 use promql_utilities::data_model::KeyByLabelNames;
-use promql_utilities::query_logics::enums::{QueryPatternType, Statistic};
-use promql_utilities::query_logics::parsing::{
-    get_metric_and_spatial_filter, get_spatial_aggregation_output_labels, get_statistics_to_compute,
-};
+use promql_utilities::query_logics::enums::Statistic;
 use tracing::warn;
 
 use crate::planner::patterns::build_patterns;
@@ -160,12 +156,6 @@ fn decompose_to_leaves(query: &str) -> Vec<String> {
 
 /// Try to extract `QueryRequirements` from a single leaf PromQL query string.
 /// Returns `None` if the query cannot be parsed or does not match any pattern.
-///
-/// TODO: this duplicates `build_query_requirements_promql` in
-/// `asap-query-engine/src/engines/simple_engine/promql.rs:614`. That function
-/// is a private `&self` method tied to `SimplePromQLEngine`. The shared logic
-/// should be extracted into a free function in `asap_types::query_requirements`
-/// and called from both sites.
 fn extract_requirements(query: &str, metric_schema: &PromQLSchema) -> Option<QueryRequirements> {
     let ast = promql_parser::parser::parse(query).ok()?;
     let patterns = build_patterns();
@@ -179,52 +169,7 @@ fn extract_requirements(query: &str, metric_schema: &PromQLSchema) -> Option<Que
         }
     })?;
 
-    let (metric, spatial_filter) = get_metric_and_spatial_filter(&match_result);
-    let statistics = get_statistics_to_compute(pattern_type, &match_result)
-        .map_err(|err| {
-            warn!(
-                query = %query,
-                error = %err,
-                "aqe_extractor: skipping matched leaf query with unsupported statistics"
-            );
-            err
-        })
-        .ok()?;
-
-    let data_range_ms = match pattern_type {
-        QueryPatternType::OnlySpatial => None,
-        _ => match_result
-            .get_range_duration()
-            .map(|d| d.num_seconds() as u64 * 1000),
-    };
-
-    let grouping_labels = match pattern_type {
-        // OnlyTemporal preserves all labels — look them up in the schema.
-        // If the metric is unknown, fall back to empty (dedup still works; cost
-        // model will treat it as a zero-group-count sketch).
-        QueryPatternType::OnlyTemporal => metric_schema
-            .get_labels(&metric)
-            .cloned()
-            .unwrap_or_else(KeyByLabelNames::empty),
-        // OnlySpatial and OneTemporalOneSpatial encode their output labels in
-        // the AST's `by (...)` / `without (...)` clause.
-        QueryPatternType::OnlySpatial | QueryPatternType::OneTemporalOneSpatial => {
-            let all_labels = metric_schema
-                .get_labels(&metric)
-                .cloned()
-                .unwrap_or_else(KeyByLabelNames::empty);
-            get_spatial_aggregation_output_labels(&match_result, &all_labels)
-        }
-    };
-
-    Some(QueryRequirements {
-        metric,
-        statistics,
-        data_range_ms,
-        grouping_labels,
-        spatial_filter_normalized: normalize_spatial_filter(&spatial_filter),
-        topk_count_events: None,
-    })
+    build_query_requirements_promql(query, &match_result, pattern_type, metric_schema)
 }
 
 #[cfg(test)]

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

@@ -10,8 +10,8 @@ use super::{
 };
 use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig};
 use crate::engines::query_result::{QueryResult, RangeVectorElement};
-use asap_types::query_requirements::QueryRequirements;
-use asap_types::utils::normalize_spatial_filter;
+use asap_types::query_requirements::build_query_requirements_promql;
+use asap_types::PromQLSchema;
 use promql_utilities::ast_matching::PromQLMatchResult;
 use promql_utilities::data_model::KeyByLabelNames;
 use promql_utilities::get_is_collapsable;
@@ -614,56 +614,6 @@ impl SimpleEngine {
         Some((output_labels, QueryResult::matrix(combined)))
     }
 
-    /// Extract QueryRequirements from a parsed PromQL match result.
-    /// Used as the fallback path when no query_configs entry is found.
-    fn build_query_requirements_promql(
-        &self,
-        match_result: &PromQLMatchResult,
-        query_pattern_type: QueryPatternType,
-    ) -> Option<QueryRequirements> {
-        let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result);
-
-        let statistics = get_statistics_to_compute(query_pattern_type, match_result)
-            .map_err(|err| {
-                warn!("{}", err);
-                err
-            })
-            .ok()?;
-
-        let data_range_ms = match query_pattern_type {
-            QueryPatternType::OnlySpatial => None,
-            _ => match_result
-                .get_range_duration()
-                .map(|d| d.num_seconds() as u64 * 1000),
-        };
-
-        let all_labels = match &self.inference_config.read().unwrap().schema {
-            SchemaConfig::PromQL(schema) => schema
-                .get_labels(&metric)
-                .cloned()
-                .unwrap_or_else(KeyByLabelNames::empty),
-            _ => KeyByLabelNames::empty(),
-        };
-
-        let grouping_labels = match query_pattern_type {
-            QueryPatternType::OnlyTemporal => all_labels,
-            QueryPatternType::OnlySpatial | QueryPatternType::OneTemporalOneSpatial => {
-                get_spatial_aggregation_output_labels(match_result, &all_labels)
-            }
-        };
-
-        Some(QueryRequirements {
-            metric,
-            statistics,
-            data_range_ms,
-            grouping_labels,
-            spatial_filter_normalized: normalize_spatial_filter(&spatial_filter),
-            // PromQL top-k does not constrain the sketch weighting; leave the
-            // count/sum discriminator unset so matching does not over-filter.
-            topk_count_events: None,
-        })
-    }
-
     // /// Try to extract sketch query components from a PromQL query string.
     // ///
     // /// Attempts the standard AST parser first. If that fails (e.g. for custom
@@ -1085,8 +1035,18 @@ impl SimpleEngine {
                 "No query_config entry for PromQL query '{}'. Attempting capability-based matching.",
                 query
             );
-            let requirements =
-                self.build_query_requirements_promql(&match_result, query_pattern_type)?;
+            let inference_config = self.inference_config.read().unwrap();
+            let empty_schema = PromQLSchema::new();
+            let metric_schema = match &inference_config.schema {
+                SchemaConfig::PromQL(schema) => schema,
+                _ => &empty_schema,
+            };
+            let requirements = build_query_requirements_promql(
+                &query,
+                &match_result,
+                query_pattern_type,
+                metric_schema,
+            )?;
             self.streaming_config
                 .read()
                 .unwrap()