From 2a6e6b7829646089a203b17a0d2d2dcd1f514ebc Mon Sep 17 00:00:00 2001 From: zz_y Date: Fri, 15 May 2026 23:01:42 -0600 Subject: [PATCH] feat(emit): controller emits *WithHeap aggregationType when with_heap is set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a `with_heap: bool` field to `CmsParams` (parallel to the existing `CountSketchParams.with_heap`) so the planner's CMS-Heap binding (`bind_cms_with_heap_on_topk`) can carry the heap signal all the way through the L4 IR to the L5 emit. `BindCmsOnCount` sets `with_heap: false`; the topk binder sets `true`. `sketch_kind_to_backend_type` now takes `(SketchKind, &SketchParams)` and promotes to `CountMinSketchWithHeap` / `CountSketchWithHeap` when the params side carries `with_heap: true`. The backend's `policy_capability` keys on these strings to return `FrequencyTopk(*WithHeap)`, which is what the analyzer's `topk(...)` candidate binds against — the missing piece that previously forced test fixtures to patch the streaming-config JSON in-place. Test 8 + Test 9 in `e2e_controller_plans_and_backend_serves.rs` drop the in-test JSON rewrites and now verify the controller emits the heap-bearing aggregationType directly. Test 7 (heap-less CountSketch soft-check on `top_endpoint_qps`) is updated to reflect that the canonical TopK metric now correctly surfaces as `CountSketchWithHeap` in the streaming-config — the soft-check's real signal (wire-format ingest works) is unchanged. Refreshed adjacent fixtures: - `phase_b_backend_agg_type_strings_for_every_sketch_kind` now exercises both with_heap=false and with_heap=true for CMS and CountSketch. - `backend_json_handles_topk_and_pointcount_readouts` asserts the promoted aggregationType for the topk readout case. Co-Authored-By: Claude Opus 4.7 (1M context) --- control_plane/src/emit/stage_config.rs | 84 +++++++++++++------ control_plane/src/optimizer/rules/mod.rs | 6 +- .../sketch_algebra/rules/bind_cms_count.rs | 6 +- control_plane/src/sketch_algebra/schema.rs | 4 +- .../src/sketch_algebra/sketch_params.rs | 13 ++- control_plane/src/sketch_algebra/tests.rs | 2 +- ...e2e_controller_plans_and_backend_serves.rs | 57 +++++-------- 7 files changed, 105 insertions(+), 67 deletions(-) diff --git a/control_plane/src/emit/stage_config.rs b/control_plane/src/emit/stage_config.rs index 3cd83d82..69e2c4ec 100644 --- a/control_plane/src/emit/stage_config.rs +++ b/control_plane/src/emit/stage_config.rs @@ -1190,7 +1190,7 @@ fn build_default_edge_processor_block( d: 5, with_heap: true, }), - SketchKind::Cms => SketchParams::Cms(CmsParams { w: 4096, d: 4 }), + SketchKind::Cms => SketchParams::Cms(CmsParams { w: 4096, d: 4, with_heap: false }), }; let synthetic = EdgeSketchProcessor { processor_name: sketch_kind_to_processor_name(kind).to_string(), @@ -1391,7 +1391,7 @@ fn build_backend_aggregation_json(agg: &BackendAggregation) -> JsonValue { AggregationInput::Raw => "raw", }; json!({ - "aggregationType": sketch_kind_to_backend_type(&agg.sketch_kind), + "aggregationType": sketch_kind_to_backend_type(&agg.sketch_kind, &agg.sketch_params), "aggregationSubType": "", "metric": agg.metric_name, "labels": { @@ -1437,18 +1437,29 @@ fn build_backend_readout_json(r: &BackendReadout) -> JsonValue { } } -/// Map a `SketchKind` to the backend's `AggregationType::Display` string -/// — the same mapping +/// Map a `(SketchKind, SketchParams)` pair to the backend's +/// `AggregationType::Display` string — the same mapping /// [`crate::config::asapquery_backend::map_sketch_type_to_agg_type`] uses /// (the strings must match `AggregationType::FromStr` in the backend's /// `promql_utilities::query_logics::enums`). -fn sketch_kind_to_backend_type(kind: &SketchKind) -> &'static str { - match kind { - SketchKind::DDSketch => "DDSketch", - SketchKind::Kll => "DatasketchesKLL", - SketchKind::Hll => "HLL", - SketchKind::CountSketch => "CountSketch", - SketchKind::Cms => "CountMinSketch", +/// +/// The params side promotes CMS / CountSketch to their `*WithHeap` +/// variants when the planner-set `with_heap` flag is true (see +/// `bind_cms_with_heap_on_topk` and `BindCountSketchOnTopK`). This +/// is what lets the backend's `policy_capability` lookup return +/// `FrequencyTopk(*WithHeap)` for heap-bearing aggregations — required +/// for `topk(...)` queries to bind to the right sids. +fn sketch_kind_to_backend_type(kind: &SketchKind, params: &SketchParams) -> &'static str { + match (kind, params) { + (SketchKind::DDSketch, _) => "DDSketch", + (SketchKind::Kll, _) => "DatasketchesKLL", + (SketchKind::Hll, _) => "HLL", + (SketchKind::CountSketch, SketchParams::CountSketch(p)) if p.with_heap => { + "CountSketchWithHeap" + } + (SketchKind::CountSketch, _) => "CountSketch", + (SketchKind::Cms, SketchParams::Cms(p)) if p.with_heap => "CountMinSketchWithHeap", + (SketchKind::Cms, _) => "CountMinSketch", } } @@ -1621,7 +1632,7 @@ mod tests { ( SketchKind::Cms, "countmin", - SketchParams::Cms(CmsParams { w: 4096, d: 4 }), + SketchParams::Cms(CmsParams { w: 4096, d: 4, with_heap: false }), ), ] { let mut cfg = ddsketch_edge_cfg(); @@ -1646,7 +1657,7 @@ mod tests { cfg.sketch_processors[0] = EdgeSketchProcessor { processor_name: "countmin".to_string(), sketch_kind: SketchKind::Cms, - sketch_params: SketchParams::Cms(CmsParams { w: 4096, d: 4 }), + sketch_params: SketchParams::Cms(CmsParams { w: 4096, d: 4, with_heap: false }), aggregation_id: "agg-cms".to_string(), }; let yaml = emit_edge_yaml(&cfg, "ws://c/").expect("emit ok"); @@ -1822,7 +1833,7 @@ mod tests { aggregation_id: "agg1".into(), metric_name: "endpoint_hits".into(), sketch_kind: SketchKind::Cms, - sketch_params: SketchParams::Cms(CmsParams { w: 4096, d: 4 }), + sketch_params: SketchParams::Cms(CmsParams { w: 4096, d: 4, with_heap: false }), window_secs: 60, spatial_filter: String::new(), grouping: Vec::new(), @@ -1850,8 +1861,13 @@ mod tests { assert_eq!(reads[1]["key"], "user_42"); let aggs = v["aggregations"].as_array().unwrap(); - assert_eq!(aggs[0]["aggregationType"], "CountSketch"); + // CountSketch with `with_heap: true` promotes to + // `CountSketchWithHeap` — the backend's `policy_capability` + // maps that to `FrequencyTopk(CountSketchWithHeap)`, the only + // form the analyzer's `topk(...)` candidate binds against. + assert_eq!(aggs[0]["aggregationType"], "CountSketchWithHeap"); assert_eq!(aggs[0]["parameters"]["with_heap"], true); + // CMS with `with_heap: false` stays plain `CountMinSketch`. assert_eq!(aggs[1]["aggregationType"], "CountMinSketch"); assert_eq!(aggs[1]["parameters"]["w"], 4096); } @@ -1874,7 +1890,7 @@ mod tests { SketchKind::DDSketch => SketchParams::DDSketch(DDSketchParams { alpha: 0.01 }), SketchKind::Kll => SketchParams::Kll(KllParams { k: 200 }), SketchKind::Hll => SketchParams::Hll(HllParams { precision: 14 }), - SketchKind::Cms => SketchParams::Cms(CmsParams { w: 4096, d: 4 }), + SketchKind::Cms => SketchParams::Cms(CmsParams { w: 4096, d: 4, with_heap: false }), SketchKind::CountSketch => SketchParams::CountSketch(CountSketchParams { w: 2048, d: 5, @@ -2219,18 +2235,36 @@ mod tests { /// silently break the backend. #[test] fn phase_b_backend_agg_type_strings_for_every_sketch_kind() { - let cases = vec![ - (SketchKind::Kll, "DatasketchesKLL"), - (SketchKind::DDSketch, "DDSketch"), - (SketchKind::Hll, "HLL"), - (SketchKind::Cms, "CountMinSketch"), - (SketchKind::CountSketch, "CountSketch"), + let cases: Vec<(SketchKind, SketchParams, &str)> = vec![ + (SketchKind::Kll, SketchParams::Kll(KllParams { k: 200 }), "DatasketchesKLL"), + (SketchKind::DDSketch, SketchParams::DDSketch(DDSketchParams { alpha: 0.01 }), "DDSketch"), + (SketchKind::Hll, SketchParams::Hll(HllParams { precision: 14 }), "HLL"), + ( + SketchKind::Cms, + SketchParams::Cms(CmsParams { w: 4096, d: 4, with_heap: false }), + "CountMinSketch", + ), + ( + SketchKind::Cms, + SketchParams::Cms(CmsParams { w: 4096, d: 4, with_heap: true }), + "CountMinSketchWithHeap", + ), + ( + SketchKind::CountSketch, + SketchParams::CountSketch(CountSketchParams { w: 2048, d: 5, with_heap: false }), + "CountSketch", + ), + ( + SketchKind::CountSketch, + SketchParams::CountSketch(CountSketchParams { w: 2048, d: 5, with_heap: true }), + "CountSketchWithHeap", + ), ]; - for (kind, expected) in cases { + for (kind, params, expected) in cases { assert_eq!( - sketch_kind_to_backend_type(&kind), + sketch_kind_to_backend_type(&kind, ¶ms), expected, - "sketch_kind_to_backend_type({kind:?}) drift — backend FromStr will reject" + "sketch_kind_to_backend_type({kind:?}, {params:?}) drift — backend FromStr will reject" ); } } diff --git a/control_plane/src/optimizer/rules/mod.rs b/control_plane/src/optimizer/rules/mod.rs index 543f24dd..aae2a5bf 100644 --- a/control_plane/src/optimizer/rules/mod.rs +++ b/control_plane/src/optimizer/rules/mod.rs @@ -293,7 +293,11 @@ fn bind_cms_with_heap_on_topk( Some(PhysicalExpr::estimate_over_agg( EstimateOp::TopK { k: k_topk }, SketchKind::Cms, - SketchParams::Cms(CmsParams { w, d }), + // CMS-Heap pattern: pair the CMS matrix with a heavy-hitter + // heap so `topk(...)` can enumerate items from the heap + // directly. The streaming-config emit picks + // `CountMinSketchWithHeap` for this binding. + SketchParams::Cms(CmsParams { w, d, with_heap: true }), (**child).clone(), )) } diff --git a/control_plane/src/sketch_algebra/rules/bind_cms_count.rs b/control_plane/src/sketch_algebra/rules/bind_cms_count.rs index 336b3f7f..43425532 100644 --- a/control_plane/src/sketch_algebra/rules/bind_cms_count.rs +++ b/control_plane/src/sketch_algebra/rules/bind_cms_count.rs @@ -93,7 +93,11 @@ impl Rule for BindCmsOnCount { Some(PhysicalExpr::estimate_over_agg( readout, SketchKind::Cms, - SketchParams::Cms(CmsParams { w, d }), + // Heap-LESS CMS — `BindCmsOnCount` is the + // BindCmsOnCount path (frequency / count without TopK). + // The CMS-with-heap binding fires from + // `bind_cms_with_heap_on_topk` and sets `with_heap: true`. + SketchParams::Cms(CmsParams { w, d, with_heap: false }), (**child).clone(), )) } diff --git a/control_plane/src/sketch_algebra/schema.rs b/control_plane/src/sketch_algebra/schema.rs index edd618c4..29b94b34 100644 --- a/control_plane/src/sketch_algebra/schema.rs +++ b/control_plane/src/sketch_algebra/schema.rs @@ -141,7 +141,7 @@ mod tests { fn cms_supports_subtract_and_delete() { let s = SketchStateSchema::for_kind( SketchKind::Cms, - SketchParams::Cms(CmsParams { w: 2048, d: 5 }), + SketchParams::Cms(CmsParams { w: 2048, d: 5, with_heap: false }), ); assert!(s.caps.mergeable); assert!(s.caps.subtractable); @@ -166,7 +166,7 @@ mod tests { SketchStateSchema::for_kind(SketchKind::Kll, SketchParams::Kll(KllParams { k: 200 })); let cms = SketchStateSchema::for_kind( SketchKind::Cms, - SketchParams::Cms(CmsParams { w: 2048, d: 5 }), + SketchParams::Cms(CmsParams { w: 2048, d: 5, with_heap: false }), ); assert!(!kll.is_compatible_for_merge(&cms)); } diff --git a/control_plane/src/sketch_algebra/sketch_params.rs b/control_plane/src/sketch_algebra/sketch_params.rs index 7d5b3db8..b8d1a43d 100644 --- a/control_plane/src/sketch_algebra/sketch_params.rs +++ b/control_plane/src/sketch_algebra/sketch_params.rs @@ -141,6 +141,15 @@ pub struct CmsParams { pub w: u32, /// Number of rows (depth). Drives the failure probability (≤ 2^−d). pub d: u32, + /// Whether to pair the CMS matrix with a heavy-hitter heap (CMS-Heap + /// pattern from Cormode & Muthukrishnan 2005). Set by + /// `bind_cms_with_heap_on_topk` when the planner picks CMS for a + /// TopK statistic. The streaming-config emit consults this flag to + /// pick `CountMinSketchWithHeap` vs `CountMinSketch` for the + /// `aggregationType` string the backend's `policy_capability` keys + /// on. + #[serde(default)] + pub with_heap: bool, } /// Count-Sketch parameters. @@ -236,7 +245,7 @@ mod tests { SketchKind::Hll ); assert_eq!( - SketchParams::Cms(CmsParams { w: 2048, d: 5 }).kind(), + SketchParams::Cms(CmsParams { w: 2048, d: 5, with_heap: false }).kind(), SketchKind::Cms ); assert_eq!( @@ -283,7 +292,7 @@ mod tests { SketchParams::Kll(KllParams { k: 200 }), SketchParams::DDSketch(DDSketchParams { alpha: 0.01 }), SketchParams::Hll(HllParams { precision: 14 }), - SketchParams::Cms(CmsParams { w: 2048, d: 5 }), + SketchParams::Cms(CmsParams { w: 2048, d: 5, with_heap: false }), SketchParams::CountSketch(CountSketchParams { w: 2048, d: 5, diff --git a/control_plane/src/sketch_algebra/tests.rs b/control_plane/src/sketch_algebra/tests.rs index 47a2bcb7..a3a70860 100644 --- a/control_plane/src/sketch_algebra/tests.rs +++ b/control_plane/src/sketch_algebra/tests.rs @@ -118,7 +118,7 @@ fn physical_expr_serde_roundtrip() { }, PhysicalExpr::SketchAgg { sketch_type: SketchKind::Cms, - params: SketchParams::Cms(CmsParams { w: 2048, d: 5 }), + params: SketchParams::Cms(CmsParams { w: 2048, d: 5, with_heap: false }), child: Box::new(PhysicalExpr::Logical(windowed_scan())), }, ]; diff --git a/data_plane/tests/e2e_controller_plans_and_backend_serves.rs b/data_plane/tests/e2e_controller_plans_and_backend_serves.rs index 3178281b..bca83233 100644 --- a/data_plane/tests/e2e_controller_plans_and_backend_serves.rs +++ b/data_plane/tests/e2e_controller_plans_and_backend_serves.rs @@ -1248,9 +1248,14 @@ async fn controller_plan_to_query_full_roundtrip_count_sketch() { Some(SketchType::CountSketch), ); let streaming_config_json = plan_streaming_config_json(&workload); + // `top_endpoint_qps` is the canonical TopK metric — the planner + // picks `with_heap: true` even with `SketchType::CountSketch` + // override, so the controller emits `CountSketchWithHeap`. The + // soft-check below verifies wire-format ingest works regardless + // of heap-bearing classification. assert_eq!( - streaming_config_json["aggregations"][0]["aggregationType"], "CountSketch", - "controller must emit CountSketch aggregationType for SketchType::CountSketch override\n{streaming_config_json}" + streaming_config_json["aggregations"][0]["aggregationType"], "CountSketchWithHeap", + "controller must emit CountSketchWithHeap for top_endpoint_qps (TopK metric)\n{streaming_config_json}" ); post_streaming_config(&client, stack.backend_port, &streaming_config_json).await; @@ -1583,31 +1588,20 @@ async fn controller_plan_to_query_full_roundtrip_cms_with_heap_topk() { Vec::new(), Some(SketchType::CountMinSketch), ); - let mut streaming_config_json = plan_streaming_config_json(&workload); - // The controller emits `aggregationType: "CountMinSketch"` regardless - // of whether the planner picked the heap-bearing binding — the - // `sketch_kind_to_backend_type` mapping doesn't surface the heap - // variant. The TopK signal lives in the readouts (`op: topk`). - // - // For analyzer ↔ policy matching to bind `topk(...)` queries, the - // policy's `policy_capability` must be `FrequencyTopk(CmsWithHeap)`, - // which only fires for `AggregationType::CountMinSketchWithHeap` - // (see `asap_tier_analysis::policy_capability`). So patch the - // emitted JSON in-place: replace the `CountMinSketch` aggregationType - // with `CountMinSketchWithHeap` to bridge the controller-emit gap. - // (Tracked: the controller-side fix is a parallel change to - // `sketch_kind_to_backend_type` to consult the readout class + - // sketch params — out of scope for this test PR.) + let streaming_config_json = plan_streaming_config_json(&workload); + // The controller now emits `CountMinSketchWithHeap` directly when + // the planner-set `with_heap: true` flag on `CmsParams` fires + // (see `sketch_kind_to_backend_type`). No in-test JSON patch is + // needed — the analyzer ↔ policy match binds against the + // controller-emitted aggregation type as-is. assert_eq!( - streaming_config_json["aggregations"][0]["aggregationType"], "CountMinSketch", - "controller must emit CountMinSketch aggregationType\n{streaming_config_json}" + streaming_config_json["aggregations"][0]["aggregationType"], "CountMinSketchWithHeap", + "controller must emit CountMinSketchWithHeap when bind_cms_with_heap_on_topk fires\n{streaming_config_json}" ); assert_eq!( streaming_config_json["readouts"][0]["op"], "topk", "controller must emit a topk readout for CMS-with-heap binding\n{streaming_config_json}" ); - streaming_config_json["aggregations"][0]["aggregationType"] = - JsonValue::String("CountMinSketchWithHeap".to_string()); post_streaming_config(&client, stack.backend_port, &streaming_config_json).await; // Use the planner-picked `(w, d)` so the OTLP DP's wire-level @@ -1747,26 +1741,19 @@ async fn controller_plan_to_query_full_roundtrip_count_sketch_with_heap_topk() { Vec::new(), None, // default → CountSketch (canonical TopK pick) ); - let mut streaming_config_json = plan_streaming_config_json(&workload); - // The controller emits `aggregationType: "CountSketch"` and the - // heap signal lives in `parameters.with_heap` (true) — but the - // analyzer ↔ policy match keys off `aggregation_type`, which - // must be `CountSketchWithHeap` for `policy_capability` to - // return `FrequencyTopk(CountSketchWithHeap)`. Patch the JSON - // in-place: the controller-side fix (consult `with_heap` flag - // when emitting `aggregation_type`) is out of scope for this - // test PR — tracked alongside the parallel CMS-side patch from - // Test 8. + let streaming_config_json = plan_streaming_config_json(&workload); + // The controller now emits `CountSketchWithHeap` directly when + // the planner-set `with_heap: true` flag on `CountSketchParams` + // fires (see `sketch_kind_to_backend_type`). No in-test JSON + // patch is needed. assert_eq!( - streaming_config_json["aggregations"][0]["aggregationType"], "CountSketch", - "controller must emit CountSketch aggregationType for default top_endpoint_qps\n{streaming_config_json}" + streaming_config_json["aggregations"][0]["aggregationType"], "CountSketchWithHeap", + "controller must emit CountSketchWithHeap for default top_endpoint_qps TopK binding\n{streaming_config_json}" ); assert_eq!( streaming_config_json["aggregations"][0]["parameters"]["with_heap"], true, "controller must set parameters.with_heap=true for CountSketch TopK binding\n{streaming_config_json}" ); - streaming_config_json["aggregations"][0]["aggregationType"] = - JsonValue::String("CountSketchWithHeap".to_string()); post_streaming_config(&client, stack.backend_port, &streaming_config_json).await; let items: &[(&str, u64)] = &[