diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs index a1fa5144..f255dcef 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs @@ -312,6 +312,22 @@ impl AggregationType { } } + /// Single-subpopulation types whose merge of two groups' states is the + /// state of their union, so a query can read one at a coarser grouping + /// (a roll-up). Increase is excluded: its merge joins pieces of one + /// counter in time. + pub const MERGEABLE_ACROSS_GROUPS: [AggregationType; 5] = [ + AggregationType::Sum, + AggregationType::MinMax, + AggregationType::DatasketchesKLL, + AggregationType::DDSketch, + AggregationType::HLL, + ]; + + pub fn merges_across_groups(self) -> bool { + Self::MERGEABLE_ACROSS_GROUPS.contains(&self) + } + /// Returns `true` if this type produces keyed (multi-population) accumulators. pub fn is_keyed(self) -> bool { matches!( diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index 1b282d6b..74428090 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -8,7 +8,7 @@ use indexmap::IndexMap; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::AggregationType; use rqe_optimizer::milp::MilpSolution; -use rqe_optimizer::{Capability, Deployment}; +use rqe_optimizer::{Capability, Deployment, LabelSet, Raqe}; use serde_json::Value; use thiserror::Error; @@ -48,10 +48,22 @@ pub enum MilpOutputError { #[error("query {0:?} is served by two deployments; a query string can name only one")] QueryOnTwoDeployments(String), #[error( - "query {query:?} is grouped by {raqe:?} but served by a deployment grouped by \ - {deployment:?}: a roll-up, which the engine can't read yet" + "query {query:?} is grouped by {raqe:?} but served by a {aggregation_type} deployment \ + grouped by {deployment:?}: only aggregations that merge across groups roll up \ + ({mergeable:?})", + mergeable = AggregationType::MERGEABLE_ACROSS_GROUPS )] RollupUnsupported { + query: String, + aggregation_type: AggregationType, + raqe: Vec, + deployment: Vec, + }, + #[error( + "query {query:?} is grouped by {raqe:?}, which is not a subset of its \ + deployment's grouping {deployment:?}" + )] + GroupingOutsideDeployment { query: String, raqe: Vec, deployment: Vec, @@ -117,20 +129,8 @@ pub fn deployments_to_planner_output( "one deployment per Raqe" ); reject_unwritable_queries(config)?; - // sketch-bench lets a deployment serve a RAQE grouped by a subset of its - // grouping (a roll-up, sketch-bench#190). The engine reads a deployment - // only at its own grouping, and `aggregation_config` takes the labels - // from the first RAQE served, so refuse a roll-up rather than write - // configs at the wrong grouping. for (raqe, &index) in workload.raqes.iter().zip(served_by) { - let deployment = &deployments[index]; - if raqe.grouping_labels != deployment.grouping_labels { - return Err(MilpOutputError::RollupUnsupported { - query: raqe.id.clone(), - raqe: raqe.grouping_labels.iter().cloned().collect(), - deployment: deployment.grouping_labels.iter().cloned().collect(), - }); - } + check_grouping(raqe, &deployments[index])?; } // A SQL aggregation names its table and value column, and its metric is @@ -169,6 +169,11 @@ pub fn deployments_to_planner_output( for (raqe, &index) in served_by.iter().enumerate() { let item = workload.raqe_items[raqe]; let key = &item_keys[&(index, item)]; + check_rollup( + &workload.raqes[raqe], + &deployments[index], + aggregations[key].aggregation_type, + )?; for query in &workload.items[item].query_strings { match queries.entry(query.clone()) { Entry::Vacant(entry) => { @@ -212,9 +217,46 @@ pub fn deployments_to_planner_output( })) } +fn label_names(set: &LabelSet) -> Vec { + set.iter().cloned().collect() +} + +/// A deployment can serve a Raqe grouped by its own grouping or a subset of +/// it (a roll-up), never one grouped by labels it doesn't keep. +fn check_grouping(raqe: &Raqe, deployment: &Deployment) -> Result<(), MilpOutputError> { + if raqe.grouping_labels.is_subset(&deployment.grouping_labels) { + return Ok(()); + } + Err(MilpOutputError::GroupingOutsideDeployment { + query: raqe.id.clone(), + raqe: label_names(&raqe.grouping_labels), + deployment: label_names(&deployment.grouping_labels), + }) +} + +/// A Raqe grouped by a strict subset of its deployment's grouping is a +/// roll-up. The engine merges such aggregations across the deployment's +/// groups at query time, so those need only a query_config naming the +/// aggregation. +fn check_rollup( + raqe: &Raqe, + deployment: &Deployment, + aggregation_type: AggregationType, +) -> Result<(), MilpOutputError> { + if raqe.grouping_labels == deployment.grouping_labels || aggregation_type.merges_across_groups() + { + return Ok(()); + } + Err(MilpOutputError::RollupUnsupported { + query: raqe.id.clone(), + aggregation_type, + raqe: label_names(&raqe.grouping_labels), + deployment: label_names(&deployment.grouping_labels), + }) +} + /// The aggregation `deployment` runs for `item`, one of the optimizer items -/// it serves. Items share the deployment's capability and grouping -/// (`plan_to_planner_output` refuses roll-ups); `item` decides the labels. +/// it serves. fn aggregation_config( deployment: &Deployment, item: &OptimizerItem, @@ -283,7 +325,8 @@ fn aggregation_config( } }; - let (grouping_labels, aggregated_labels) = subpopulation_labels(item, aggregation_type); + let (grouping_labels, aggregated_labels) = + subpopulation_labels(deployment, item, aggregation_type); Ok(IntermediateAggConfig { aggregation_type, @@ -309,20 +352,34 @@ fn aggregation_config( }) } -/// The (grouping, aggregated) labels `item`'s query needs from an aggregation -/// of `aggregation_type`. +/// The (grouping, aggregated) labels `item`'s query needs from `deployment`. +/// A roll-up query is grouped by a subset of the deployment's grouping, so a +/// non-top-k aggregation is always built at the deployment's grouping. A +/// top-k heap is bucketed and keyed by the item's own labels (its `topk by` +/// buckets; all labels, or a topk-of-sum's inner `sum by` labels), which can +/// differ between items on one deployment. Top-k never rolls up. fn subpopulation_labels( + deployment: &Deployment, item: &OptimizerItem, aggregation_type: AggregationType, ) -> (KeyByLabelNames, KeyByLabelNames) { let requirements = &item.requirements; + let deployment_labels = + KeyByLabelNames::new(deployment.grouping_labels.iter().cloned().collect()); + let (subpopulation_labels, topk_by_labels) = match deployment.capability { + Capability::TopKByValue | Capability::TopKByCount => ( + &requirements.grouping_labels, + requirements.topk_by_labels.as_ref(), + ), + _ => (&deployment_labels, None), + }; let mut grouping_labels = KeyByLabelNames::empty(); let mut aggregated_labels = KeyByLabelNames::empty(); set_subpopulation_labels( requirements.statistics[0], aggregation_type, - &requirements.grouping_labels, - requirements.topk_by_labels.as_ref(), + subpopulation_labels, + topk_by_labels, &mut KeyByLabelNames::empty(), &mut grouping_labels, &mut aggregated_labels, @@ -603,6 +660,50 @@ mod tests { } } + /// Cardinality only comes from SQL `COUNT(DISTINCT ...)`, so the HLL + /// roll-up is planned from a real SQL cardinality query. + #[test] + fn hll_count_distinct_rollup_is_deployed_at_the_deployment_grouping() { + let query = "SELECT job, COUNT(DISTINCT requests) FROM http WHERE \ + time BETWEEN DATEADD(s, -60, NOW()) AND NOW() GROUP BY job"; + let sql: SQLControllerConfig = serde_yaml::from_str(&format!( + "tables:\n - name: http\n time_column: time\n value_columns: [requests]\n \ + metadata_columns: [instance, job]\nquery_groups:\n{}", + group(query, 0.99) + )) + .unwrap(); + let facts = parse_workload_facts( + &FACTS.replace("http_requests_total", "http.requests"), + &MilpConfig::Sql(&sql).hints().unwrap(), + SCRAPE_MS, + ) + .unwrap(); + let workload = build_milp_workload(MilpConfig::Sql(&sql), &facts, SCRAPE_MS).unwrap(); + let raqe = &workload.raqes[0]; + assert_eq!(raqe.capability, Capability::Cardinality); + assert_eq!(raqe.grouping_labels, LabelSet::from(["job".into()])); + // The test cost table has no HLL rows to solve with, so reuse a + // solved deployment's shape as an HLL deployment on the table. + let (_, _, mut solution) = plan(&group( + "quantile by (job) (0.99, http_requests_total)", + 0.99, + )); + widen_to_instance_job(&mut solution); + let mut deployment = solution.deployments[0].deployment.clone(); + deployment.metric = raqe.metric.clone(); + deployment.capability = Capability::Cardinality; + deployment.config.sketch = "hll".into(); + deployment.config.sketch_config["params"] = json!({"lg_k": 12}); + + let output = + deployments_to_planner_output(MilpConfig::Sql(&sql), &workload, &[deployment], &[0]) + .unwrap(); + let a = aggregation_in(&output, query, QueryLanguage::sql); + assert_eq!(a.aggregation_type, AggregationType::HLL); + assert_eq!(a.grouping_labels, instance_job()); + assert!(a.aggregated_labels.is_empty()); + } + /// The heap `query`'s deployment needs for it: `m ยท k` for the windows /// its lookback merges and its own `k` (`Deployment::heap_needed`). fn heap_needed(workload: &MilpWorkload, solution: &MilpSolution, query: &str) -> u64 { @@ -823,22 +924,72 @@ mod tests { ); } + fn instance_job() -> KeyByLabelNames { + KeyByLabelNames::new(vec!["instance".into(), "job".into()]) + } + + /// The solution with its only deployment grouped by (instance, job). + fn widen_to_instance_job(solution: &mut MilpSolution) { + let [planned] = solution.deployments.as_mut_slice() else { + panic!("expected one deployment, got {:?}", solution.deployments); + }; + planned.deployment.grouping_labels = LabelSet::from(["instance".into(), "job".into()]); + } + + /// A sketch `by (job)` served from a deployment by (instance, job) is + /// written at the deployment's grouping; the engine merges across + /// `instance` at query time. + #[test] + fn sketch_rollup_is_deployed_at_the_deployment_grouping() { + let query = "quantile by (job) (0.99, http_requests_total)"; + for (variant, capability, params, aggregation_type) in [ + ( + "kll-percall", + Capability::Quantile, + json!({"k": 200}), + AggregationType::DatasketchesKLL, + ), + ( + "dd", + Capability::Quantile, + json!({"alpha": 0.01}), + AggregationType::DDSketch, + ), + ] { + let (config, workload, mut solution) = plan(&group(query, 0.99)); + assert_eq!( + workload.raqes[0].grouping_labels, + LabelSet::from(["job".into()]) + ); + widen_to_instance_job(&mut solution); + let deployment = &mut solution.deployments[0].deployment; + deployment.config.sketch = variant.into(); + deployment.config.sketch_config["params"] = params; + deployment.capability = capability; + + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); + assert_eq!(output.streaming_aggregation_count(), 1, "{variant}"); + let a = aggregation_for(&output, query); + assert_eq!(a.aggregation_type, aggregation_type, "{variant}"); + assert_eq!(a.grouping_labels, instance_job(), "{variant}"); + assert!(a.aggregated_labels.is_empty(), "{variant}"); + } + } + + /// Exact accumulators are keyed types the engine doesn't roll up. #[test] - fn rollup_is_refused() { - // A deployment by (instance, job) serving a query by (job): the - // engine would build the aggregation at the query's grouping. + fn exact_sum_rollup_is_refused() { let query = "sum by (job) (http_requests_total)"; let (config, workload, mut solution) = plan(&group(query, 0.99)); - solution.deployments[0] - .deployment - .grouping_labels - .insert("instance".into()); + widen_to_instance_job(&mut solution); let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) .err() .unwrap(); assert!( - matches!(&err, MilpOutputError::RollupUnsupported { query: q, raqe, deployment } - if q.contains(query) && raqe == &["job"] && deployment == &["instance", "job"]), + matches!(&err, MilpOutputError::RollupUnsupported { + query: q, aggregation_type: AggregationType::MultipleSum, raqe, deployment + } if q.contains(query) && raqe == &["job"] && deployment == &["instance", "job"]), "{err}" ); } @@ -871,6 +1022,64 @@ mod tests { } } + #[test] + fn grouping_outside_the_deployment_is_refused() { + let query = "quantile by (job) (0.99, http_requests_total)"; + let (config, workload, mut solution) = plan(&group(query, 0.99)); + let grouping = &mut solution.deployments[0].deployment.grouping_labels; + *grouping = LabelSet::from(["instance".into()]); + let err = plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution) + .err() + .unwrap(); + assert!( + matches!(&err, MilpOutputError::GroupingOutsideDeployment { raqe, deployment, .. } + if raqe == &["job"] && deployment == &["instance"]), + "{err}" + ); + } + + /// A deployment serving an exact-grouping query and a roll-up is one + /// aggregation at the deployment's grouping, whichever query comes first. + #[test] + fn shared_deployment_takes_its_own_grouping_in_either_order() { + let fine = "quantile by (instance, job) (0.99, http_requests_total)"; + let coarse = "quantile by (job) (0.99, http_requests_total)"; + for reverse in [false, true] { + let (config, mut workload, mut solution) = + plan(&(group(fine, 0.99) + &group(coarse, 0.99))); + // Raqes are sorted by query string; reversing the index-aligned + // lists puts the other query first. + if reverse { + workload.raqes.reverse(); + workload.raqe_items.reverse(); + solution.raqes.reverse(); + } + let fine_raqe = workload + .raqes + .iter() + .position(|r| r.id.contains(fine)) + .unwrap(); + let shared = solution.deployments[solution.raqes[fine_raqe].deployment].clone(); + solution.deployments = vec![shared]; + for planned in &mut solution.raqes { + planned.deployment = 0; + } + + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); + assert_eq!(output.streaming_aggregation_count(), 1); + for query in [fine, coarse] { + let a = aggregation_for(&output, query); + assert_eq!(a.aggregation_type, AggregationType::DatasketchesKLL); + assert_eq!( + a.grouping_labels, + instance_job(), + "{query}, reverse {reverse}" + ); + } + } + } + #[test] fn query_split_across_deployments_is_an_error() { // The same query string at two cadences is two items; if they land on diff --git a/asap-planner-rs/src/optimizer/saved_plan.rs b/asap-planner-rs/src/optimizer/saved_plan.rs index 4a75e352..654034e9 100644 --- a/asap-planner-rs/src/optimizer/saved_plan.rs +++ b/asap-planner-rs/src/optimizer/saved_plan.rs @@ -391,6 +391,47 @@ mod tests { assert!(yaml.contains("alpha: 0.02"), "{yaml}"); } + /// A DDSketch grouped by (job, label_0) serves `quantile by (job)` as a + /// roll-up: one aggregation at the deployment's grouping. + #[test] + fn writes_a_sketch_roll_up_at_the_deployment_grouping() { + let table = r#"{"workloads": [{"e2e": { + "queries": [ + {"promql": "quantile by (job) (0.5, data)", "interval_secs": 1, "rqes": ["p50"]} + ], + "metrics": {"data": {"type": "gauge", "labels": [ + {"name": "job", "cardinality": 10}, + {"name": "label_0", "cardinality": 1000, "child_of": "job"} + ]}} + }}]}"#; + let mut p50 = pick("p50", 0, "dd", json!({"alpha": 0.02}), 1); + p50["deployment_grouping"] = json!(["job", "label_0"]); + let plan = json!({"plans": [ + {"method": "asap", "weights": "cpu", "kind": "cheapest", "chosen": [p50]}, + ]}); + + let output = run_on(table, plan.to_string().as_bytes(), None).unwrap(); + + let streaming = output + .to_streaming_config(asap_types::enums::QueryLanguage::promql) + .unwrap(); + let aggregations: Vec<_> = streaming.aggregation_configs.values().collect(); + let [aggregation] = aggregations.as_slice() else { + panic!("expected one aggregation, got {aggregations:?}"); + }; + assert_eq!( + aggregation.aggregation_type, + promql_utilities::query_logics::enums::AggregationType::DDSketch + ); + assert_eq!( + aggregation.grouping_labels, + promql_utilities::data_model::KeyByLabelNames::new(vec![ + "job".into(), + "label_0".into() + ]) + ); + } + #[test] fn sla_ms_selects_the_plan_at_that_sla() { let yaml = run(plans(shared_p50()), Some(100.0)) @@ -402,12 +443,13 @@ mod tests { assert!(err.contains("0 plans match"), "{err}"); } + /// The exact sum is a keyed MultipleSum, which the engine doesn't roll up. #[test] - fn refuses_a_roll_up() { + fn refuses_an_exact_sum_roll_up() { let mut plan = plans(shared_p50()); plan["plans"][0]["chosen"][0]["deployment_grouping"] = json!(["job", "label_0"]); let err = refusal(plan, None); - assert!(err.contains("roll-up"), "{err}"); + assert!(err.contains("merge across groups"), "{err}"); } #[test] diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index b66fec41..6b0d29aa 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -190,17 +190,6 @@ struct RangeQueryReads { keys: Option, } -/// Single-subpopulation types whose merge of two groups' states is the state -/// of their union. Increase is excluded: its merge joins pieces of one -/// counter in time. -const ROLLUP_AGGREGATION_TYPES: [AggregationType; 5] = [ - AggregationType::Sum, - AggregationType::MinMax, - AggregationType::DatasketchesKLL, - AggregationType::DDSketch, - AggregationType::HLL, -]; - type BucketMap = HashMap>>; type TopkPartition = (u64, Vec); type TopkCandidates<'a> = HashMap>; @@ -2586,7 +2575,7 @@ impl SimpleEngine { context.metric, output, context.agg_info.aggregation_id_for_value, grouping )) })?; - if !ROLLUP_AGGREGATION_TYPES.contains(&value_type) { + if !value_type.merges_across_groups() { return Err(QueryExecutionError::Native(format!( "query on metric {} is grouped by {:?}, a roll-up of aggregation {} (grouped by \ {:?}), but {} can't be merged across groups",