feat: all-five-sketch query_statistic + cold-store wiring (P1)

Cargo.toml

@@ -27,7 +27,7 @@ thiserror = "1.0"
 anyhow = "1.0"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-clap = { version = "4.0", features = ["derive"] }
+clap = { version = "4.0", features = ["derive", "env"] }
 chrono = { version = "0.4", features = ["serde"] }
 promql-parser = "0.5.0"
 tokio = { version = "1.0", features = ["full"] }

TODO.md

@@ -6,6 +6,49 @@ deferred to future work.
 
 See the design source at [`docs/design-sketch-db.md`](docs/design-sketch-db.md).
 
+## All-five-sketch query path verification (2026-04-30)
+
+Each sketch type now has a runtime-verified PromQL → backend path
+through the modified-OTLP wire format (typed `Metric.data =
+{DDSketch | KLLSketch | HLLSketch | CountSketch | CountMinSketch}`
+data points). Specifically:
+
+- **`query_statistic` for every sketch accumulator.** Implemented
+  on `DDSketchAccumulator` (Quantile / Sum / Count / Min / Max),
+  `HllSketchAccumulator` (Cardinality, with `Count` accepted as a
+  Cardinality alias for the existing PromQL `count(...)` path),
+  `CountSketchAccumulator` (Topk / Count / Sum, no-key fallback
+  returns row-mean total), and `CountMinSketchAccumulator` (Count
+  / Sum, no-key fallback returns the min-row sum — the canonical
+  CMS total-event estimator that's exact when each insert
+  increments one cell per row).
+- **`accumulator_factory.rs`**: `DDSketchAccumulatorUpdater` wired
+  in (alongside CMS / CountSketch / KLL / HLL updaters) so the
+  precompute_engine recognises `AggregationType::DDSketch` from the
+  `streaming.yaml` schema.
+- **Modified-OTLP envelope decoders** for each sketch type land via
+  the agent processors using sketchlib-go's `SerializePortable` /
+  `SerializeMsgpack`; the backend's `from_sketchlib_proto_bytes` /
+  `from_msgpack_bytes` constructors round-trip through
+  `SketchEnvelope { sketch_state: Some(SketchState::*(state)) }`.
+
+### Known reconciliation gap (cleanup, not a blocker)
+
+- `compatible_agg_types` in
+  [`asap_types/src/capability_matching.rs`](asap-common/dependencies/rs/asap_types/src/capability_matching.rs)
+  does not list `CountMinSketch` under `Statistic::Sum`, but
+  [`promql_utilities/src/query_logics/logics.rs`](asap-common/dependencies/rs/promql_utilities/src/query_logics/logics.rs)
+  treats CMS as the canonical approximator for both `Sum` and
+  `Count`. The runtime e2e succeeds because the inference YAML's
+  exact-match `find_query_config` path bypasses
+  `find_compatible_aggregation`. Two tables → one table is the
+  right cleanup.
+- CMS query without a paired `SetAggregator` /
+  `DeltaSetAggregator` returns total volume, not per-key
+  frequency. To drive `topk(N, …)` over CMS-tracked keys we need
+  a key-aggregator processor on the agent. Tracked as a paper
+  follow-up; out of scope for v1.
+
 ## For paper submission (blocker)
 
 ### 1. Cold-query fallback — §5.2 of the sketch-DB design — **done (local-FS cold store)**
@@ -30,10 +73,15 @@ Follow-ups (not paper-blocking):
   of `DataCollector/TODO.md`). The three-way query harness in
   [#66](https://github.com/ProjectASAP/ASAPQuery-backend/pull/66)
   (`benchmarks/run_full_eval.sh`) is the runner that will produce
-  this number once `asap-query-engine`'s `main.rs` wires the
-  `ASAP_COLD_STORE_ROOT` flag into the `prometheus_promql_with_cold`
-  constructor — until then it exercises only the Prom-forwarding
-  fallback leg.
+  this number.
+- ~~`asap-query-engine` `main.rs` wiring of `ASAP_COLD_STORE_ROOT`~~
+  **done (P1, 2026-04-30).** `--cold-store-root` flag with
+  `env = "ASAP_COLD_STORE_ROOT"` plumbed into a
+  `build_adapter_config` helper that selects
+  `prometheus_promql_with_cold` when set. Four unit tests pin
+  the wiring matrix (cold × forward). Combine with
+  `--forward-unsupported-queries` to keep Prom as the tail of
+  the chain; without it, unsupported shapes return empty.
 
 ### 2. Accuracy-profile library per sketch type
 

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

@@ -9,7 +9,7 @@ use query_engine_rust::precompute_engine::config::{LateDataPolicy, PrecomputeEng
 use query_engine_rust::precompute_engine::output_sink::{RawPassthroughSink, StoreOutputSink};
 use query_engine_rust::precompute_engine::PrecomputeEngine;
 use query_engine_rust::stores::SimpleMapStore;
-use query_engine_rust::{HttpServer, HttpServerConfig};
+use query_engine_rust::{HttpServer, HttpServerConfig, OtlpReceiver, OtlpReceiverConfig};
 use std::sync::Arc;
 use tracing::info;
 use tracing_subscriber::fmt::format::FmtSpan;
@@ -95,6 +95,68 @@ struct Args {
     /// Tier-2 part-cache byte budget, in MiB. 0 disables.
     #[arg(long)]
     persistence_part_cache_mb: Option<u64>,
+
+    /// Root of the §5.2 cold-tier raw-sample store. When set,
+    /// capability-miss queries first try the hour-bucketed JSONL
+    /// layout under this root. Combine with
+    /// `--forward-unsupported-queries` to keep Prometheus as the
+    /// tail of the fallback chain. Reads from `ASAP_COLD_STORE_ROOT`
+    /// so containerised deploys can wire it via env (matches the
+    /// backend Docker image's environment in
+    /// `deploy/docker-compose/base.yml`).
+    #[arg(long, env = "ASAP_COLD_STORE_ROOT")]
+    cold_store_root: Option<std::path::PathBuf>,
+
+    /// Upstream Prometheus URL for the tail of the fallback chain.
+    /// Only consulted when `--forward-unsupported-queries` is set.
+    #[arg(long, default_value = "http://localhost:9090")]
+    prometheus_server: String,
+
+    /// Forward unsupported PromQL shapes to Prometheus rather than
+    /// returning empty.
+    #[arg(long, default_value_t = false)]
+    forward_unsupported_queries: bool,
+
+    /// Path to the inference config YAML — maps query patterns to
+    /// aggregation IDs so the query engine can pick the right
+    /// stored sketch for an incoming PromQL. Without it the query
+    /// engine starts with an empty pattern table and every query
+    /// "no matches" → falls through to cold/Prom fallback. Same
+    /// schema as `query_engine_rust --config`.
+    #[arg(long)]
+    inference_config: Option<String>,
+
+    /// Prometheus-equivalent scrape interval (seconds). Used by
+    /// SimpleEngine when computing the instant-query lookback
+    /// window: each `query` resolves the metric over the last
+    /// `scrape_interval` seconds. For tumbling-window
+    /// aggregations this MUST be ≥ the window size in
+    /// `streaming-config`, otherwise a window's bucket end
+    /// timestamp falls outside the lookback and the query
+    /// returns empty even though data is in the store. Default 30
+    /// matches the e2e harness's 30s window. Old default was 15
+    /// (kept as a deprecated alias).
+    #[arg(long, default_value_t = 30)]
+    prometheus_scrape_interval: u64,
+
+    /// Enable OTLP metrics ingest (gRPC + HTTP). Required for the
+    /// e2e harness's warm-tier sketch path: `query_engine_rust`'s
+    /// patched proto deserialiser handles `DDSketch` / `KLLSketch` /
+    /// `CountSketch` / `CountMinSketch` / `HLLSketch` types that
+    /// the stock OTel `prometheusremotewrite` exporter would
+    /// otherwise drop.
+    #[arg(long, default_value_t = false)]
+    enable_otel_ingest: bool,
+
+    /// OTLP gRPC listen port (only consulted when
+    /// `--enable-otel-ingest` is set).
+    #[arg(long, default_value_t = 4317)]
+    otel_grpc_port: u16,
+
+    /// OTLP HTTP listen port (only consulted when
+    /// `--enable-otel-ingest` is set).
+    #[arg(long, default_value_t = 4318)]
+    otel_http_port: u16,
 }
 
 #[tokio::main]
@@ -178,23 +240,40 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 
     // Optionally start the query HTTP server
     if args.query_port > 0 {
-        let inference_config =
-            InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::CircularBuffer);
+        let inference_config = match args.inference_config.as_deref() {
+            Some(path) => query_engine_rust::utils::file_io::read_inference_config(
+                path,
+                QueryLanguage::promql,
+            )?,
+            None => InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::CircularBuffer),
+        };
+        info!(
+            "Loaded inference config with {} query configs",
+            inference_config.query_configs.len()
+        );
         let query_engine = Arc::new(SimpleEngine::new(
             store.clone(),
             inference_config,
             streaming_config.clone(),
-            15, // default prometheus scrape interval
+            args.prometheus_scrape_interval, // default 30s (matches e2e window size)
             QueryLanguage::promql,
         ));
+        if let Some(root) = args.cold_store_root.as_deref() {
+            info!(
+                cold_store_root = %root.display(),
+                prom_tail = args.forward_unsupported_queries,
+                "Cold-tier fallback enabled (§5.2 cold store)",
+            );
+        }
+        let adapter_config = AdapterConfig::from_prom_with_optional_cold(
+            args.prometheus_server.clone(),
+            args.forward_unsupported_queries,
+            args.cold_store_root.as_deref(),
+        );
         let http_config = HttpServerConfig {
             port: args.query_port,
             handle_http_requests: true,
-            adapter_config: AdapterConfig {
-                protocol: query_engine_rust::data_model::QueryProtocol::PrometheusHttp,
-                language: QueryLanguage::promql,
-                fallback: None,
-            },
+            adapter_config,
         };
         let http_server = HttpServer::new(http_config, query_engine, store.clone(), None);
         tokio::spawn(async move {
@@ -227,15 +306,52 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
             Arc::new(StoreOutputSink::new(store))
         };
 
-    // Build and run the engine
+    // Build the engine. Snapshot `ingest_state` BEFORE starting the
+    // engine — once `engine.run()` is awaited it owns the engine
+    // and we can't pull the handle out for the OTLP receiver.
     let engine = PrecomputeEngine::new(
         engine_config,
         query_engine_rust::data_model::HotReloadStreamingConfig::from_arc(streaming_config),
         output_sink,
     );
+    let ingest_state = if args.enable_otel_ingest {
+        Some(engine.ingest_state())
+    } else {
+        None
+    };
+
+    // Spawn the OTLP receiver alongside the engine when requested.
+    // Without it the warm-tier sketch path doesn't get fed: the
+    // gateway's PRW translator drops `DDSketch` / `HLLSketch` types,
+    // so the only way to deliver sketches to the backend is OTLP.
+    let otel_handle = if let Some(ingest_state) = ingest_state {
+        let otel_config = OtlpReceiverConfig {
+            grpc_port: args.otel_grpc_port,
+            http_port: args.otel_http_port,
+        };
+        info!(
+            grpc_port = args.otel_grpc_port,
+            http_port = args.otel_http_port,
+            "Starting OTLP receiver wired to precompute engine",
+        );
+        let receiver = OtlpReceiver::with_ingest_state(otel_config, ingest_state);
+        Some(tokio::spawn(async move {
+            if let Err(e) = receiver.run().await {
+                tracing::error!("OTLP receiver error: {}", e);
+            }
+        }))
+    } else {
+        None
+    };
 
     info!("Starting precompute engine...");
-    engine.run().await?;
+    let run_result = engine.run().await;
+
+    if let Some(h) = otel_handle {
+        h.abort();
+        let _ = h.await;
+    }
 
+    run_result?;
     Ok(())
 }

asap-query-engine/src/drivers/query/adapters/config.rs

@@ -102,6 +102,33 @@ impl AdapterConfig {
         )
     }
 
+    /// Pick between [`Self::prometheus_promql`] and
+    /// [`Self::prometheus_promql_with_cold`] based on whether the
+    /// caller has a cold-store root configured (`--cold-store-root`
+    /// CLI flag or `ASAP_COLD_STORE_ROOT` env var).
+    ///
+    /// Wired into both binaries that face deployment:
+    /// `query_engine_rust` (`src/main.rs`) and `precompute_engine`
+    /// (`src/bin/precompute_engine.rs`). Centralised here so the
+    /// behaviour matrix only lives in one place.
+    pub fn from_prom_with_optional_cold(
+        prometheus_server: String,
+        forward_unsupported: bool,
+        cold_store_root: Option<&std::path::Path>,
+    ) -> Self {
+        match cold_store_root {
+            Some(root) => {
+                let prom = if forward_unsupported {
+                    Some(prometheus_server)
+                } else {
+                    None
+                };
+                Self::prometheus_promql_with_cold(root.to_path_buf(), prom)
+            }
+            None => Self::prometheus_promql(prometheus_server, forward_unsupported),
+        }
+    }
+
     /// Create a configuration for ClickHouse HTTP with SQL
     /// Convenience constructor for ClickHouse adapter
     pub fn clickhouse_sql(base_url: String, database: String, forward_unsupported: bool) -> Self {

asap-query-engine/src/main.rs

@@ -83,6 +83,21 @@ struct Args {
     #[arg(long)]
     forward_unsupported_queries: bool,
 
+    /// Root directory of the §5.2 cold-tier raw-sample store.
+    /// When set, capability-miss queries first try the
+    /// hour-bucketed JSONL layout under this root
+    /// (`raw/<metric>/YYYY/MM/DD/HH/part-NNNNNN.jsonl` —
+    /// byte-identical to the S3 key layout, see
+    /// `drivers::query::fallback::cold_store::format`) and only
+    /// fall through to Prometheus when the cold adapter can't
+    /// answer the query shape. Combine with
+    /// `--forward-unsupported-queries` to keep Prom as the tail
+    /// of the chain; without it, unsupported shapes return empty
+    /// instead of forwarding. Also reads from `ASAP_COLD_STORE_ROOT`
+    /// so containerised deploys can wire it via env.
+    #[arg(long, env = "ASAP_COLD_STORE_ROOT")]
+    cold_store_root: Option<std::path::PathBuf>,
+
     /// Kafka broker address
     #[arg(long, default_value = "localhost:9092")]
     kafka_broker: String,
@@ -617,9 +632,17 @@ async fn main() -> Result<()> {
     //);
 
     // Original Prometheus config (commented out temporarily):
-    let adapter_config = AdapterConfig::prometheus_promql(
+    if let Some(root) = args.cold_store_root.as_deref() {
+        info!(
+            cold_store_root = %root.display(),
+            prom_tail = args.forward_unsupported_queries,
+            "Cold-tier fallback enabled (§5.2 cold store)",
+        );
+    }
+    let adapter_config = AdapterConfig::from_prom_with_optional_cold(
         args.prometheus_server.clone(),
         args.forward_unsupported_queries,
+        args.cold_store_root.as_deref(),
     );
 
     let http_config = HttpServerConfig {
@@ -911,3 +934,63 @@ fn setup_logging(
     info!("Logs will be written to: {}/query_engine.log", output_dir);
     Ok(guard)
 }
+
+#[cfg(test)]
+mod tests {
+    use query_engine_rust::drivers::AdapterConfig;
+
+    #[test]
+    fn no_cold_no_forward_yields_no_fallback() {
+        let cfg = AdapterConfig::from_prom_with_optional_cold(
+            "http://prom:9090".into(),
+            false,
+            None,
+        );
+        assert!(
+            cfg.fallback.is_none(),
+            "without cold-store and without forward, no fallback should be installed",
+        );
+    }
+
+    #[test]
+    fn no_cold_with_forward_yields_prom_fallback() {
+        let cfg = AdapterConfig::from_prom_with_optional_cold(
+            "http://prom:9090".into(),
+            true,
+            None,
+        );
+        assert!(
+            cfg.fallback.is_some(),
+            "forward_unsupported=true must install Prom fallback",
+        );
+    }
+
+    #[test]
+    fn cold_store_set_installs_fallback_even_without_forward() {
+        // Key wiring claim: setting --cold-store-root alone is
+        // sufficient to engage the §5.2 cold path. The only thing
+        // forward_unsupported_queries adds in that case is the Prom
+        // tail of the chain.
+        let tmp = tempfile::TempDir::new().unwrap();
+        let cfg = AdapterConfig::from_prom_with_optional_cold(
+            "http://prom:9090".into(),
+            false,
+            Some(tmp.path()),
+        );
+        assert!(
+            cfg.fallback.is_some(),
+            "cold-store-root must install ColdFallback regardless of forward_unsupported",
+        );
+    }
+
+    #[test]
+    fn cold_store_with_forward_installs_full_chain() {
+        let tmp = tempfile::TempDir::new().unwrap();
+        let cfg = AdapterConfig::from_prom_with_optional_cold(
+            "http://prom:9090".into(),
+            true,
+            Some(tmp.path()),
+        );
+        assert!(cfg.fallback.is_some());
+    }
+}