feat: CountSketch + KLL + DDSketch + HLL decoders for modified-OTLP hot path (PR C)

asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs

@@ -273,10 +273,12 @@ pub enum AggregationType {
     HydraKLL,
     CountMinSketch,
     CountMinSketchWithHeap,
+    CountSketch,
     // ---------- cardinality / set tracking ----------
     SetAggregator,
     DeltaSetAggregator,
     HLL,
+    DDSketch,
     // ---------- legacy config wrapper names ----------
     SingleSubpopulation,
     MultipleSubpopulation,
@@ -295,9 +297,11 @@ impl AggregationType {
             AggregationType::HydraKLL => "HydraKLL",
             AggregationType::CountMinSketch => "CountMinSketch",
             AggregationType::CountMinSketchWithHeap => "CountMinSketchWithHeap",
+            AggregationType::CountSketch => "CountSketch",
             AggregationType::SetAggregator => "SetAggregator",
             AggregationType::DeltaSetAggregator => "DeltaSetAggregator",
             AggregationType::HLL => "HLL",
+            AggregationType::DDSketch => "DDSketch",
             AggregationType::SingleSubpopulation => "SingleSubpopulation",
             AggregationType::MultipleSubpopulation => "MultipleSubpopulation",
         }
@@ -313,6 +317,7 @@ impl AggregationType {
                 | AggregationType::MultipleMinMax
                 | AggregationType::CountMinSketch
                 | AggregationType::CountMinSketchWithHeap
+                | AggregationType::CountSketch
                 | AggregationType::HydraKLL
         )
     }
@@ -326,6 +331,7 @@ impl AggregationType {
                 | AggregationType::MultipleIncrease
                 | AggregationType::CountMinSketch
                 | AggregationType::CountMinSketchWithHeap
+                | AggregationType::CountSketch
         )
     }
 
@@ -360,9 +366,11 @@ impl FromStr for AggregationType {
             "HydraKLL" => Ok(AggregationType::HydraKLL),
             "CountMinSketch" => Ok(AggregationType::CountMinSketch),
             "CountMinSketchWithHeap" => Ok(AggregationType::CountMinSketchWithHeap),
+            "CountSketch" => Ok(AggregationType::CountSketch),
             "SetAggregator" => Ok(AggregationType::SetAggregator),
             "DeltaSetAggregator" => Ok(AggregationType::DeltaSetAggregator),
             "HLL" | "HyperLogLog" => Ok(AggregationType::HLL),
+            "DDSketch" | "DdSketch" => Ok(AggregationType::DDSketch),
             "SingleSubpopulation" => Ok(AggregationType::SingleSubpopulation),
             "MultipleSubpopulation" => Ok(AggregationType::MultipleSubpopulation),
             // Legacy accumulator-suffixed aliases
@@ -384,6 +392,9 @@ impl FromStr for AggregationType {
                 Ok(AggregationType::CountMinSketch)
             }
             "CountMinSketchWithHeapAccumulator" => Ok(AggregationType::CountMinSketchWithHeap),
+            "CountSketchAccumulator" | "CS" | "cs" | "count_sketch" => {
+                Ok(AggregationType::CountSketch)
+            }
             "SetAggregatorAccumulator" => Ok(AggregationType::SetAggregator),
             "DeltaSetAggregatorAccumulator" => Ok(AggregationType::DeltaSetAggregator),
             _ => Err(format!("Unknown aggregation type: '{s}'")),

asap-common/sketch-core/src/count_sketch.rs

@@ -0,0 +1,165 @@
+//! Count Sketch (a.k.a. Count-Min-style signed-counter sketch) —
+//! element-wise mergeable frequency estimator.
+//!
+//! Parallel to `count_min::CountMinSketch` but with **signed** counters,
+//! matching the `asap_sketchlib::proto::sketchlib::CountSketchState` wire
+//! format that DataCollector's `countsketchprocessor` emits via the
+//! modified OTLP `Metric.data = CountSketch{…}` variant.
+//!
+//! This is the minimal surface needed for PR C-CountSketch in the
+//! modified-OTLP hot path: construct from a decoded proto state, merge
+//! element-wise with another sketch, emit the matrix for queries and
+//! serialization. The richer query semantics of Count Sketch (median-
+//! of-estimators heavy-hitter tracking, `TopKState` integration, etc.)
+//! are intentionally deferred to a follow-up — the wire format already
+//! carries the matrix losslessly, so the merge/store round-trip works
+//! with just a matrix today.
+
+use serde::{Deserialize, Serialize};
+
+/// Minimal Count Sketch state — a flat `rows × cols` matrix of signed
+/// counts. Element-wise mergeable (sum over aligned cells).
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CountSketch {
+    pub row_num: usize,
+    pub col_num: usize,
+    /// Row-major matrix of signed counts. `matrix[r][c]` is the value of
+    /// hash row `r`, column `c`.
+    pub matrix: Vec<Vec<f64>>,
+}
+
+impl CountSketch {
+    /// Construct an all-zero sketch with the given dimensions.
+    pub fn new(row_num: usize, col_num: usize) -> Self {
+        Self {
+            row_num,
+            col_num,
+            matrix: vec![vec![0.0; col_num]; row_num],
+        }
+    }
+
+    /// Construct from a pre-built matrix (used by the modified-OTLP
+    /// proto-decode path).
+    pub fn from_legacy_matrix(matrix: Vec<Vec<f64>>, row_num: usize, col_num: usize) -> Self {
+        debug_assert_eq!(matrix.len(), row_num, "row count mismatch");
+        debug_assert!(
+            matrix.iter().all(|r| r.len() == col_num),
+            "column count mismatch in at least one row"
+        );
+        Self {
+            row_num,
+            col_num,
+            matrix,
+        }
+    }
+
+    /// Borrow the inner matrix.
+    pub fn sketch(&self) -> &Vec<Vec<f64>> {
+        &self.matrix
+    }
+
+    /// Merge one other sketch into self via element-wise addition. Both
+    /// operands must have identical dimensions.
+    pub fn merge(
+        &mut self,
+        other: &CountSketch,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+        if self.row_num != other.row_num || self.col_num != other.col_num {
+            return Err(format!(
+                "CountSketch dimension mismatch: self={}x{}, other={}x{}",
+                self.row_num, self.col_num, other.row_num, other.col_num
+            )
+            .into());
+        }
+        for r in 0..self.row_num {
+            for c in 0..self.col_num {
+                self.matrix[r][c] += other.matrix[r][c];
+            }
+        }
+        Ok(())
+    }
+
+    /// Merge a slice of references into a single new sketch. All inputs
+    /// must share the same dimensions; returns `Err` on mismatch or an
+    /// empty input.
+    pub fn merge_refs(
+        inputs: &[&CountSketch],
+    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
+        let first = inputs
+            .first()
+            .ok_or("CountSketch::merge_refs called with empty input")?;
+        let mut merged = CountSketch::new(first.row_num, first.col_num);
+        for cs in inputs {
+            merged.merge(cs)?;
+        }
+        Ok(merged)
+    }
+
+    /// Serialize to MessagePack bytes (used by the legacy Arroyo path
+    /// and by PR I's `_ENCODING_MSGPACK` variant when that lands).
+    pub fn serialize_msgpack(&self) -> Vec<u8> {
+        rmp_serde::to_vec(self).unwrap_or_default()
+    }
+
+    /// Deserialize from MessagePack bytes.
+    pub fn deserialize_msgpack(
+        buffer: &[u8],
+    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
+        Ok(rmp_serde::from_slice(buffer)?)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_new_empty() {
+        let cs = CountSketch::new(2, 3);
+        assert_eq!(cs.row_num, 2);
+        assert_eq!(cs.col_num, 3);
+        assert_eq!(cs.sketch(), &vec![vec![0.0, 0.0, 0.0], vec![0.0, 0.0, 0.0]]);
+    }
+
+    #[test]
+    fn test_from_legacy_matrix() {
+        let m = vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]];
+        let cs = CountSketch::from_legacy_matrix(m.clone(), 2, 3);
+        assert_eq!(cs.sketch(), &m);
+    }
+
+    #[test]
+    fn test_merge_element_wise() {
+        let mut a = CountSketch::from_legacy_matrix(vec![vec![1.0, 2.0], vec![3.0, 4.0]], 2, 2);
+        let b = CountSketch::from_legacy_matrix(vec![vec![-1.0, -2.0], vec![-3.0, -4.0]], 2, 2);
+        a.merge(&b).unwrap();
+        assert_eq!(a.sketch(), &vec![vec![0.0, 0.0], vec![0.0, 0.0]]);
+    }
+
+    #[test]
+    fn test_merge_dimension_mismatch() {
+        let mut a = CountSketch::new(2, 3);
+        let b = CountSketch::new(3, 3);
+        assert!(a.merge(&b).is_err());
+    }
+
+    #[test]
+    fn test_merge_refs() {
+        let a = CountSketch::from_legacy_matrix(vec![vec![1.0, 2.0]], 1, 2);
+        let b = CountSketch::from_legacy_matrix(vec![vec![3.0, 4.0]], 1, 2);
+        let c = CountSketch::from_legacy_matrix(vec![vec![5.0, 6.0]], 1, 2);
+        let merged = CountSketch::merge_refs(&[&a, &b, &c]).unwrap();
+        assert_eq!(merged.sketch(), &vec![vec![9.0, 12.0]]);
+    }
+
+    #[test]
+    fn test_msgpack_round_trip() {
+        let original =
+            CountSketch::from_legacy_matrix(vec![vec![1.5, -2.5], vec![3.5, -4.5]], 2, 2);
+        let bytes = original.serialize_msgpack();
+        let decoded = CountSketch::deserialize_msgpack(&bytes).unwrap();
+        assert_eq!(decoded.sketch(), original.sketch());
+        assert_eq!(decoded.row_num, original.row_num);
+        assert_eq!(decoded.col_num, original.col_num);
+    }
+}