Add correctness tests: new store vs legacy per all Store trait methods

Tests in store_correctness_tests.rs verify that SimpleMapStorePerKey and
SimpleMapStoreGlobal produce identical results to LegacySimpleMapStorePerKey
for range queries, exact queries, earliest-timestamp lookups, and batch insert.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: zzylol

asap-query-engine/src/tests/mod.rs

@@ -3,6 +3,7 @@ pub mod datafusion;
 pub mod elastic_forwarding_tests;
 pub mod prometheus_forwarding_tests;
 pub mod query_equivalence_tests;
+pub mod store_correctness_tests;
 pub mod trait_design_tests;
 
 #[cfg(test)]

asap-query-engine/src/tests/store_correctness_tests.rs

@@ -0,0 +1,379 @@
+//! Store Correctness Tests
+//!
+//! Verifies that `SimpleMapStorePerKey` and `SimpleMapStoreGlobal` produce
+//! identical results to `LegacySimpleMapStorePerKey` for range queries,
+//! exact queries, and earliest-timestamp lookups.
+
+#![allow(deprecated)]
+
+use crate::data_model::{CleanupPolicy, KeyByLabelValues, LockStrategy, PrecomputedOutput};
+use crate::precompute_operators::sum_accumulator::SumAccumulator;
+use crate::stores::simple_map_store::per_key_legacy::LegacySimpleMapStorePerKey;
+use crate::stores::simple_map_store::SimpleMapStore;
+use crate::stores::{Store, TimestampedBucketsMap};
+use crate::AggregateCore;
+use promql_utilities::data_model::KeyByLabelNames;
+use sketch_db_common::aggregation_config::AggregationConfig;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+fn make_streaming_config(
+    num_aggregates_to_retain: Option<u64>,
+) -> Arc<crate::data_model::StreamingConfig> {
+    let mut configs = HashMap::new();
+    configs.insert(
+        1u64,
+        AggregationConfig {
+            aggregation_id: 1,
+            aggregation_type: "SumAccumulator".to_string(),
+            aggregation_sub_type: String::new(),
+            parameters: HashMap::new(),
+            grouping_labels: KeyByLabelNames::empty(),
+            aggregated_labels: KeyByLabelNames::empty(),
+            rollup_labels: KeyByLabelNames::empty(),
+            original_yaml: String::new(),
+            window_size: 1000,
+            slide_interval: 1000,
+            window_type: "tumbling".to_string(),
+            tumbling_window_size: 1000,
+            spatial_filter: String::new(),
+            spatial_filter_normalized: String::new(),
+            metric: "test_metric".to_string(),
+            num_aggregates_to_retain,
+            read_count_threshold: None,
+            table_name: None,
+            value_column: None,
+        },
+    );
+    Arc::new(crate::data_model::StreamingConfig::new(configs))
+}
+
+/// One row of test data: (start_ts, end_ts, label, sum_value)
+struct Row {
+    start: u64,
+    end: u64,
+    label: String,
+    value: f64,
+}
+
+fn populate_new(store: &SimpleMapStore, rows: &[Row]) {
+    for r in rows {
+        let key = KeyByLabelValues::new_with_labels(vec![r.label.clone()]);
+        let output = PrecomputedOutput::new(r.start, r.end, Some(key), 1);
+        let acc: Box<dyn AggregateCore> = Box::new(SumAccumulator::with_sum(r.value));
+        store.insert_precomputed_output(output, acc).unwrap();
+    }
+}
+
+fn populate_legacy(store: &LegacySimpleMapStorePerKey, rows: &[Row]) {
+    for r in rows {
+        let key = KeyByLabelValues::new_with_labels(vec![r.label.clone()]);
+        let output = PrecomputedOutput::new(r.start, r.end, Some(key), 1);
+        let acc: Box<dyn AggregateCore> = Box::new(SumAccumulator::with_sum(r.value));
+        store.insert_precomputed_output(output, acc).unwrap();
+    }
+}
+
+/// Extract (label, sum) pairs from a TimestampedBucketsMap, sorted for comparison.
+fn extract_sums(map: &TimestampedBucketsMap) -> Vec<(Vec<String>, u64, u64, f64)> {
+    let mut out = Vec::new();
+    for (key, buckets) in map {
+        let labels = key
+            .as_ref()
+            .map(|k| k.labels.clone())
+            .unwrap_or_default();
+        for ((start, end), acc) in buckets {
+            let sum = acc
+                .as_any()
+                .downcast_ref::<SumAccumulator>()
+                .map(|s| s.sum)
+                .unwrap_or(f64::NAN);
+            out.push((labels.clone(), *start, *end, sum));
+        }
+    }
+    out.sort_by(|a, b| {
+        a.0.cmp(&b.0)
+            .then(a.1.cmp(&b.1))
+            .then(a.2.cmp(&b.2))
+    });
+    out
+}
+
+/// Assert that two `TimestampedBucketsMap` results are identical.
+fn assert_maps_equal(new: &TimestampedBucketsMap, legacy: &TimestampedBucketsMap, ctx: &str) {
+    let new_sums = extract_sums(new);
+    let legacy_sums = extract_sums(legacy);
+    assert_eq!(
+        new_sums, legacy_sums,
+        "{ctx}: result mismatch\n  new:    {new_sums:?}\n  legacy: {legacy_sums:?}",
+    );
+}
+
+/// Build a standard set of rows: 5 time windows × 3 labels.
+fn make_rows() -> Vec<Row> {
+    let mut rows = Vec::new();
+    for i in 0u64..5 {
+        let start = i * 1000;
+        let end = start + 1000;
+        for (j, label) in ["host-a", "host-b", "host-c"].iter().enumerate() {
+            rows.push(Row {
+                start,
+                end,
+                label: label.to_string(),
+                value: (i * 3 + j as u64) as f64 + 1.0,
+            });
+        }
+    }
+    rows
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // --- range query (query_precomputed_output) ---
+
+    #[test]
+    fn per_key_range_query_matches_legacy_full_range() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let legacy = LegacySimpleMapStorePerKey::new(config.clone(), CleanupPolicy::NoCleanup);
+        populate_legacy(&legacy, &rows);
+
+        let new = SimpleMapStore::new_with_strategy(
+            config,
+            CleanupPolicy::NoCleanup,
+            LockStrategy::PerKey,
+        );
+        populate_new(&new, &rows);
+
+        let start = 0;
+        let end = 5000;
+        let new_result = new
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+        let legacy_result = legacy
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+
+        assert_maps_equal(&new_result, &legacy_result, "per_key range full");
+    }
+
+    #[test]
+    fn global_range_query_matches_legacy_full_range() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let legacy = LegacySimpleMapStorePerKey::new(config.clone(), CleanupPolicy::NoCleanup);
+        populate_legacy(&legacy, &rows);
+
+        let new = SimpleMapStore::new_with_strategy(
+            config,
+            CleanupPolicy::NoCleanup,
+            LockStrategy::Global,
+        );
+        populate_new(&new, &rows);
+
+        let start = 0;
+        let end = 5000;
+        let new_result = new
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+        let legacy_result = legacy
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+
+        assert_maps_equal(&new_result, &legacy_result, "global range full");
+    }
+
+    #[test]
+    fn per_key_range_query_matches_legacy_partial_range() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let legacy = LegacySimpleMapStorePerKey::new(config.clone(), CleanupPolicy::NoCleanup);
+        populate_legacy(&legacy, &rows);
+
+        let new = SimpleMapStore::new_with_strategy(
+            config,
+            CleanupPolicy::NoCleanup,
+            LockStrategy::PerKey,
+        );
+        populate_new(&new, &rows);
+
+        // Query only the middle three windows
+        let start = 1000;
+        let end = 3999;
+        let new_result = new
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+        let legacy_result = legacy
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+
+        assert_maps_equal(&new_result, &legacy_result, "per_key range partial");
+    }
+
+    // --- exact query (query_precomputed_output_exact) ---
+
+    #[test]
+    fn per_key_exact_query_matches_legacy() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let legacy = LegacySimpleMapStorePerKey::new(config.clone(), CleanupPolicy::NoCleanup);
+        populate_legacy(&legacy, &rows);
+
+        let new = SimpleMapStore::new_with_strategy(
+            config,
+            CleanupPolicy::NoCleanup,
+            LockStrategy::PerKey,
+        );
+        populate_new(&new, &rows);
+
+        // Exact match on the third window [2000, 3000)
+        let exact_start = 2000;
+        let exact_end = 3000;
+        let new_result = new
+            .query_precomputed_output_exact("test_metric", 1, exact_start, exact_end)
+            .unwrap();
+        let legacy_result = legacy
+            .query_precomputed_output_exact("test_metric", 1, exact_start, exact_end)
+            .unwrap();
+
+        assert_maps_equal(&new_result, &legacy_result, "per_key exact");
+    }
+
+    #[test]
+    fn global_exact_query_matches_legacy() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let legacy = LegacySimpleMapStorePerKey::new(config.clone(), CleanupPolicy::NoCleanup);
+        populate_legacy(&legacy, &rows);
+
+        let new = SimpleMapStore::new_with_strategy(
+            config,
+            CleanupPolicy::NoCleanup,
+            LockStrategy::Global,
+        );
+        populate_new(&new, &rows);
+
+        let exact_start = 2000;
+        let exact_end = 3000;
+        let new_result = new
+            .query_precomputed_output_exact("test_metric", 1, exact_start, exact_end)
+            .unwrap();
+        let legacy_result = legacy
+            .query_precomputed_output_exact("test_metric", 1, exact_start, exact_end)
+            .unwrap();
+
+        assert_maps_equal(&new_result, &legacy_result, "global exact");
+    }
+
+    #[test]
+    fn exact_query_returns_empty_when_no_match() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let new =
+            SimpleMapStore::new_with_strategy(config, CleanupPolicy::NoCleanup, LockStrategy::PerKey);
+        populate_new(&new, &rows);
+
+        // No window at [9000, 10000)
+        let result = new
+            .query_precomputed_output_exact("test_metric", 1, 9000, 10000)
+            .unwrap();
+
+        assert!(result.is_empty(), "Expected empty result for unmatched exact query");
+    }
+
+    // --- earliest timestamp ---
+
+    #[test]
+    fn per_key_earliest_timestamp_matches_legacy() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let legacy = LegacySimpleMapStorePerKey::new(config.clone(), CleanupPolicy::NoCleanup);
+        populate_legacy(&legacy, &rows);
+
+        let new = SimpleMapStore::new_with_strategy(
+            config,
+            CleanupPolicy::NoCleanup,
+            LockStrategy::PerKey,
+        );
+        populate_new(&new, &rows);
+
+        let new_ts = new.get_earliest_timestamp_per_aggregation_id().unwrap();
+        let legacy_ts = legacy.get_earliest_timestamp_per_aggregation_id().unwrap();
+
+        assert_eq!(new_ts, legacy_ts, "earliest timestamp mismatch (per_key)");
+    }
+
+    #[test]
+    fn global_earliest_timestamp_matches_legacy() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let legacy = LegacySimpleMapStorePerKey::new(config.clone(), CleanupPolicy::NoCleanup);
+        populate_legacy(&legacy, &rows);
+
+        let new = SimpleMapStore::new_with_strategy(
+            config,
+            CleanupPolicy::NoCleanup,
+            LockStrategy::Global,
+        );
+        populate_new(&new, &rows);
+
+        let new_ts = new.get_earliest_timestamp_per_aggregation_id().unwrap();
+        let legacy_ts = legacy.get_earliest_timestamp_per_aggregation_id().unwrap();
+
+        assert_eq!(new_ts, legacy_ts, "earliest timestamp mismatch (global)");
+    }
+
+    // --- batch insert produces same results as single insert ---
+
+    #[test]
+    fn batch_insert_matches_single_insert_per_key() {
+        let rows = make_rows();
+        let config = make_streaming_config(None);
+
+        let single =
+            SimpleMapStore::new_with_strategy(config.clone(), CleanupPolicy::NoCleanup, LockStrategy::PerKey);
+        populate_new(&single, &rows);
+
+        let batch =
+            SimpleMapStore::new_with_strategy(config, CleanupPolicy::NoCleanup, LockStrategy::PerKey);
+        let batch_inputs: Vec<_> = rows
+            .iter()
+            .map(|r| {
+                let key = KeyByLabelValues::new_with_labels(vec![r.label.clone()]);
+                let output = PrecomputedOutput::new(r.start, r.end, Some(key), 1);
+                let acc: Box<dyn AggregateCore> = Box::new(SumAccumulator::with_sum(r.value));
+                (output, acc)
+            })
+            .collect();
+        batch.insert_precomputed_output_batch(batch_inputs).unwrap();
+
+        let start = 0;
+        let end = 5000;
+        let single_result = single
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+        let batch_result = batch
+            .query_precomputed_output("test_metric", 1, start, end)
+            .unwrap();
+
+        assert_maps_equal(&single_result, &batch_result, "batch vs single insert");
+    }
+}