refactor: restructure precompute engine from per-series to per-group accumulators

Match Arroyo's GROUP BY semantics: the ingest handler now extracts
grouping label values from each series and groups samples by
(agg_id, group_key) before routing to workers. The router hashes
by group key so all series sharing the same grouping labels land on
the same worker and feed a single shared accumulator.

For the quickstart (189K series, 7 pattern groups), this reduces
accumulator count from 189K to 7, store writes/sec from 189K to 7,
and eliminates query-time fan-in merge.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: zz_y

asap-query-engine/src/main.rs

@@ -11,6 +11,7 @@ use query_engine_rust::data_model::enums::{InputFormat, LockStrategy, StreamingE
 use query_engine_rust::drivers::AdapterConfig;
 use query_engine_rust::precompute_engine::config::LateDataPolicy;
 use query_engine_rust::utils::file_io::{read_inference_config, read_streaming_config};
+use query_engine_rust::precompute_engine::PrecomputeWorkerDiagnostics;
 use query_engine_rust::{
     HttpServer, HttpServerConfig, KafkaConsumer, KafkaConsumerConfig, OtlpReceiver,
     OtlpReceiverConfig, PrecomputeEngine, PrecomputeEngineConfig, Result, SimpleEngine,
@@ -323,16 +324,29 @@ async fn main() -> Result<()> {
         let output_sink = Arc::new(StoreOutputSink::new(store.clone()));
         let engine =
             PrecomputeEngine::new(precompute_config, streaming_config.clone(), output_sink);
+        let worker_diagnostics = engine.diagnostics();
         info!(
             "Starting precompute engine on port {}",
             args.prometheus_remote_write_port
         );
+
+        // Spawn periodic memory diagnostics logger
+        let diag_store = store.clone();
+        tokio::spawn(async move {
+            spawn_memory_diagnostics(diag_store, Some(worker_diagnostics)).await;
+        });
+
         Some(tokio::spawn(async move {
             if let Err(e) = engine.run().await {
                 error!("Precompute engine error: {}", e);
             }
         }))
     } else {
+        // Even without precompute, log store diagnostics
+        let diag_store = store.clone();
+        tokio::spawn(async move {
+            spawn_memory_diagnostics(diag_store, None).await;
+        });
         None
     };
 
@@ -396,6 +410,59 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
+/// Periodic memory diagnostics logger — runs every 30 seconds.
+async fn spawn_memory_diagnostics(
+    store: Arc<SimpleMapStore>,
+    worker_diagnostics: Option<Arc<PrecomputeWorkerDiagnostics>>,
+) {
+    use std::sync::atomic::Ordering;
+
+    let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(30));
+    loop {
+        interval.tick().await;
+
+        // 1. Store diagnostics
+        let store_diag = store.diagnostic_info();
+        info!(
+            "[MEMORY_DIAG] Store: {} aggregation(s), {} total time_map entries, {:.2} KB total sketch bytes",
+            store_diag.num_aggregations,
+            store_diag.total_time_map_entries,
+            store_diag.total_sketch_bytes as f64 / 1024.0,
+        );
+        for agg in &store_diag.per_aggregation {
+            info!(
+                "[MEMORY_DIAG]   agg_id={}: time_map_len={}, read_counts_len={}, aggregate_objects={}, sketch_bytes={:.2} KB",
+                agg.aggregation_id,
+                agg.time_map_len,
+                agg.read_counts_len,
+                agg.num_aggregate_objects,
+                agg.sketch_bytes as f64 / 1024.0,
+            );
+        }
+
+        // 2. Worker diagnostics (precompute engine only)
+        if let Some(ref diag) = worker_diagnostics {
+            let total_groups: usize = diag
+                .worker_group_counts
+                .iter()
+                .map(|c| c.load(Ordering::Relaxed))
+                .sum();
+            info!(
+                "[MEMORY_DIAG] PrecomputeEngine: {} total groups across {} workers",
+                total_groups,
+                diag.worker_group_counts.len(),
+            );
+            for (i, counter) in diag.worker_group_counts.iter().enumerate() {
+                info!(
+                    "[MEMORY_DIAG]   worker_{}: group_states_len={}",
+                    i,
+                    counter.load(Ordering::Relaxed),
+                );
+            }
+        }
+    }
+}
+
 fn setup_logging(
     output_dir: &str,
     log_level: &str,

asap-query-engine/src/precompute_engine/engine.rs

@@ -9,18 +9,25 @@ use crate::precompute_engine::worker::{Worker, WorkerRuntimeConfig};
 use axum::{routing::post, Router};
 use sketch_db_common::aggregation_config::AggregationConfig;
 use std::collections::HashMap;
+use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 use tokio::net::TcpListener;
 use tokio::sync::mpsc;
 use tracing::{info, warn};
 
+/// Shared diagnostic counters readable from outside the engine.
+pub struct PrecomputeWorkerDiagnostics {
+    pub worker_group_counts: Vec<Arc<AtomicUsize>>,
+}
+
 /// The top-level precompute engine orchestrator.
 ///
 /// Creates worker threads, the series router, and the Axum ingest server.
 pub struct PrecomputeEngine {
     config: PrecomputeEngineConfig,
     streaming_config: Arc<StreamingConfig>,
     output_sink: Arc<dyn OutputSink>,
+    diagnostics: Arc<PrecomputeWorkerDiagnostics>,
 }
 
 impl PrecomputeEngine {
@@ -29,13 +36,25 @@ impl PrecomputeEngine {
         streaming_config: Arc<StreamingConfig>,
         output_sink: Arc<dyn OutputSink>,
     ) -> Self {
+        let worker_group_counts = (0..config.num_workers)
+            .map(|_| Arc::new(AtomicUsize::new(0)))
+            .collect();
+        let diagnostics = Arc::new(PrecomputeWorkerDiagnostics {
+            worker_group_counts,
+        });
         Self {
             config,
             streaming_config,
             output_sink,
+            diagnostics,
         }
     }
 
+    /// Get a handle to worker diagnostics, readable even after `run()` starts.
+    pub fn diagnostics(&self) -> Arc<PrecomputeWorkerDiagnostics> {
+        self.diagnostics.clone()
+    }
+
     /// Start the precompute engine. This spawns worker tasks and the HTTP
     /// ingest server, then blocks until shutdown.
     pub async fn run(self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
@@ -63,6 +82,10 @@ impl PrecomputeEngine {
             .map(|(&id, cfg)| (id, Arc::new(cfg.clone())))
             .collect();
 
+        // Build a Vec<Arc<AggregationConfig>> for the ingest handler
+        let agg_configs_vec: Vec<Arc<AggregationConfig>> =
+            agg_configs.values().cloned().collect();
+
         // Spawn workers
         let mut worker_handles = Vec::with_capacity(num_workers);
         for (id, rx) in receivers.into_iter().enumerate() {
@@ -78,6 +101,7 @@ impl PrecomputeEngine {
                     raw_mode_aggregation_id: self.config.raw_mode_aggregation_id,
                     late_data_policy: self.config.late_data_policy,
                 },
+                self.diagnostics.worker_group_counts[id].clone(),
             );
             let handle = tokio::spawn(async move {
                 worker.run().await;
@@ -94,6 +118,8 @@ impl PrecomputeEngine {
         let ingest_state = Arc::new(IngestState {
             router,
             samples_ingested: std::sync::atomic::AtomicU64::new(0),
+            agg_configs: agg_configs_vec,
+            pass_raw_samples: self.config.pass_raw_samples,
         });
 
         // Start flush timer

asap-query-engine/src/precompute_engine/ingest_handler.rs

@@ -1,7 +1,9 @@
 use crate::drivers::ingest::prometheus_remote_write::decode_prometheus_remote_write;
 use crate::drivers::ingest::victoriametrics_remote_write::decode_victoriametrics_remote_write;
-use crate::precompute_engine::series_router::SeriesRouter;
+use crate::precompute_engine::series_router::{SeriesRouter, WorkerMessage};
+use crate::precompute_engine::worker::{extract_metric_name, parse_labels_from_series_key};
 use axum::{body::Bytes, extract::State, http::StatusCode};
+use sketch_db_common::aggregation_config::AggregationConfig;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Instant;
@@ -11,9 +13,28 @@ use tracing::warn;
 pub(crate) struct IngestState {
     pub(crate) router: SeriesRouter,
     pub(crate) samples_ingested: std::sync::atomic::AtomicU64,
+    /// Aggregation configs for group-key extraction.
+    pub(crate) agg_configs: Vec<Arc<AggregationConfig>>,
+    /// When true, skip group-key extraction and pass raw samples through.
+    pub(crate) pass_raw_samples: bool,
 }
 
-/// Shared logic: group decoded samples by series key and route to workers.
+/// Extract the group key (grouping label values joined by semicolons)
+/// for a given series key and aggregation config.
+fn extract_group_key(series_key: &str, config: &AggregationConfig) -> String {
+    let labels = parse_labels_from_series_key(series_key);
+    let mut values = Vec::new();
+    for label_name in &config.grouping_labels.labels {
+        if let Some(val) = labels.get(label_name.as_str()) {
+            values.push(*val);
+        } else {
+            values.push("");
+        }
+    }
+    values.join(";")
+}
+
+/// Shared logic: group decoded samples by (agg_id, group_key) and route to workers.
 async fn route_decoded_samples(
     state: &IngestState,
     samples: Vec<crate::drivers::ingest::prometheus_remote_write::DecodedSample>,
@@ -28,25 +49,71 @@ async fn route_decoded_samples(
         .samples_ingested
         .fetch_add(count, std::sync::atomic::Ordering::Relaxed);
 
-    // Group samples by series key for batch routing
-    let mut by_series: HashMap<&str, Vec<(i64, f64)>> = HashMap::new();
+    if state.pass_raw_samples {
+        // Raw mode: group by series key and send as RawSamples
+        let mut by_series: HashMap<&str, Vec<(i64, f64)>> = HashMap::new();
+        for s in &samples {
+            by_series
+                .entry(&s.labels)
+                .or_default()
+                .push((s.timestamp_ms, s.value));
+        }
+        let messages: Vec<WorkerMessage> = by_series
+            .into_iter()
+            .map(|(k, v)| WorkerMessage::RawSamples {
+                series_key: k.to_string(),
+                samples: v,
+                ingest_received_at,
+            })
+            .collect();
+
+        if let Err(e) = state
+            .router
+            .route_group_batch(messages, ingest_received_at)
+            .await
+        {
+            warn!("Batch routing error: {}", e);
+            return StatusCode::INTERNAL_SERVER_ERROR;
+        }
+        return StatusCode::NO_CONTENT;
+    }
+
+    // Group-by mode: for each sample, find matching agg configs and group by
+    // (agg_id, group_key). This is the equivalent of Arroyo's GROUP BY.
+    //
+    // Key: (agg_id, group_key) → Vec<(series_key, timestamp_ms, value)>
+    let mut by_group: HashMap<(u64, String), Vec<(String, i64, f64)>> = HashMap::new();
+
     for s in &samples {
-        by_series
-            .entry(&s.labels)
-            .or_default()
-            .push((s.timestamp_ms, s.value));
+        let metric_name = extract_metric_name(&s.labels);
+        for config in &state.agg_configs {
+            if config.metric != metric_name
+                && config.spatial_filter_normalized != metric_name
+                && config.spatial_filter != metric_name
+            {
+                continue;
+            }
+            let group_key = extract_group_key(&s.labels, config);
+            by_group
+                .entry((config.aggregation_id, group_key))
+                .or_default()
+                .push((s.labels.clone(), s.timestamp_ms, s.value));
+        }
     }
 
-    // Convert to owned keys for batch routing
-    let by_series_owned: HashMap<String, Vec<(i64, f64)>> = by_series
+    let messages: Vec<WorkerMessage> = by_group
         .into_iter()
-        .map(|(k, v)| (k.to_string(), v))
+        .map(|((agg_id, group_key), samples)| WorkerMessage::GroupSamples {
+            agg_id,
+            group_key,
+            samples,
+            ingest_received_at,
+        })
         .collect();
 
-    // Route all series to workers concurrently
     if let Err(e) = state
         .router
-        .route_batch(by_series_owned, ingest_received_at)
+        .route_group_batch(messages, ingest_received_at)
         .await
     {
         warn!("Batch routing error: {}", e);