fix(warm-engine): answer replay queries under empty-schema deploy (was status=error)

asap-common/dependencies/rs/asap_types/src/capability_matching.rs

@@ -146,10 +146,20 @@ pub fn compatible_agg_types(stat: Statistic) -> &'static [AggregationType] {
         // Count: exact via MultipleSum (the planner's canonical pick for
         // Count-Exact uses `MultipleSum` with sub_type="count"); approximate
         // via CountMinSketch / CountMinSketchWithHeap.
+        //
+        // HLL is also valid here: the warm-tier MVP demo
+        // (ProjectASAP/ASAPCollector#46) plans `unique_users_per_min`
+        // as an HLL agg and the replay client queries it with
+        // `count(unique_users_per_min)`. `HllSketchAccumulator`
+        // answers `Statistic::Count` as a cardinality alias —
+        // see `precompute_operators/hll_sketch_accumulator.rs:220`.
+        // Without HLL listed here the warm engine returns `status=error`
+        // for every count-of-HLL replay row.
         Statistic::Count => &[
             AggregationType::MultipleSum,
             AggregationType::CountMinSketch,
             AggregationType::CountMinSketchWithHeap,
+            AggregationType::HLL,
         ],
         Statistic::Min | Statistic::Max => {
             &[AggregationType::MinMax, AggregationType::MultipleMinMax]
@@ -190,7 +200,19 @@ pub fn compatible_agg_types(stat: Statistic) -> &'static [AggregationType] {
             AggregationType::DeltaSetAggregator,
             AggregationType::HLL,
         ],
-        Statistic::Topk => &[AggregationType::CountMinSketchWithHeap],
+        // Topk: `CountMinSketchWithHeap` is the canonical CMS-Heap
+        // pattern. CountSketch is the second-tier reservoir-style
+        // approximator the MVP demo's controller plans for
+        // `top_endpoint_qps` (median-of-row estimator over a
+        // signed-counter matrix). `CountSketchAccumulator` answers
+        // `Statistic::Topk` directly — see
+        // `precompute_operators/count_sketch_accumulator.rs:284`.
+        // Without CountSketch listed here, `topk(K, top_endpoint_qps)`
+        // capability-misses and the warm engine returns `status=error`.
+        Statistic::Topk => &[
+            AggregationType::CountMinSketchWithHeap,
+            AggregationType::CountSketch,
+        ],
     }
 }
 
@@ -295,11 +317,29 @@ pub fn window_compatible(config: &AggregationConfig, data_range_ms: Option<u64>)
     }
 }
 
-/// Label compatibility: strict exact match.
-/// TODO: relax to superset (config.grouping_labels ⊇ req.grouping_labels) for
-/// simple accumulators (Sum, MinMax, Increase).
+/// Label compatibility: config can serve a query whose grouping is a
+/// **subset** (including equality) of the config's grouping_labels.
+///
+/// Pre-fix this was strict-exact: `config_labels == req_labels`. The
+/// MVP demo (ProjectASAP/ASAPCollector#46) replays
+/// `count(unique_users_per_min)` / `topk(5, top_endpoint_qps)` with
+/// no `by (...)` modifier, which translates to `req.grouping_labels =
+/// []`. The corresponding agg configs are per-zone (`[zone]` grouping).
+/// Pre-fix every such replay row capability-missed and the warm engine
+/// returned `status=error`. Post-fix the engine accepts the agg, runs
+/// the per-zone accumulators through the merge path
+/// (`execute_and_merge_store_queries` produces a per-key map; the
+/// downstream merge collapses them to the requested `[]` grouping —
+/// HLL/CMS/CountSketch all support natural across-key merge, and
+/// scalar accumulators like Sum / Increase reduce by addition).
+///
+/// Direction is asymmetric: `config ⊇ req` is OK (engine merges away
+/// the extra labels), but `req ⊃ config` is NOT — the engine cannot
+/// invent a label that the materialised agg never partitioned by.
 pub fn labels_compatible(config_labels: &KeyByLabelNames, req_labels: &KeyByLabelNames) -> bool {
-    config_labels == req_labels
+    let req: std::collections::HashSet<&String> = req_labels.labels.iter().collect();
+    let cfg: std::collections::HashSet<&String> = config_labels.labels.iter().collect();
+    req.is_subset(&cfg)
 }
 
 /// Spatial filter compatibility.
@@ -725,8 +765,17 @@ mod tests {
     }
 
     #[test]
-    fn label_strict_superset_rejected() {
-        // Config has {job, instance}, query wants only {job} — strict mode rejects
+    fn label_superset_config_accepts_subset_query() {
+        // Config has `{job, instance}`, query wants only `{job}`.
+        //
+        // Pre-fix `labels_compatible` did strict-eq and rejected this,
+        // which broke the MVP demo (ProjectASAP/ASAPCollector#46): the
+        // agent's per-zone HLL agg has `grouping_labels = [zone]`, the
+        // replay client's `count(unique_users_per_min)` has no `by`
+        // modifier (req grouping = `[]`). Post-fix the agg can serve
+        // the broader-aggregation query — the engine's merge path
+        // collapses the extra label dimension before the result
+        // surface. See `labels_compatible` rustdoc.
         let configs = single_config(make_config(
             1,
             "cpu",
@@ -741,6 +790,38 @@ mod tests {
             &configs,
             &req("cpu", &[Statistic::Sum], Some(300_000), &["job"], ""),
         );
+        assert!(
+            result.is_some(),
+            "post-fix: a config with `[job, instance]` grouping must serve a `[job]`-only req \
+             via the merge path",
+        );
+    }
+
+    #[test]
+    fn label_subset_config_rejects_superset_query() {
+        // Config has only `[job]`, query wants `[job, instance]`.
+        // The engine cannot invent a partition the agg never
+        // materialised, so this remains incompatible.
+        let configs = single_config(make_config(
+            1,
+            "cpu",
+            "Sum",
+            "",
+            300,
+            "tumbling",
+            &["job"],
+            "",
+        ));
+        let result = find_compatible_aggregation(
+            &configs,
+            &req(
+                "cpu",
+                &[Statistic::Sum],
+                Some(300_000),
+                &["job", "instance"],
+                "",
+            ),
+        );
         assert!(result.is_none());
     }
 
@@ -1113,6 +1194,29 @@ mod tests {
                 .contains(&AggregationType::CountMinSketchWithHeap),
             "CountMinSketchWithHeap must be a compatible type for Topk",
         );
+        // CountSketch → Topk (warm-engine-error-on-replay-queries fix).
+        // Required so the MVP demo's `topk(5, top_endpoint_qps)` —
+        // which routes through the agent's `countsketchprocessor`
+        // and lands as a CountSketch-only config — resolves
+        // through capability matching. Without this, the warm
+        // engine returned `status=error` for every topk replay row.
+        assert!(
+            compatible_agg_types(Statistic::Topk).contains(&AggregationType::CountSketch),
+            "CountSketch must be a compatible type for Topk (warm-engine-error fix)",
+        );
+        // HLL → Count (warm-engine-error-on-replay-queries fix). The
+        // MVP demo's `count(unique_users_per_min)` is structurally a
+        // PromQL `Statistic::Count` (the AggregationOperator::Count
+        // → Statistic::Count mapping in
+        // `promql_utilities::query_logics::enums`); the
+        // `HllSketchAccumulator` answers it as a cardinality alias
+        // (`hll_sketch_accumulator.rs:220`). Without HLL listed
+        // here, capability matching missed and the warm engine
+        // returned `status=error` for every count-of-HLL replay row.
+        assert!(
+            compatible_agg_types(Statistic::Count).contains(&AggregationType::HLL),
+            "HLL must be a compatible type for Count (warm-engine-error fix)",
+        );
     }
 
     /// Phase-3.1 regression test for the canonical MVP-demo failure

asap-query-engine/src/engines/simple/engine.rs

@@ -495,6 +495,78 @@ impl SimpleEngine {
         result
     }
 
+    /// Resolve the canonical "all labels" set for a metric, with a
+    /// streaming-config fallback for schema-empty deploys.
+    ///
+    /// The user-facing `inference_config.schema` is the source of truth
+    /// for "what labels does this metric carry" — but the production
+    /// warm-tier deploy launches with `--streaming-config` only and no
+    /// `--config`, so the schema is empty. Pre-fix every query lookup
+    /// in `build_promql_execution_context_tail` and
+    /// `build_query_requirements_promql` returned `None` /
+    /// `KeyByLabelNames::empty()` for that metric, killing capability
+    /// matching (`req.grouping_labels = []` strict-mismatches every
+    /// agg config's `[zone]`) and the downstream context build (the
+    /// `None` short-circuits the whole query). See
+    /// `tests::warm_engine_replay_regression_tests::production_conditions_*`.
+    ///
+    /// Fallback rules:
+    /// 1. Look up the metric in `inference_config.schema`. Return its
+    ///    labels if present.
+    /// 2. Otherwise scan the current `StreamingConfig` snapshot for
+    ///    every agg config whose `metric == name`. Union their
+    ///    `grouping_labels` (the per-series partition the agg
+    ///    materialises) and return that. The union preserves order
+    ///    of first-appearance and de-dupes — `KeyByLabelNames`
+    ///    equality is strict, so we have to keep insertion order
+    ///    deterministic across config swaps.
+    /// 3. Returns `None` only if no agg config references the metric
+    ///    AND the schema is empty. Callers translate that into the
+    ///    same "metric unknown" outcome as before this helper landed.
+    fn resolve_metric_labels(&self, metric: &str) -> Option<KeyByLabelNames> {
+        // (1) schema lookup — user-supplied source of truth.
+        if let SchemaConfig::PromQL(schema) = &self.inference_config.schema {
+            if let Some(labels) = schema.get_labels(metric).cloned() {
+                return Some(labels);
+            }
+        }
+
+        // (2) streaming-config fallback — derived from whatever agg
+        // configs the controller / static YAML registered for the
+        // metric. Produces the union of `grouping_labels` across all
+        // matching aggs in deterministic insertion order.
+        let snap = self.streaming_config_snapshot();
+        let mut seen = std::collections::HashSet::new();
+        let mut union: Vec<String> = Vec::new();
+        // Sort by aggregation_id so the resulting label vector is
+        // stable across re-runs even though `aggregation_configs` is
+        // a `HashMap`. Without this ordering, two engines holding
+        // bit-identical configs could produce different
+        // `KeyByLabelNames` instances and `labels_compatible`'s
+        // strict-eq would flake intermittently.
+        let mut agg_ids: Vec<u64> = snap.aggregation_configs.keys().copied().collect();
+        agg_ids.sort_unstable();
+        for id in agg_ids {
+            let cfg = match snap.get_aggregation_config(id) {
+                Some(c) => c,
+                None => continue,
+            };
+            if cfg.metric != metric {
+                continue;
+            }
+            for label in &cfg.grouping_labels.labels {
+                if seen.insert(label.clone()) {
+                    union.push(label.clone());
+                }
+            }
+        }
+        if union.is_empty() {
+            None
+        } else {
+            Some(KeyByLabelNames::new(union))
+        }
+    }
+
     /// Convert query timestamp (seconds) to data timestamp (milliseconds)
     pub fn convert_query_time_to_data_time(query_time: f64) -> u64 {
         (query_time * 1000.0) as u64
@@ -1502,11 +1574,15 @@ impl SimpleEngine {
     ) -> Option<QueryExecutionContext> {
         let (metric, spatial_filter) = get_metric_and_spatial_filter(match_result);
 
-        let promql_schema = match &self.inference_config.schema {
-            SchemaConfig::PromQL(schema) => schema,
-            _ => return None,
-        };
-        let all_labels = match promql_schema.get_labels(&metric).cloned() {
+        // Resolve the metric's "all labels" set. Falls back to a
+        // streaming-config-derived label union when the schema is
+        // empty for this metric — the production warm-tier deploy
+        // launches with `--streaming-config` only and an empty
+        // schema, and pre-fix every query for a streaming-config-
+        // registered metric blew up here on the schema lookup. See
+        // `Self::resolve_metric_labels` and the
+        // `production_conditions_*` regression tests for context.
+        let all_labels = match self.resolve_metric_labels(&metric) {
             Some(labels) => labels,
             None => {
                 warn!("No metric configuration found for '{}'", metric);
@@ -1952,13 +2028,16 @@ impl SimpleEngine {
                 .map(|d| d.num_seconds() as u64 * 1000),
         };
 
-        let all_labels = match &self.inference_config.schema {
-            SchemaConfig::PromQL(schema) => schema
-                .get_labels(&metric)
-                .cloned()
-                .unwrap_or_else(KeyByLabelNames::empty),
-            _ => KeyByLabelNames::empty(),
-        };
+        // Resolve the metric's "all labels" set with the same
+        // schema-empty fallback used by
+        // `build_promql_execution_context_tail`. Without this
+        // fallback the schema-empty production deploy returns
+        // `KeyByLabelNames::empty()` for every metric, and
+        // `labels_compatible`'s strict-eq mismatches every agg
+        // config's `[zone]` → capability-miss → `status=error`.
+        let all_labels = self
+            .resolve_metric_labels(&metric)
+            .unwrap_or_else(KeyByLabelNames::empty);
 
         let grouping_labels = match query_pattern_type {
             QueryPatternType::OnlyTemporal => all_labels,