feat(query-engine): roll up single-subpopulation aggregations to a coarser grouping

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

@@ -188,6 +188,17 @@ struct RangeQueryReads {
     keys: Option<TimestampedBucketsMap>,
 }
 
+/// 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<u64, Vec<Arc<dyn AggregateCore>>>;
 type TopkPartition = (u64, Vec<String>);
 type TopkCandidates<'a> = HashMap<TopkPartition, Vec<(&'a KeyByLabelValues, f64)>>;
@@ -2616,10 +2627,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<Option<Vec<usize>>, 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::<Option<Vec<usize>>>()
+            .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<TimestampedBucketsMap, QueryExecutionError> {
+        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<RangeQueryReads, QueryExecutionError> {
+        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(
@@ -2644,6 +2728,14 @@ impl SimpleEngine {
             values.len(),
             values.values().map(|buckets| buckets.len()).sum::<usize>()
         );
+        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) => {
@@ -3174,6 +3266,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<dyn AggregateCore>) {
+        (
+            (start, start + 1_000),
+            Arc::new(SumAccumulator::with_sum(value)),
+        )
+    }
+
+    fn key(labels: &[&str]) -> Option<KeyByLabelValues> {
+        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};

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;