diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index b2e2e14d..2fcc3894 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -188,6 +188,17 @@ 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>; @@ -2602,10 +2613,83 @@ impl SimpleEngine { )) } + /// For a query grouped by labels `Q` served from a single-subpopulation + /// aggregation grouped by `G ⊋ Q` (a roll-up), the position in `G` of + /// each label of `Q`. `None` when the query reads the aggregation at its + /// own grouping. Keyed and top-k aggregations are never rolled up here. + fn rollup_projection( + context: &QueryExecutionContext, + ) -> Result>, QueryExecutionError> { + let value_type = context.agg_info.aggregation_type_for_value; + let single_subpopulation = !value_type.is_keyed() + && context.store_plan.keys_query.is_none() + && context.aggregated_labels.labels.is_empty() + && context.metadata.statistic_to_compute != Statistic::Topk; + let grouping = &context.grouping_labels.labels; + let output = &context.metadata.query_output_labels.labels; + if !single_subpopulation || grouping == output { + return Ok(None); + } + let projection = output + .iter() + .map(|label| grouping.iter().position(|g| g == label)) + .collect::>>() + .ok_or_else(|| { + QueryExecutionError::Native(format!( + "query on metric {} is grouped by {:?}, which aggregation {} (grouped by \ + {:?}) can't serve", + context.metric, output, context.agg_info.aggregation_id_for_value, grouping + )) + })?; + if !ROLLUP_AGGREGATION_TYPES.contains(&value_type) { + 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", + context.metric, + output, + context.agg_info.aggregation_id_for_value, + grouping, + value_type.as_str() + ))); + } + Ok(Some(projection)) + } + + /// Re-keys every store group by its projection onto the query's labels, + /// so the per-step merge folds all subpopulations of an output group. + fn roll_up_groups( + values: TimestampedBucketsMap, + projection: &[usize], + grouping_label_count: usize, + ) -> Result { + let mut rolled: TimestampedBucketsMap = HashMap::new(); + for (key, buckets) in values { + let key = key + .filter(|key| key.labels.len() == grouping_label_count) + .ok_or_else(|| { + QueryExecutionError::Native(format!( + "store key doesn't have the aggregation's {grouping_label_count} \ + grouping label values" + )) + })?; + let projected = KeyByLabelValues::new_with_labels( + projection.iter().map(|&i| key.labels[i].clone()).collect(), + ); + rolled.entry(Some(projected)).or_default().extend(buckets); + } + // The store sorts each group's buckets chronologically; keep that + // order across the merged groups. + for buckets in rolled.values_mut() { + buckets.sort_by_key(|(range, _)| *range); + } + Ok(rolled) + } + fn read_range_query_inputs( &self, context: &RangeQueryExecutionContext, ) -> Result { + let rollup_projection = Self::rollup_projection(&context.base)?; let lookback_ms = (context.lookback_bucket_count as u64) * context.tumbling_window_ms; let values = if context.window_type == WindowType::Sliding { self.execute_sliding_cover_query( @@ -2630,6 +2714,14 @@ impl SimpleEngine { values.len(), values.values().map(|buckets| buckets.len()).sum::() ); + let values = match rollup_projection { + Some(projection) => Self::roll_up_groups( + values, + &projection, + context.base.grouping_labels.labels.len(), + )?, + None => values, + }; let keys = match &context.base.store_plan.keys_query { Some(query) if context.keys_window_type == Some(WindowType::Sliding) => { @@ -3147,6 +3239,85 @@ impl SimpleEngine { } } +#[cfg(test)] +mod roll_up_groups_tests { + use super::SimpleEngine; + use crate::data_model::{AggregateCore, KeyByLabelValues}; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::stores::TimestampedBucketsMap; + use std::sync::Arc; + + fn bucket(start: u64, value: f64) -> ((u64, u64), Arc) { + ( + (start, start + 1_000), + Arc::new(SumAccumulator::with_sum(value)), + ) + } + + fn key(labels: &[&str]) -> Option { + Some(KeyByLabelValues::new_with_labels( + labels.iter().map(|label| label.to_string()).collect(), + )) + } + + #[test] + fn groups_with_the_same_projection_share_one_chronological_bucket_list() { + let values: TimestampedBucketsMap = [ + (key(&["/x", "a"]), vec![bucket(2_000, 1.0)]), + ( + key(&["/y", "a"]), + vec![bucket(1_000, 2.0), bucket(2_000, 4.0)], + ), + (key(&["/x", "b"]), vec![bucket(1_000, 8.0)]), + ] + .into_iter() + .collect(); + + let rolled = SimpleEngine::roll_up_groups(values, &[1], 2).unwrap(); + + let ranges = |labels: &[&str]| -> Vec<(u64, u64)> { + rolled[&key(labels)] + .iter() + .map(|(range, _)| *range) + .collect() + }; + assert_eq!(rolled.len(), 2); + assert_eq!( + ranges(&["a"]), + vec![(1_000, 2_000), (2_000, 3_000), (2_000, 3_000)] + ); + assert_eq!(ranges(&["b"]), vec![(1_000, 2_000)]); + } + + #[test] + fn empty_projection_rolls_everything_into_one_group() { + let values: TimestampedBucketsMap = [ + (key(&["/x", "a"]), vec![bucket(1_000, 1.0)]), + (key(&["/x", "b"]), vec![bucket(1_000, 2.0)]), + ] + .into_iter() + .collect(); + + let rolled = SimpleEngine::roll_up_groups(values, &[], 2).unwrap(); + + assert_eq!(rolled.len(), 1); + assert_eq!(rolled[&key(&[])].len(), 2); + } + + #[test] + fn a_store_key_without_every_grouping_value_is_an_error() { + for bad_key in [None, key(&["a"])] { + let values: TimestampedBucketsMap = [(bad_key.clone(), vec![bucket(1_000, 1.0)])] + .into_iter() + .collect(); + assert!( + SimpleEngine::roll_up_groups(values, &[0], 2).is_err(), + "{bad_key:?} should be rejected" + ); + } + } +} + #[cfg(test)] mod topk_metadata_tests { use super::{QueryExecutionError, SimpleEngine}; diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index a7f3a5f5..47cab84a 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -12,6 +12,7 @@ pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; pub mod query_time_aggregation_public_tests; pub mod range_query_arithmetic_tests; +pub mod rollup_tests; pub mod sql_pattern_matching_tests; pub mod stage_e_instant_range_equivalence_tests; pub mod store_correctness_tests; diff --git a/asap-query-engine/src/tests/rollup_tests.rs b/asap-query-engine/src/tests/rollup_tests.rs new file mode 100644 index 00000000..62eb0cd1 --- /dev/null +++ b/asap-query-engine/src/tests/rollup_tests.rs @@ -0,0 +1,488 @@ +//! Roll-ups: a query grouped by `Q` served from a single-subpopulation +//! aggregation grouped by `G ⊋ Q`, by merging the states of every `G`-group +//! that projects onto the same `Q`-group. + +#[cfg(test)] +mod tests { + use crate::data_model::{ + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + KeyByLabelValues, Measurement, PrecomputedOutput, QueryConfig, QueryLanguage, SchemaConfig, + StreamingConfig, WindowType, + }; + use crate::engines::query_result::QueryResult; + use crate::engines::simple_engine::SimpleEngine; + use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; + use crate::precompute_operators::ddsketch_accumulator::DDSketchAccumulator; + use crate::precompute_operators::hll_accumulator::HllAccumulator; + use crate::precompute_operators::increase_accumulator::IncreaseAccumulator; + use crate::precompute_operators::min_max_accumulator::MinMaxAccumulator; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::stores::simple_map_store::SimpleMapStore; + use crate::stores::Store; + use crate::tests::test_utilities::config_builders::create_test_schema; + use crate::tests::test_utilities::engine_factories::{ + create_engine_multi_timestamp_with_window, create_engine_single_pop, AccumulatorData, + }; + use crate::AggregateCore; + use promql_utilities::data_model::KeyByLabelNames; + use std::collections::HashMap; + use std::sync::Arc; + + const QUERY_TIME_S: f64 = 1_000.0; + + /// Store key values, ordered like the aggregation's grouping label names + /// (which are kept sorted). + fn key(labels: &[&str]) -> Option> { + Some(labels.iter().map(|label| label.to_string()).collect()) + } + + /// Instant rows as (labels, value), sorted by labels. + fn instant_rows(engine: &SimpleEngine, query: &str) -> Vec<(Vec, f64)> { + let (_, result) = engine + .handle_query_promql(query.to_string(), QUERY_TIME_S) + .expect("query should execute") + .expect("query should be served locally"); + let QueryResult::Vector(vector) = result else { + panic!("expected an instant vector"); + }; + let mut rows: Vec<_> = vector + .values + .into_iter() + .map(|element| (element.labels.labels, element.value)) + .collect(); + rows.sort_by(|left, right| left.0.cmp(&right.0)); + rows + } + + fn sums(values: &[(&[&str], f64)]) -> AccumulatorData { + values + .iter() + .map(|(labels, value)| { + ( + key(labels), + Box::new(SumAccumulator::with_sum(*value)) as Box, + ) + }) + .collect() + } + + fn kll(values: impl IntoIterator) -> DatasketchesKLLAccumulator { + let mut sketch = DatasketchesKLLAccumulator::new(200); + values.into_iter().for_each(|value| sketch.update(value)); + sketch + } + + fn dd(values: impl IntoIterator) -> DDSketchAccumulator { + let mut sketch = DDSketchAccumulator::new(0.01); + values.into_iter().for_each(|value| sketch.update(value)); + sketch + } + + #[test] + fn sum_by_subset_equals_sum_deployed_at_that_grouping() { + let query = "sum by (service) (requests)"; + let fine = create_engine_single_pop( + "requests", + AggregationType::Sum, + vec!["endpoint", "service"], + sums(&[ + (&["/x", "a"], 1.0), + (&["/y", "a"], 2.0), + (&["/x", "b"], 4.0), + ]), + query, + ); + let coarse = create_engine_single_pop( + "requests", + AggregationType::Sum, + vec!["service"], + sums(&[(&["a"], 3.0), (&["b"], 4.0)]), + query, + ); + + let rows = instant_rows(&fine, query); + assert_eq!( + rows, + vec![(vec!["a".to_string()], 3.0), (vec!["b".to_string()], 4.0)] + ); + assert_eq!(rows, instant_rows(&coarse, query)); + } + + #[test] + fn sum_without_grouping_rolls_every_group_into_one() { + let query = "sum(requests)"; + let engine = create_engine_single_pop( + "requests", + AggregationType::Sum, + vec!["endpoint", "service"], + sums(&[ + (&["/x", "a"], 1.0), + (&["/y", "a"], 2.0), + (&["/x", "b"], 4.0), + ]), + query, + ); + + assert_eq!(instant_rows(&engine, query), vec![(vec![], 7.0)]); + } + + /// Dropping a label from the middle of the grouping keeps the remaining + /// values aligned with the output label names. + #[test] + fn roll_up_drops_a_middle_grouping_label() { + let query = "sum by (az, service) (requests)"; + let engine = create_engine_single_pop( + "requests", + AggregationType::Sum, + vec!["az", "endpoint", "service"], + sums(&[ + (&["eu", "/x", "a"], 1.0), + (&["eu", "/y", "a"], 2.0), + (&["us", "/x", "a"], 4.0), + (&["eu", "/x", "b"], 8.0), + ]), + query, + ); + + let (labels, _) = engine + .handle_query_promql(query.to_string(), QUERY_TIME_S) + .expect("query should execute") + .expect("query should be served locally"); + assert_eq!(labels.labels, vec!["az".to_string(), "service".to_string()]); + let row = + |az: &str, service: &str, value| (vec![az.to_string(), service.to_string()], value); + assert_eq!( + instant_rows(&engine, query), + vec![ + row("eu", "a", 3.0), + row("eu", "b", 8.0), + row("us", "a", 4.0) + ] + ); + } + + #[test] + fn max_by_subset_is_the_max_over_the_subgroups() { + let query = "max by (service) (latency)"; + let data: AccumulatorData = [ + (&["/x", "a"][..], 5.0), + (&["/y", "a"][..], 9.0), + (&["/x", "b"][..], 2.0), + ] + .into_iter() + .map(|(labels, value)| { + ( + key(labels), + Box::new(MinMaxAccumulator::with_value(value, "max".to_string()).unwrap()) + as Box, + ) + }) + .collect(); + let engine = create_engine_single_pop( + "latency", + AggregationType::MinMax, + vec!["endpoint", "service"], + data, + query, + ); + + assert_eq!( + instant_rows(&engine, query), + vec![(vec!["a".to_string()], 9.0), (vec!["b".to_string()], 2.0)] + ); + } + + /// A merged sketch answers for the union of its subgroups. + #[test] + fn kll_quantile_by_subset_is_within_rank_error_of_the_union() { + let query = "quantile by (service) (0.5, latency)"; + let a_x: Vec = (0..500).map(f64::from).collect(); + let a_y: Vec = (500..1_000).map(f64::from).collect(); + let b_x: Vec = (0..100).map(|v| f64::from(v) * 10.0).collect(); + let data: AccumulatorData = vec![ + (key(&["/x", "a"]), Box::new(kll(a_x.clone()))), + (key(&["/y", "a"]), Box::new(kll(a_y.clone()))), + (key(&["/x", "b"]), Box::new(kll(b_x.clone()))), + ]; + let engine = create_engine_single_pop( + "latency", + AggregationType::DatasketchesKLL, + vec!["endpoint", "service"], + data, + query, + ); + + // KLL compaction is randomized, so compare against the true median + // within its rank error rather than against a second merge. + let rows = instant_rows(&engine, query); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].0, vec!["a".to_string()]); + assert!( + (rows[0].1 - 500.0).abs() <= 1_000.0 * 0.02, + "median of 0..1000 should be ~500, got {}", + rows[0].1 + ); + // Below k items a KLL is exact. + assert_eq!(rows[1], (vec!["b".to_string()], kll(b_x).get_quantile(0.5))); + } + + #[test] + fn ddsketch_quantile_by_subset_equals_the_merged_sketch() { + let query = "quantile by (service) (0.9, latency)"; + let a_x: Vec = (1..=500).map(f64::from).collect(); + let a_y: Vec = (501..=1_000).map(f64::from).collect(); + let data: AccumulatorData = vec![ + (key(&["/x", "a"]), Box::new(dd(a_x.clone()))), + (key(&["/y", "a"]), Box::new(dd(a_y.clone()))), + ]; + let engine = create_engine_single_pop( + "latency", + AggregationType::DDSketch, + vec!["endpoint", "service"], + data, + query, + ); + + let rows = instant_rows(&engine, query); + let expected = dd(a_x.into_iter().chain(a_y)).get_quantile(0.9); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].0, vec!["a".to_string()]); + assert!( + (rows[0].1 - expected).abs() <= expected * 1e-9, + "merged DDSketch p90 {} should equal the whole-data sketch's {expected}", + rows[0].1 + ); + assert!((rows[0].1 - 900.0).abs() <= 900.0 * 0.02); + } + + /// HLL merges by register max, so the rolled-up sketch is exactly the + /// sketch of the union. HLL is reached through SQL `COUNT(DISTINCT)`. + #[test] + fn hll_count_distinct_by_subset_equals_the_union_sketch() { + let sql = "SELECT COUNT(DISTINCT value) FROM events WHERE time BETWEEN \ + DATEADD(s, -1, NOW()) AND NOW() GROUP BY service"; + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs: HashMap::from([( + 1, + AggregationConfig { + aggregation_id: 1, + aggregation_type: AggregationType::HLL, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::new(vec![ + "endpoint".to_string(), + "service".to_string(), + ]), + aggregated_labels: KeyByLabelNames::empty(), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: 1_000, + slide_interval_ms: 1_000, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: "events".to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: Some("events".to_string()), + value_column: Some("value".to_string()), + }, + )]), + }); + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + let hll = |values: std::ops::Range| { + let mut sketch = HllAccumulator::with_default_precision(); + values.for_each(|value| sketch.update(f64::from(value))); + sketch + }; + let query_time_ms = (QUERY_TIME_S * 1_000.0) as u64; + for (labels, values) in [(["/x", "a"], 0..1_000), (["/y", "a"], 500..2_000)] { + store + .insert_precomputed_output( + PrecomputedOutput::new( + query_time_ms - 1_000, + query_time_ms, + key(&labels).map(|labels| KeyByLabelValues { labels }), + 1, + ), + Box::new(hll(values)), + ) + .unwrap(); + } + let inference_config = InferenceConfig { + schema: SchemaConfig::SQL(create_test_schema( + "events", + "time", + "value", + vec!["endpoint", "service"], + )), + query_configs: vec![QueryConfig::new(sql.to_string()) + .add_aggregation(AggregationReference::new(1, None))], + cleanup_policy: CleanupPolicy::NoCleanup, + }; + let engine = SimpleEngine::new( + store, + inference_config, + streaming_config, + 1_000, + QueryLanguage::sql, + ); + + let (_, result) = engine + .handle_query_sql(sql.to_string(), QUERY_TIME_S) + .expect("the SQL roll-up should be served locally"); + let QueryResult::Vector(vector) = result else { + panic!("expected an instant vector"); + }; + assert_eq!(vector.values.len(), 1); + assert_eq!(vector.values[0].labels.labels, vec!["a".to_string()]); + assert_eq!(vector.values[0].value, hll(0..2_000).estimate()); + } + + /// A range query rolls up at every step, and a subgroup that only + /// appears partway through contributes from then on. + #[test] + fn range_sum_by_subset_rolls_up_each_step() { + let query = "sum by (service) (requests)"; + let mut data = Vec::new(); + for second in 1..=4u64 { + data.push(( + second * 1_000, + key(&["/x", "a"]), + Box::new(SumAccumulator::with_sum(1.0)) as Box, + )); + if second >= 3 { + data.push(( + second * 1_000, + key(&["/new", "a"]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )); + } + } + let engine = create_engine_multi_timestamp_with_window( + "requests", + AggregationType::Sum, + vec!["endpoint", "service"], + data, + query, + 1_000, + 1_000, + WindowType::Tumbling, + ); + + let (_, result) = engine + .handle_range_query_promql(query.to_string(), 2.0, 4.0, 1.0) + .expect("query should execute") + .expect("query should be served locally"); + let QueryResult::Matrix(matrix) = result else { + panic!("expected a range vector"); + }; + assert_eq!(matrix.values.len(), 1); + assert_eq!(matrix.values[0].labels.labels, vec!["a".to_string()]); + let samples: Vec<(u64, f64)> = matrix.values[0] + .samples + .iter() + .map(|sample| (sample.timestamp, sample.value)) + .collect(); + assert_eq!(samples, vec![(2_000, 1.0), (3_000, 11.0), (4_000, 11.0)]); + } + + /// A Sliding step is complete when every cover window holds data from + /// some subgroup. A subgroup with no samples in one cover window still + /// contributes its other windows. + #[test] + fn sliding_range_roll_up_keeps_a_subgroup_missing_an_early_cover_window() { + let query = "sum(sum_over_time(requests[4s])) by (service)"; + let mut data = Vec::new(); + for second in 1..=8u64 { + data.push(( + second * 1_000, + key(&["/x", "a"]), + Box::new(SumAccumulator::with_sum(1.0)) as Box, + )); + if second >= 5 { + data.push(( + second * 1_000, + key(&["/new", "a"]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )); + } + } + let engine = create_engine_multi_timestamp_with_window( + "requests", + AggregationType::Sum, + vec!["endpoint", "service"], + data, + query, + 2_000, + 1_000, + WindowType::Sliding, + ); + + let (_, result) = engine + .handle_range_query_promql(query.to_string(), 6.0, 8.0, 2.0) + .expect("query should execute") + .expect("query should be served locally"); + let QueryResult::Matrix(matrix) = result else { + panic!("expected a range vector"); + }; + let samples: Vec<(u64, f64)> = matrix.values[0] + .samples + .iter() + .map(|sample| (sample.timestamp, sample.value)) + .collect(); + // At 6s, /new has no samples in [2s, 4s) but 20 in [4s, 6s). + assert_eq!(samples, vec![(6_000, 24.0), (8_000, 44.0)]); + } + + #[test] + fn grouping_outside_the_aggregations_labels_is_an_error() { + let query = "sum by (endpoint) (requests)"; + let engine = create_engine_single_pop( + "requests", + AggregationType::Sum, + vec!["service"], + sums(&[(&["a"], 1.0)]), + query, + ); + + let error = engine + .handle_query_promql(query.to_string(), QUERY_TIME_S) + .expect_err("a grouping the aggregation lacks must not be answered"); + assert!(error.to_string().contains("can't serve"), "{error}"); + } + + /// Increase merges pieces of one counter in time, so summing two + /// series' counters by merging them would be wrong. + #[test] + fn roll_up_of_a_type_not_mergeable_across_groups_is_an_error() { + let query = "sum by (service) (requests)"; + let data: AccumulatorData = vec![( + key(&["/x", "a"]), + Box::new(IncreaseAccumulator::new( + Measurement::new(1.0), + 0, + Measurement::new(5.0), + 1_000, + )) as Box, + )]; + let engine = create_engine_single_pop( + "requests", + AggregationType::Increase, + vec!["endpoint", "service"], + data, + query, + ); + + let error = engine + .handle_query_promql(query.to_string(), QUERY_TIME_S) + .expect_err("an Increase roll-up must not be answered"); + assert!( + error.to_string().contains("can't be merged across groups"), + "{error}" + ); + } +}