From 9ec099e802e22e1c9efa69f8344a4d15ad86bb2c Mon Sep 17 00:00:00 2001
From: zz_y <zeyingz@umd.edu>
Date: Fri, 15 May 2026 23:20:05 -0600
Subject: [PATCH] fix(precompute): factory routes CountSketch[WithHeap] to
 CmsAccumulator
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The precompute accumulator factory was silently falling through to
the catch-all `tracing::warn! + SumAccumulatorUpdater` default for
both `AggregationType::CountSketch` and `AggregationType::CountSketchWithHeap`.
Any raw-input ingest path (precompute worker + sketch_db backfill
window_builder) hitting one of those policies would produce a Sum
accumulator instead of a frequency sketch — silently wrong results.

Route both to `CmsAccumulatorUpdater` (which already handles the
shared `(rows, cols)` matrix shape) for parity with the existing
`CountMinSketch | CountMinSketchWithHeap` arm at line 729. This
is the correctness floor: a registered policy yields a working
accumulator, not a Sum default.

Limitation called out in the doc comment: like the existing
`CountMinSketchWithHeap` arm, this drops the per-policy top-k
heap. Dedicated `CountSketchAccumulatorUpdater` /
`CountMinSketchWithHeapAccumulatorUpdater` impls (the accumulator
structs exist but lack the updater wrapper) are tracked as
follow-up — the present arm only matters for Mode 2 / raw-input
ingest, which the OTLP sketch-envelope path doesn't exercise.

`config_is_keyed` updated correspondingly (the four new variants
are all multi-population by shape, parallel to CountMinSketch),
plus a test covering all four arms.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 .../precompute_engine/accumulator_factory.rs  | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/data_plane/src/precompute_engine/accumulator_factory.rs b/data_plane/src/precompute_engine/accumulator_factory.rs
index 12c76f4a..edeb2bb1 100644
--- a/data_plane/src/precompute_engine/accumulator_factory.rs
+++ b/data_plane/src/precompute_engine/accumulator_factory.rs
@@ -630,6 +630,8 @@ pub fn config_is_keyed(config: &AggregationConfig) -> bool {
             | AggregationType::MultipleMinMax
             | AggregationType::CountMinSketch
             | AggregationType::CountMinSketchWithHeap
+            | AggregationType::CountSketch
+            | AggregationType::CountSketchWithHeap
             | AggregationType::HydraKLL
     )
 }
@@ -730,6 +732,26 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box<dyn Accumul
             let (row_num, col_num) = cms_params(config);
             Box::new(CmsAccumulatorUpdater::new(row_num, col_num))
         }
+        // CountSketch + CountSketchWithHeap (raw-input ingest path):
+        // route to `CmsAccumulatorUpdater` for now — it handles the
+        // same `(rows, cols)` matrix shape. The OTLP modified-sketch
+        // wire path uses `SketchEnvelope` ingest (not raw), so this
+        // arm fires only for Mode 2 / raw-input policies.
+        //
+        // Limitation: like the `CountMinSketchWithHeap` arm above,
+        // this drops the per-policy top-k heap on the raw-input side.
+        // The heap-bearing accumulator
+        // (`count_min_sketch_with_heap_accumulator.rs`) exists but
+        // doesn't yet have an `AccumulatorUpdater` impl; same is
+        // true for the CountSketch variants. Adding dedicated
+        // updaters is tracked as a follow-up — the present arm is
+        // a correctness floor (registered policy → working
+        // accumulator) without silently falling through to
+        // `SumAccumulatorUpdater`.
+        AggregationType::CountSketch | AggregationType::CountSketchWithHeap => {
+            let (row_num, col_num) = cms_params(config);
+            Box::new(CmsAccumulatorUpdater::new(row_num, col_num))
+        }
         AggregationType::HydraKLL => {
             let (row_num, col_num, k) = hydra_kll_params(config);
             Box::new(HydraKllAccumulatorUpdater::new(row_num, col_num, k))
@@ -882,6 +904,18 @@ mod tests {
             AggregationType::CountMinSketch,
             ""
         )));
+        assert!(config_is_keyed(&make_config(
+            AggregationType::CountMinSketchWithHeap,
+            ""
+        )));
+        assert!(config_is_keyed(&make_config(
+            AggregationType::CountSketch,
+            ""
+        )));
+        assert!(config_is_keyed(&make_config(
+            AggregationType::CountSketchWithHeap,
+            ""
+        )));
         assert!(config_is_keyed(&make_config(AggregationType::HydraKLL, "")));
 
         // Verify agreement with updater.is_keyed()