Skip to content

Commit cbabfd6

Browse files
feat(query-engine): execute nested PromQL aggregation DAGs (#770)
* feat(planner): add explicit query execution plans * feat(planner): emit nested aggregation plans * feat(query-engine): validate query-time aggregation configs * test(planner): cover nested aggregation stages * feat(query-engine): execute query-time aggregation pipelines * test(query-engine): add nested aggregation e2e tracer * test(query-engine): cover nested topk e2e * test(query-engine): cover nested aggregation operators e2e * test(query-engine): assert nested aggregation e2e values * test(promql): add nested aggregation differential suite * test(query-engine): cover query-time aggregation grouping * fix(query-config): reject zero topk limits * test(query-engine): cover query-time aggregation boundaries * test(query-engine): cover query-time topk ties * feat(query-engine): execute nested aggregation DAG nodes * fix(query-engine): preserve missing query-time labels * fix(planner): preserve empty without aggregations * fix(query-engine): apply nested topk anchors first * fix(query-engine): validate nested aggregation plans * fix(query-engine): preserve nested topk labels * fix(query-engine): keep grouped topk buckets contiguous * fix(query-engine): align raw topk labels
1 parent 8132a74 commit cbabfd6

22 files changed

Lines changed: 3478 additions & 68 deletions

File tree

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

Lines changed: 93 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use std::io::BufReader;
77
use crate::aggregation_reference::AggregationReference;
88
use crate::enums::{CleanupPolicy, QueryLanguage};
99
use crate::promql_schema::PromQLSchema;
10-
use crate::query_config::QueryConfig;
10+
use crate::query_config::{QueryConfig, QueryTimeAggregation};
1111
use elastic_dsl_utilities::{ElasticIndexSchema, ElasticMappingSchema};
1212
use promql_utilities::data_model::KeyByLabelNames;
1313
use sql_utilities::sqlhelper::{SQLSchema, Table};
@@ -250,6 +250,18 @@ impl InferenceConfig {
250250
.and_then(|v| v.as_str())
251251
.ok_or_else(|| anyhow::anyhow!("Missing query field"))?
252252
.to_string();
253+
let planned_subquery = query_data
254+
.get("planned_subquery")
255+
.and_then(|v| v.as_str())
256+
.ok_or_else(|| anyhow::anyhow!("Missing planned_subquery field"))?
257+
.to_string();
258+
let query_time_aggregations = query_data
259+
.get("query_time_aggregations")
260+
.ok_or_else(|| anyhow::anyhow!("Missing query_time_aggregations field"))
261+
.and_then(|value| {
262+
serde_yaml::from_value::<Vec<QueryTimeAggregation>>(value.clone())
263+
.map_err(anyhow::Error::from)
264+
})?;
253265

254266
let aggregations = if let Some(aggregations_data) =
255267
query_data.get("aggregations").and_then(|v| v.as_sequence())
@@ -290,7 +302,12 @@ impl InferenceConfig {
290302
Vec::new()
291303
};
292304

293-
let config = QueryConfig::new(query).with_aggregations(aggregations);
305+
let config =
306+
QueryConfig::with_plan(query, planned_subquery, query_time_aggregations)
307+
.with_aggregations(aggregations);
308+
config
309+
.validate_execution_plan()
310+
.map_err(|error| anyhow::anyhow!("Invalid query execution plan: {error}"))?;
294311
configs.push(config);
295312
}
296313
configs
@@ -300,3 +317,77 @@ impl InferenceConfig {
300317
Ok(query_configs)
301318
}
302319
}
320+
321+
#[cfg(test)]
322+
mod tests {
323+
use super::*;
324+
325+
#[test]
326+
fn rejects_a_query_without_a_planned_subquery() {
327+
let data: Value = serde_yaml::from_str(
328+
r#"
329+
cleanup_policy:
330+
name: no_cleanup
331+
metrics: {}
332+
queries:
333+
- query: "sum(metric)"
334+
query_time_aggregations: []
335+
aggregations: []
336+
"#,
337+
)
338+
.unwrap();
339+
340+
let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql)
341+
.expect_err("query plans must name their planned subquery");
342+
343+
assert!(error.to_string().contains("planned_subquery"));
344+
}
345+
346+
#[test]
347+
fn rejects_a_query_without_a_query_time_pipeline() {
348+
let data: Value = serde_yaml::from_str(
349+
r#"
350+
cleanup_policy:
351+
name: no_cleanup
352+
metrics: {}
353+
queries:
354+
- query: "sum(metric)"
355+
planned_subquery: "sum(metric)"
356+
aggregations: []
357+
"#,
358+
)
359+
.unwrap();
360+
361+
let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql)
362+
.expect_err("query plans must declare their query-time pipeline");
363+
364+
assert!(error.to_string().contains("query_time_aggregations"));
365+
}
366+
367+
#[test]
368+
fn rejects_an_invalid_query_time_aggregation() {
369+
let data: Value = serde_yaml::from_str(
370+
r#"
371+
cleanup_policy:
372+
name: no_cleanup
373+
metrics: {}
374+
queries:
375+
- query: "quantile(1.5, sum(metric))"
376+
planned_subquery: "sum(metric)"
377+
query_time_aggregations:
378+
- operator: quantile
379+
grouping:
380+
mode: all
381+
labels: []
382+
parameter: 1.5
383+
aggregations: []
384+
"#,
385+
)
386+
.unwrap();
387+
388+
let error = InferenceConfig::from_yaml_data(&data, QueryLanguage::promql)
389+
.expect_err("invalid pipeline parameters must fail during config loading");
390+
391+
assert!(error.to_string().contains("quantile"));
392+
}
393+
}

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

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,133 @@ use crate::aggregation_reference::AggregationReference;
55
#[derive(Debug, Clone, Serialize, Deserialize)]
66
pub struct QueryConfig {
77
pub query: String,
8+
pub planned_subquery: String,
9+
pub query_time_aggregations: Vec<QueryTimeAggregation>,
810
pub aggregations: Vec<AggregationReference>,
911
}
1012

13+
#[derive(Debug, Clone, Serialize, Deserialize)]
14+
#[serde(rename_all = "snake_case")]
15+
pub enum QueryTimeAggregationOperator {
16+
Sum,
17+
Count,
18+
Avg,
19+
Min,
20+
Max,
21+
Quantile,
22+
Topk,
23+
}
24+
25+
#[derive(Debug, Clone, Serialize, Deserialize)]
26+
#[serde(rename_all = "snake_case")]
27+
pub enum QueryTimeGroupingMode {
28+
All,
29+
By,
30+
Without,
31+
}
32+
33+
#[derive(Debug, Clone, Serialize, Deserialize)]
34+
pub struct QueryTimeGrouping {
35+
pub mode: QueryTimeGroupingMode,
36+
pub labels: Vec<String>,
37+
}
38+
39+
#[derive(Debug, Clone, Serialize, Deserialize)]
40+
#[serde(untagged)]
41+
pub enum QueryTimeAggregationParameter {
42+
Integer(u64),
43+
Float(f64),
44+
}
45+
46+
#[derive(Debug, Clone, Serialize, Deserialize)]
47+
pub struct QueryTimeAggregation {
48+
pub operator: QueryTimeAggregationOperator,
49+
pub grouping: QueryTimeGrouping,
50+
pub parameter: Option<QueryTimeAggregationParameter>,
51+
}
52+
53+
impl QueryTimeAggregation {
54+
pub fn validate(&self) -> Result<(), String> {
55+
match self.grouping.mode {
56+
QueryTimeGroupingMode::All if !self.grouping.labels.is_empty() => {
57+
return Err("all grouping cannot name labels".to_string());
58+
}
59+
QueryTimeGroupingMode::By if self.grouping.labels.is_empty() => {
60+
return Err("by grouping must name at least one label".to_string());
61+
}
62+
_ => {}
63+
}
64+
65+
if self.grouping.labels.iter().any(|label| label.is_empty()) {
66+
return Err("grouping labels cannot be empty".to_string());
67+
}
68+
let unique_label_count = self
69+
.grouping
70+
.labels
71+
.iter()
72+
.collect::<std::collections::HashSet<_>>()
73+
.len();
74+
if unique_label_count != self.grouping.labels.len() {
75+
return Err("grouping labels must be unique".to_string());
76+
}
77+
78+
match (&self.operator, &self.parameter) {
79+
(
80+
QueryTimeAggregationOperator::Topk,
81+
Some(QueryTimeAggregationParameter::Integer(k)),
82+
) if *k > 0 => {}
83+
(
84+
QueryTimeAggregationOperator::Topk,
85+
Some(QueryTimeAggregationParameter::Integer(_)),
86+
) => {
87+
return Err("topk requires a positive integer parameter".to_string());
88+
}
89+
(QueryTimeAggregationOperator::Topk, _) => {
90+
return Err("topk requires an integer parameter".to_string());
91+
}
92+
(
93+
QueryTimeAggregationOperator::Quantile,
94+
Some(QueryTimeAggregationParameter::Float(phi)),
95+
) if phi.is_finite() && (0.0..=1.0).contains(phi) => {}
96+
(QueryTimeAggregationOperator::Quantile, _) => {
97+
return Err("quantile requires a finite parameter between 0 and 1".to_string());
98+
}
99+
(_, None) => {}
100+
_ => return Err("this aggregation does not accept a parameter".to_string()),
101+
}
102+
103+
Ok(())
104+
}
105+
}
106+
11107
impl QueryConfig {
12108
pub fn new(query: String) -> Self {
109+
Self::with_plan(query.clone(), query, Vec::new())
110+
}
111+
112+
pub fn with_plan(
113+
query: String,
114+
planned_subquery: String,
115+
query_time_aggregations: Vec<QueryTimeAggregation>,
116+
) -> Self {
13117
Self {
14118
query,
119+
planned_subquery,
120+
query_time_aggregations,
15121
aggregations: Vec::new(),
16122
}
17123
}
18124

125+
pub fn validate_execution_plan(&self) -> Result<(), String> {
126+
if self.planned_subquery.trim().is_empty() {
127+
return Err("planned_subquery cannot be empty".to_string());
128+
}
129+
for aggregation in &self.query_time_aggregations {
130+
aggregation.validate()?;
131+
}
132+
Ok(())
133+
}
134+
19135
pub fn add_aggregation(mut self, aggregation: AggregationReference) -> Self {
20136
self.aggregations.push(aggregation);
21137
self
@@ -26,3 +142,68 @@ impl QueryConfig {
26142
self
27143
}
28144
}
145+
146+
#[cfg(test)]
147+
mod tests {
148+
use super::*;
149+
150+
fn topk(k: u64) -> QueryTimeAggregation {
151+
QueryTimeAggregation {
152+
operator: QueryTimeAggregationOperator::Topk,
153+
grouping: QueryTimeGrouping {
154+
mode: QueryTimeGroupingMode::All,
155+
labels: Vec::new(),
156+
},
157+
parameter: Some(QueryTimeAggregationParameter::Integer(k)),
158+
}
159+
}
160+
161+
#[test]
162+
fn topk_requires_a_positive_k() {
163+
assert!(topk(1).validate().is_ok());
164+
assert_eq!(
165+
topk(0).validate(),
166+
Err("topk requires a positive integer parameter".into())
167+
);
168+
}
169+
170+
#[test]
171+
fn quantile_accepts_endpoints_and_rejects_out_of_range_values() {
172+
for phi in [0.0, 1.0] {
173+
assert!(QueryTimeAggregation {
174+
operator: QueryTimeAggregationOperator::Quantile,
175+
grouping: QueryTimeGrouping {
176+
mode: QueryTimeGroupingMode::All,
177+
labels: Vec::new(),
178+
},
179+
parameter: Some(QueryTimeAggregationParameter::Float(phi)),
180+
}
181+
.validate()
182+
.is_ok());
183+
}
184+
assert!(QueryTimeAggregation {
185+
operator: QueryTimeAggregationOperator::Quantile,
186+
grouping: QueryTimeGrouping {
187+
mode: QueryTimeGroupingMode::All,
188+
labels: Vec::new(),
189+
},
190+
parameter: Some(QueryTimeAggregationParameter::Float(1.01)),
191+
}
192+
.validate()
193+
.is_err());
194+
}
195+
196+
#[test]
197+
fn without_empty_labels_is_a_valid_promql_grouping() {
198+
assert!(QueryTimeAggregation {
199+
operator: QueryTimeAggregationOperator::Sum,
200+
grouping: QueryTimeGrouping {
201+
mode: QueryTimeGroupingMode::Without,
202+
labels: Vec::new(),
203+
},
204+
parameter: None,
205+
}
206+
.validate()
207+
.is_ok());
208+
}
209+
}

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

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use std::collections::HashMap;
77
use crate::config::input::ElasticDSLControllerConfig;
88
use crate::error::ControllerError;
99
use crate::generator::{
10-
build_aggregation_entry, build_queries_yaml, GeneratorOutput, KEY_AGGREGATIONS,
10+
build_aggregation_entry, build_queries_yaml, GeneratorOutput, QueryPlanEntry, KEY_AGGREGATIONS,
1111
KEY_CLEANUP_POLICY, KEY_NAME, KEY_QUERIES,
1212
};
1313
use crate::planner::agg_config::IntermediateAggConfig;
@@ -82,8 +82,8 @@ pub fn generate_elastic_plan(
8282

8383
// Dedup map: identifying_key -> IntermediateAggConfig
8484
let mut dedup_map: IndexMap<String, IntermediateAggConfig> = IndexMap::new();
85-
// query_string -> Vec<(key, cleanup_param)>
86-
let mut query_keys_map: IndexMap<String, Vec<(String, Option<u64>)>> = IndexMap::new();
85+
// query_string -> explicit physical query plan
86+
let mut query_plan_map: IndexMap<String, QueryPlanEntry> = IndexMap::new();
8787
// index -> schema builder derived from the queries targeting that index
8888
let mut index_schema_builders: IndexMap<String, ElasticIndexSchemaBuilder> = IndexMap::new();
8989

@@ -127,7 +127,10 @@ pub fn generate_elastic_plan(
127127
keys_for_query.push((key.clone(), cleanup_param));
128128
dedup_map.entry(key).or_insert(config_item);
129129
}
130-
query_keys_map.insert(query_string.clone(), keys_for_query);
130+
query_plan_map.insert(
131+
query_string.clone(),
132+
QueryPlanEntry::fully_planned(query_string.clone(), keys_for_query),
133+
);
131134
}
132135
}
133136

@@ -140,7 +143,7 @@ pub fn generate_elastic_plan(
140143
let streaming_yaml = build_elastic_streaming_yaml(&dedup_map, &id_map)?;
141144
let inference_yaml = build_elastic_inference_yaml(
142145
cleanup_policy,
143-
&query_keys_map,
146+
&query_plan_map,
144147
&id_map,
145148
&index_schema_builders,
146149
)?;
@@ -150,7 +153,7 @@ pub fn generate_elastic_plan(
150153
streaming_yaml,
151154
inference_yaml,
152155
aggregation_count: dedup_map.len(),
153-
query_count: query_keys_map.len(),
156+
query_count: query_plan_map.len(),
154157
})
155158
}
156159

@@ -174,7 +177,7 @@ fn build_elastic_streaming_yaml(
174177

175178
fn build_elastic_inference_yaml(
176179
cleanup_policy: CleanupPolicy,
177-
query_keys_map: &IndexMap<String, Vec<(String, Option<u64>)>>,
180+
query_plan_map: &IndexMap<String, QueryPlanEntry>,
178181
id_map: &HashMap<String, u32>,
179182
index_schema_builders: &IndexMap<String, ElasticIndexSchemaBuilder>,
180183
) -> Result<YamlValue, ControllerError> {
@@ -191,7 +194,7 @@ fn build_elastic_inference_yaml(
191194
);
192195
root.insert(
193196
YamlValue::String(KEY_QUERIES.to_string()),
194-
YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_keys_map, id_map)),
197+
YamlValue::Sequence(build_queries_yaml(cleanup_policy, query_plan_map, id_map)),
195198
);
196199
root.insert(
197200
YamlValue::String("indices".to_string()),

0 commit comments

Comments
 (0)