feat(query): heap-bearing CMS + CountSketch topk strict-success

control_plane/src/asap_tier_analysis.rs

@@ -446,6 +446,9 @@ pub fn policy_capability(cfg: &asap_types::AggregationConfig) -> Option<Capabili
         AggregationType::CountMinSketchWithHeap => {
             Some(Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap))
         }
+        AggregationType::CountSketchWithHeap => {
+            Some(Capability::FrequencyTopk(SketchKindHandle::CountSketchWithHeap))
+        }
         // Keyed-multi-population variants. The capability the policy
         // *provides* is the multi-pop variant itself; the matching
         // predicate (`Capability::is_satisfied_by`) recognises that

control_plane/src/query_parser/promql.rs

@@ -197,17 +197,20 @@ fn walk_qe(expr: &Expr, ctx: WalkCtx) -> anyhow::Result<QueryExpr> {
         // (sum / avg / min / max / etc.). PromQL's `count(metric)`
         // operates on the result-set's LABEL-SETS — the inner
         // selector is just "the things to count," not a value to sum.
-        // Synthesizing `Aggregate(Sum)` underneath an outer count would
-        // collect a redundant `ExactAgg(Sum)` candidate alongside the
-        // intended `CardinalityApprox` one, and the engine's
-        // "all candidates must succeed" semantic surfaces a
-        // `CapabilityMiss` when no Sum policy is registered for the
-        // metric (e.g. an HLL-only deploy).
+        // PromQL's `topk(k, metric)` operates on the SERIES — the
+        // inner selector is the population to rank, not a value to
+        // sum. Synthesizing `Aggregate(Sum)` underneath either outer
+        // would collect a redundant `ExactAgg(Sum)` candidate
+        // alongside the intended `CardinalityApprox` /
+        // `FrequencyTopk(*WithHeap)` one; the engine's
+        // "all candidates must succeed" semantic then surfaces a
+        // `CapabilityMiss` when no Sum policy is registered (e.g. an
+        // HLL-only or CMS-with-heap-only deploy).
         Expr::VectorSelector(vs) => {
             let (name, filters) = extract_vs_info(vs);
             let source   = QueryExpr::Source(QeSourceSpec { name });
             let filtered = apply_qe_filters(source, filters);
-            if ctx.outer_count {
+            if ctx.outer_count || ctx.topk.is_some() {
                 Ok(filtered)
             } else {
                 Ok(QueryExpr::Aggregate {

control_plane/src/sketch_algebra/capability.rs

@@ -327,10 +327,15 @@ pub fn capability_for(intent: &AggIntent) -> Option<Capability> {
                 None
             } else {
                 // Top-k is intrinsically heavy-hitter — only heap-bearing
-                // handles can enumerate the items. `CmsWithHeap` is the
-                // canonical handle today; `is_satisfied_by` accepts
-                // either heap-bearing variant against an `Any` required.
-                Some(Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap))
+                // handles can enumerate the items. The analyzer doesn't
+                // care which heap-bearing variant answers (CmsWithHeap or
+                // CountSketchWithHeap both work — the reducer dispatches
+                // both through `decode_cms_with_heap_from_msgpack` and
+                // produces top-k items either way). Return `Any` so
+                // `is_satisfied_by`'s `handles_compatible_for_topk`
+                // wildcard accepts whichever variant the ingest tier
+                // chose to register.
+                Some(Capability::FrequencyTopk(SketchKindHandle::Any))
             }
         }
         AggIntent::Frequency { accuracy } => {
@@ -769,14 +774,19 @@ mod tests {
     }
 
     #[test]
-    fn capability_for_topk_returns_frequency_topk_cms_with_heap() {
+    fn capability_for_topk_returns_frequency_topk_any() {
+        // The analyzer no longer pins a concrete heap-bearing variant
+        // for top-k — `Any` lets `handles_compatible_for_topk` accept
+        // either `CmsWithHeap` or `CountSketchWithHeap` registered at
+        // ingest time. Both variants share the heap envelope and the
+        // reducer dispatches them identically.
         let intent = AggIntent::TopK {
             k: 10,
             accuracy: AccuracyTarget::Epsilon(0.05),
         };
         assert_eq!(
             capability_for(&intent),
-            Some(Capability::FrequencyTopk(SketchKindHandle::CmsWithHeap))
+            Some(Capability::FrequencyTopk(SketchKindHandle::Any))
         );
     }
 

crates/promql_utilities/src/query_logics/enums.rs

@@ -274,6 +274,7 @@ pub enum AggregationType {
     CountMinSketch,
     CountMinSketchWithHeap,
     CountSketch,
+    CountSketchWithHeap,
     // ---------- cardinality / set tracking ----------
     SetAggregator,
     DeltaSetAggregator,
@@ -298,6 +299,7 @@ impl AggregationType {
             AggregationType::CountMinSketch => "CountMinSketch",
             AggregationType::CountMinSketchWithHeap => "CountMinSketchWithHeap",
             AggregationType::CountSketch => "CountSketch",
+            AggregationType::CountSketchWithHeap => "CountSketchWithHeap",
             AggregationType::SetAggregator => "SetAggregator",
             AggregationType::DeltaSetAggregator => "DeltaSetAggregator",
             AggregationType::HLL => "HLL",
@@ -318,6 +320,7 @@ impl AggregationType {
                 | AggregationType::CountMinSketch
                 | AggregationType::CountMinSketchWithHeap
                 | AggregationType::CountSketch
+                | AggregationType::CountSketchWithHeap
                 | AggregationType::HydraKLL
         )
     }
@@ -332,6 +335,7 @@ impl AggregationType {
                 | AggregationType::CountMinSketch
                 | AggregationType::CountMinSketchWithHeap
                 | AggregationType::CountSketch
+                | AggregationType::CountSketchWithHeap
         )
     }
 
@@ -367,6 +371,7 @@ impl FromStr for AggregationType {
             "CountMinSketch" => Ok(AggregationType::CountMinSketch),
             "CountMinSketchWithHeap" => Ok(AggregationType::CountMinSketchWithHeap),
             "CountSketch" => Ok(AggregationType::CountSketch),
+            "CountSketchWithHeap" => Ok(AggregationType::CountSketchWithHeap),
             "SetAggregator" => Ok(AggregationType::SetAggregator),
             "DeltaSetAggregator" => Ok(AggregationType::DeltaSetAggregator),
             "HLL" | "HyperLogLog" => Ok(AggregationType::HLL),
@@ -395,6 +400,7 @@ impl FromStr for AggregationType {
             "CountSketchAccumulator" | "CS" | "cs" | "count_sketch" => {
                 Ok(AggregationType::CountSketch)
             }
+            "CountSketchWithHeapAccumulator" => Ok(AggregationType::CountSketchWithHeap),
             "SetAggregatorAccumulator" => Ok(AggregationType::SetAggregator),
             "DeltaSetAggregatorAccumulator" => Ok(AggregationType::DeltaSetAggregator),
             _ => Err(format!("Unknown aggregation type: '{s}'")),

data_plane/src/drivers/ingest/otel.rs

@@ -1271,7 +1271,7 @@ fn aggregation_type_for_sketch_handle(
         SketchKindHandle::Kll => Some(AggregationType::DatasketchesKLL),
         SketchKindHandle::Hll => Some(AggregationType::HLL),
         SketchKindHandle::CountSketch => Some(AggregationType::CountSketch),
-        SketchKindHandle::CountSketchWithHeap => Some(AggregationType::CountSketch),
+        SketchKindHandle::CountSketchWithHeap => Some(AggregationType::CountSketchWithHeap),
         SketchKindHandle::CountMin => Some(AggregationType::CountMinSketch),
         SketchKindHandle::CmsWithHeap => Some(AggregationType::CountMinSketchWithHeap),
         SketchKindHandle::Any => None,
@@ -1307,8 +1307,15 @@ fn sketch_config_to_params(
         }
         SketchConfig::CountSketch { rows, cols }
         | SketchConfig::CountMin { rows, cols } => {
-            params.insert("rows".to_string(), serde_json::json!(*rows));
-            params.insert("cols".to_string(), serde_json::json!(*cols));
+            // Canonical key mapping (matches the controller's
+            // `sketch_params_to_json` in
+            // `control_plane::emit::stage_config`): `w` is the
+            // matrix width (=cols), `d` is the depth (=rows). The
+            // controller writes `{w, d}` into the streaming-config
+            // `parameters`, so the policy_fp content match has to
+            // probe the same keys.
+            params.insert("w".to_string(), serde_json::json!(*cols));
+            params.insert("d".to_string(), serde_json::json!(*rows));
         }
     }
     params
@@ -1377,7 +1384,26 @@ fn sketch_kind_handle_for(
         SketchKind::DdSketch => SketchKindHandle::DDSketch,
         SketchKind::Kll => SketchKindHandle::Kll,
         SketchKind::Hll => SketchKindHandle::Hll,
-        SketchKind::CountSketch => SketchKindHandle::CountSketch,
+        SketchKind::CountSketch => {
+            // Mirror the CountMin branch: CountSketch-with-heap
+            // payloads share the same outer msgpack envelope
+            // (`CountMinSketchWithHeapSerialized` — see the comment
+            // in `sketch_reducer.rs` at the dispatch site, which
+            // notes both heap-bearing variants reuse this wire
+            // shape since the heap is the distinguishing payload).
+            // Auto-promote to `CountSketchWithHeap` when the bytes
+            // decode AND the heap is non-empty; otherwise stay with
+            // vanilla `CountSketch`.
+            if dp.encoding == ENCODING_MSGPACK {
+                use asap_sketchlib::sketches::countminsketch_topk::CountMinSketchWithHeap;
+                if let Ok(cms) = CountMinSketchWithHeap::deserialize_msgpack(&dp.sketch) {
+                    if !cms.topk_heap_items().is_empty() {
+                        return SketchKindHandle::CountSketchWithHeap;
+                    }
+                }
+            }
+            SketchKindHandle::CountSketch
+        }
         SketchKind::CountMin => {
             // Try a no-cost peek: msgpack-encoded CMS-with-heap payloads
             // round-trip through asap_sketchlib's
@@ -2081,15 +2107,17 @@ mod policy_fp_lookup_tests {
             rows: 4,
             cols: 256,
         });
-        assert_eq!(cs.get("rows"), Some(&serde_json::json!(4)));
-        assert_eq!(cs.get("cols"), Some(&serde_json::json!(256)));
+        // Canonical keys: w (=cols, width) and d (=rows, depth) —
+        // matches `control_plane::emit::stage_config::sketch_params_to_json`.
+        assert_eq!(cs.get("w"), Some(&serde_json::json!(256)));
+        assert_eq!(cs.get("d"), Some(&serde_json::json!(4)));
 
         let cm = sketch_config_to_params(&SketchConfig::CountMin {
             rows: 4,
             cols: 256,
         });
-        assert_eq!(cm.get("rows"), Some(&serde_json::json!(4)));
-        assert_eq!(cm.get("cols"), Some(&serde_json::json!(256)));
+        assert_eq!(cm.get("w"), Some(&serde_json::json!(256)));
+        assert_eq!(cm.get("d"), Some(&serde_json::json!(4)));
     }
 }
 

data_plane/src/query_engines/asap_query_engine/engine.rs

@@ -6411,10 +6411,10 @@ mod analyzer_parity_tests {
     ctrl   OK [metric=http_requests_total gbk=[\"region\", \"zone\"] cap=ExactAgg(Sum) fn=sum args=[] range_s=0]
     engine OK pattern=only_spatial stats=[sum] metric=http_requests_total fn= agg_op=sum range_s=- range_ms=None spatial=\"\" grouping=[\"region\", \"zone\"]
 ─── q06: topk(5, http_requests_total)
-    ctrl   OK [metric=http_requests_total gbk=[] cap=FrequencyTopk(CmsWithHeap) fn=topk args=[5.0] range_s=0 | metric=http_requests_total gbk=[] cap=ExactAgg(Sum) fn=topk args=[5.0] range_s=0]
+    ctrl   OK [metric=http_requests_total gbk=[] cap=FrequencyTopk(Any) fn=topk args=[5.0] range_s=0]
     engine OK pattern=only_spatial stats=[topk] metric=http_requests_total fn= agg_op=topk range_s=- range_ms=None spatial=\"\" grouping=[]
 ─── q07: topk(10, sum by (svc) (m))
-    ctrl   OK [metric=m gbk=[\"svc\"] cap=FrequencyTopk(CmsWithHeap) fn=topk args=[10.0] range_s=0 | metric=m gbk=[\"svc\"] cap=ExactAgg(Sum) fn=topk args=[10.0] range_s=0]
+    ctrl   OK [metric=m gbk=[\"svc\"] cap=FrequencyTopk(Any) fn=topk args=[10.0] range_s=0]
     engine MISS(NoPattern)
 ─── q08: count_over_time(http_requests_total[5m])
     ctrl   MISS(UnsupportedAggIntent(\"count\"))

data_plane/src/storage_engines/sketch_db/accuracy.rs

@@ -213,6 +213,26 @@ impl AccuracyProfile {
                 }
             }
 
+            // CountSketchWithHeap: CountSketch frequency estimator
+            // paired with a top-k heap. Mirrors the
+            // CountMinSketchWithHeap branch above — the CountSketch
+            // half gives ε_point = 1/√w; the heap half gives
+            // ε_heap = 1/heap_size for retention. Report the
+            // tighter (max) of the two.
+            AggregationType::CountSketchWithHeap => {
+                let (rows, cols) = cms_params(config);
+                let heap = cms_heap_size(config);
+                let cs_epsilon = 1.0 / (cols as f64).max(1.0).sqrt();
+                let heap_epsilon = 1.0 / (heap as f64).max(1.0);
+                let epsilon = cs_epsilon.max(heap_epsilon);
+                let delta = 0.5_f64.powi(rows as i32);
+                Self {
+                    epsilon,
+                    delta,
+                    kind: AccuracyKind::TopK,
+                }
+            }
+
             // HLL: std-dev ≈ 1.04/√m, m = 2^precision. Report
             // this as relative error ε; δ is the Gaussian
             // std-dev convention (stored as 0 because our δ