feat(query-engine): execute nested PromQL aggregation DAGs

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

@@ -7,7 +7,7 @@ use std::io::BufReader;
 use crate::aggregation_reference::AggregationReference;
 use crate::enums::{CleanupPolicy, QueryLanguage};
 use crate::promql_schema::PromQLSchema;
-use crate::query_config::QueryConfig;
+use crate::query_config::{QueryConfig, QueryTimeAggregation};
 use elastic_dsl_utilities::{ElasticIndexSchema, ElasticMappingSchema};
 use promql_utilities::data_model::KeyByLabelNames;
 use sql_utilities::sqlhelper::{SQLSchema, Table};
@@ -250,6 +250,18 @@ impl InferenceConfig {
                     .and_then(|v| v.as_str())
                     .ok_or_else(|| anyhow::anyhow!("Missing query field"))?
                     .to_string();
+                let planned_subquery = query_data
+                    .get("planned_subquery")
+                    .and_then(|v| v.as_str())
+                    .ok_or_else(|| anyhow::anyhow!("Missing planned_subquery field"))?
+                    .to_string();
+                let query_time_aggregations = query_data
+                    .get("query_time_aggregations")
+                    .ok_or_else(|| anyhow::anyhow!("Missing query_time_aggregations field"))
+                    .and_then(|value| {
+                        serde_yaml::from_value::<Vec<QueryTimeAggregation>>(value.clone())
+                            .map_err(anyhow::Error::from)
+                    })?;
 
                 let aggregations = if let Some(aggregations_data) =
                     query_data.get("aggregations").and_then(|v| v.as_sequence())
@@ -290,7 +302,12 @@ impl InferenceConfig {
                     Vec::new()
                 };
 
-                let config = QueryConfig::new(query).with_aggregations(aggregations);
+                let config =
+                    QueryConfig::with_plan(query, planned_subquery, query_time_aggregations)
+                        .with_aggregations(aggregations);
+                config
+                    .validate_execution_plan()
+                    .map_err(|error| anyhow::anyhow!("Invalid query execution plan: {error}"))?;
                 configs.push(config);
             }
             configs
@@ -300,3 +317,77 @@ impl InferenceConfig {
         Ok(query_configs)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn rejects_a_query_without_a_planned_subquery() {
+        let data: Value = serde_yaml::from_str(
+            r#"
+cleanup_policy:
+  name: no_cleanup
+metrics: {}
+queries:
+  - query: "sum(metric)"
+    query_time_aggregations: []
+    aggregations: []
+"#,
+        )
+        .unwrap();
+
+        let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql)
+            .expect_err("query plans must name their planned subquery");
+
+        assert!(error.to_string().contains("planned_subquery"));
+    }
+
+    #[test]
+    fn rejects_a_query_without_a_query_time_pipeline() {
+        let data: Value = serde_yaml::from_str(
+            r#"
+cleanup_policy:
+  name: no_cleanup
+metrics: {}
+queries:
+  - query: "sum(metric)"
+    planned_subquery: "sum(metric)"
+    aggregations: []
+"#,
+        )
+        .unwrap();
+
+        let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql)
+            .expect_err("query plans must declare their query-time pipeline");
+
+        assert!(error.to_string().contains("query_time_aggregations"));
+    }
+
+    #[test]
+    fn rejects_an_invalid_query_time_aggregation() {
+        let data: Value = serde_yaml::from_str(
+            r#"
+cleanup_policy:
+  name: no_cleanup
+metrics: {}
+queries:
+  - query: "quantile(1.5, sum(metric))"
+    planned_subquery: "sum(metric)"
+    query_time_aggregations:
+      - operator: quantile
+        grouping:
+          mode: all
+          labels: []
+        parameter: 1.5
+    aggregations: []
+"#,
+        )
+        .unwrap();
+
+        let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql)
+            .expect_err("invalid pipeline parameters must fail during config loading");
+
+        assert!(error.to_string().contains("quantile"));
+    }
+}

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

@@ -5,17 +5,133 @@ use crate::aggregation_reference::AggregationReference;
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct QueryConfig {
     pub query: String,
+    pub planned_subquery: String,
+    pub query_time_aggregations: Vec<QueryTimeAggregation>,
     pub aggregations: Vec<AggregationReference>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum QueryTimeAggregationOperator {
+    Sum,
+    Count,
+    Avg,
+    Min,
+    Max,
+    Quantile,
+    Topk,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum QueryTimeGroupingMode {
+    All,
+    By,
+    Without,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct QueryTimeGrouping {
+    pub mode: QueryTimeGroupingMode,
+    pub labels: Vec<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum QueryTimeAggregationParameter {
+    Integer(u64),
+    Float(f64),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct QueryTimeAggregation {
+    pub operator: QueryTimeAggregationOperator,
+    pub grouping: QueryTimeGrouping,
+    pub parameter: Option<QueryTimeAggregationParameter>,
+}
+
+impl QueryTimeAggregation {
+    pub fn validate(&self) -> Result<(), String> {
+        match self.grouping.mode {
+            QueryTimeGroupingMode::All if !self.grouping.labels.is_empty() => {
+                return Err("all grouping cannot name labels".to_string());
+            }
+            QueryTimeGroupingMode::By if self.grouping.labels.is_empty() => {
+                return Err("by grouping must name at least one label".to_string());
+            }
+            _ => {}
+        }
+
+        if self.grouping.labels.iter().any(|label| label.is_empty()) {
+            return Err("grouping labels cannot be empty".to_string());
+        }
+        let unique_label_count = self
+            .grouping
+            .labels
+            .iter()
+            .collect::<std::collections::HashSet<_>>()
+            .len();
+        if unique_label_count != self.grouping.labels.len() {
+            return Err("grouping labels must be unique".to_string());
+        }
+
+        match (&self.operator, &self.parameter) {
+            (
+                QueryTimeAggregationOperator::Topk,
+                Some(QueryTimeAggregationParameter::Integer(k)),
+            ) if *k > 0 => {}
+            (
+                QueryTimeAggregationOperator::Topk,
+                Some(QueryTimeAggregationParameter::Integer(_)),
+            ) => {
+                return Err("topk requires a positive integer parameter".to_string());
+            }
+            (QueryTimeAggregationOperator::Topk, _) => {
+                return Err("topk requires an integer parameter".to_string());
+            }
+            (
+                QueryTimeAggregationOperator::Quantile,
+                Some(QueryTimeAggregationParameter::Float(phi)),
+            ) if phi.is_finite() && (0.0..=1.0).contains(phi) => {}
+            (QueryTimeAggregationOperator::Quantile, _) => {
+                return Err("quantile requires a finite parameter between 0 and 1".to_string());
+            }
+            (_, None) => {}
+            _ => return Err("this aggregation does not accept a parameter".to_string()),
+        }
+
+        Ok(())
+    }
+}
+
 impl QueryConfig {
     pub fn new(query: String) -> Self {
+        Self::with_plan(query.clone(), query, Vec::new())
+    }
+
+    pub fn with_plan(
+        query: String,
+        planned_subquery: String,
+        query_time_aggregations: Vec<QueryTimeAggregation>,
+    ) -> Self {
         Self {
             query,
+            planned_subquery,
+            query_time_aggregations,
             aggregations: Vec::new(),
         }
     }
 
+    pub fn validate_execution_plan(&self) -> Result<(), String> {
+        if self.planned_subquery.trim().is_empty() {
+            return Err("planned_subquery cannot be empty".to_string());
+        }
+        for aggregation in &self.query_time_aggregations {
+            aggregation.validate()?;
+        }
+        Ok(())
+    }
+
     pub fn add_aggregation(mut self, aggregation: AggregationReference) -> Self {
         self.aggregations.push(aggregation);
         self
@@ -26,3 +142,68 @@ impl QueryConfig {
         self
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn topk(k: u64) -> QueryTimeAggregation {
+        QueryTimeAggregation {
+            operator: QueryTimeAggregationOperator::Topk,
+            grouping: QueryTimeGrouping {
+                mode: QueryTimeGroupingMode::All,
+                labels: Vec::new(),
+            },
+            parameter: Some(QueryTimeAggregationParameter::Integer(k)),
+        }
+    }
+
+    #[test]
+    fn topk_requires_a_positive_k() {
+        assert!(topk(1).validate().is_ok());
+        assert_eq!(
+            topk(0).validate(),
+            Err("topk requires a positive integer parameter".into())
+        );
+    }
+
+    #[test]
+    fn quantile_accepts_endpoints_and_rejects_out_of_range_values() {
+        for phi in [0.0, 1.0] {
+            assert!(QueryTimeAggregation {
+                operator: QueryTimeAggregationOperator::Quantile,
+                grouping: QueryTimeGrouping {
+                    mode: QueryTimeGroupingMode::All,
+                    labels: Vec::new(),
+                },
+                parameter: Some(QueryTimeAggregationParameter::Float(phi)),
+            }
+            .validate()
+            .is_ok());
+        }
+        assert!(QueryTimeAggregation {
+            operator: QueryTimeAggregationOperator::Quantile,
+            grouping: QueryTimeGrouping {
+                mode: QueryTimeGroupingMode::All,
+                labels: Vec::new(),
+            },
+            parameter: Some(QueryTimeAggregationParameter::Float(1.01)),
+        }
+        .validate()
+        .is_err());
+    }
+
+    #[test]
+    fn without_empty_labels_is_a_valid_promql_grouping() {
+        assert!(QueryTimeAggregation {
+            operator: QueryTimeAggregationOperator::Sum,
+            grouping: QueryTimeGrouping {
+                mode: QueryTimeGroupingMode::Without,
+                labels: Vec::new(),
+            },
+            parameter: None,
+        }
+        .validate()
+        .is_ok());
+    }
+}

asap-planner-rs/src/elastic_dsl/generator.rs

@@ -7,7 +7,7 @@ use std::collections::HashMap;
 use crate::config::input::ElasticDSLControllerConfig;
 use crate::error::ControllerError;
 use crate::generator::{
-    build_aggregation_entry, build_queries_yaml, GeneratorOutput, KEY_AGGREGATIONS,
+    build_aggregation_entry, build_queries_yaml, GeneratorOutput, QueryPlanEntry, KEY_AGGREGATIONS,
     KEY_CLEANUP_POLICY, KEY_NAME, KEY_QUERIES,
 };
 use crate::planner::agg_config::IntermediateAggConfig;
@@ -82,8 +82,8 @@ pub fn generate_elastic_plan(
 
     // Dedup map: identifying_key -> IntermediateAggConfig
     let mut dedup_map: IndexMap<String, IntermediateAggConfig> = IndexMap::new();
-    // query_string -> Vec<(key, cleanup_param)>
-    let mut query_keys_map: IndexMap<String, Vec<(String, Option<u64>)>> = IndexMap::new();
+    // query_string -> explicit physical query plan
+    let mut query_plan_map: IndexMap<String, QueryPlanEntry> = IndexMap::new();
     // index -> schema builder derived from the queries targeting that index
     let mut index_schema_builders: IndexMap<String, ElasticIndexSchemaBuilder> = IndexMap::new();
 
@@ -127,7 +127,10 @@ pub fn generate_elastic_plan(
                 keys_for_query.push((key.clone(), cleanup_param));
                 dedup_map.entry(key).or_insert(config_item);
             }
-            query_keys_map.insert(query_string.clone(), keys_for_query);
+            query_plan_map.insert(
+                query_string.clone(),
+                QueryPlanEntry::fully_planned(query_string.clone(), keys_for_query),
+            );
         }
     }
 
@@ -140,7 +143,7 @@ pub fn generate_elastic_plan(
     let streaming_yaml = build_elastic_streaming_yaml(&dedup_map, &id_map)?;
     let inference_yaml = build_elastic_inference_yaml(
         cleanup_policy,
-        &query_keys_map,
+        &query_plan_map,
         &id_map,
         &index_schema_builders,
     )?;
@@ -150,7 +153,7 @@ pub fn generate_elastic_plan(
         streaming_yaml,
         inference_yaml,
         aggregation_count: dedup_map.len(),
-        query_count: query_keys_map.len(),
+        query_count: query_plan_map.len(),
     })
 }
 
@@ -174,7 +177,7 @@ fn build_elastic_streaming_yaml(
 
 fn build_elastic_inference_yaml(
     cleanup_policy: CleanupPolicy,
-    query_keys_map: &IndexMap<String, Vec<(String, Option<u64>)>>,
+    query_plan_map: &IndexMap<String, QueryPlanEntry>,
     id_map: &HashMap<String, u32>,
     index_schema_builders: &IndexMap<String, ElasticIndexSchemaBuilder>,
 ) -> Result<YamlValue, ControllerError> {
@@ -191,7 +194,7 @@ fn build_elastic_inference_yaml(
     );
     root.insert(
         YamlValue::String(KEY_QUERIES.to_string()),
-        YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_keys_map, id_map)),
+        YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_plan_map, id_map)),
     );
     root.insert(
         YamlValue::String("indices".to_string()),