fix(ingest): envelope-aware decoders for KLL / HLL / DDSketch / CountSketch

asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs

@@ -61,11 +61,34 @@ impl CountSketchAccumulator {
     /// `CountSketch::from_legacy_matrix` after reshaping the flat
     /// `counts_int` / `counts_float` field into a `Vec<Vec<f64>>`.
     pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
-        use asap_sketchlib::proto::sketchlib::{CountSketchState, CounterType};
+        use asap_sketchlib::proto::sketchlib::{
+            sketch_envelope, CountSketchState, CounterType, SketchEnvelope,
+        };
         use prost::Message;
 
-        let state = CountSketchState::decode(buffer)
-            .map_err(|e| format!("decode CountSketchState: {e}"))?;
+        // DataCollector's countsketchprocessor wraps the state in a
+        // `SketchEnvelope{count_sketch: CountSketchState}` via
+        // sketchlib-go's `SerializePortableFO` + `proto.Marshal`. Try
+        // decoding as envelope first, fall back to bare
+        // `CountSketchState` for callers (e.g. unit tests) that
+        // encode the state directly. Mirrors the PR #14 fix on
+        // `CountMinSketchAccumulator::from_sketchlib_proto_bytes`.
+        let state = match SketchEnvelope::decode(buffer) {
+            Ok(env) => match env.sketch_state {
+                Some(sketch_envelope::SketchState::CountSketch(st)) => st,
+                Some(other) => {
+                    return Err(format!(
+                        "SketchEnvelope contains non-CountSketch sketch: {:?}",
+                        std::mem::discriminant(&other)
+                    )
+                    .into());
+                }
+                None => CountSketchState::decode(buffer)
+                    .map_err(|e| format!("decode CountSketchState: {e}"))?,
+            },
+            Err(_) => CountSketchState::decode(buffer)
+                .map_err(|e| format!("decode CountSketchState: {e}"))?,
+        };
         let rows = state.rows as usize;
         let cols = state.cols as usize;
         if rows == 0 || cols == 0 {
@@ -241,6 +264,55 @@ mod tests {
         assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]);
     }
 
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrapped() {
+        // Mirrors what DataCollector's countsketchprocessor emits:
+        // the state wrapped in a `SketchEnvelope{count_sketch: ...}`
+        // via sketchlib-go's `SerializePortableFO` + `proto.Marshal`.
+        use asap_sketchlib::proto::sketchlib::{
+            sketch_envelope, CountSketchState, CounterType, SketchEnvelope,
+        };
+        use prost::Message;
+
+        let state = CountSketchState {
+            rows: 2,
+            cols: 3,
+            counter_type: CounterType::Int64 as i32,
+            counts_int: vec![1, -2, 3, -4, 5, -6],
+            counts_float: Vec::new(),
+            ..Default::default()
+        };
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::CountSketch(state)),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes)
+            .expect("envelope-wrapped decode should succeed");
+        let matrix = acc.inner.sketch();
+        assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]);
+        assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]);
+    }
+
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() {
+        // An envelope carrying a non-CountSketch sketch should be
+        // rejected with a clear error rather than silently producing
+        // garbage.
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope};
+        use prost::Message;
+
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes);
+        assert!(result.is_err(), "wrong-sketch envelope should error");
+    }
+
     #[test]
     fn test_from_sketchlib_proto_bytes_float64() {
         use asap_sketchlib::proto::sketchlib::CounterType;

asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs

@@ -107,10 +107,29 @@ impl DatasketchesKLLAccumulator {
     /// already have accepted from the source. Bit-identical
     /// reconstruction is tracked as a sketchlib upstream follow-up.
     pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
-        use asap_sketchlib::proto::sketchlib::KllState;
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope};
         use prost::Message;
 
-        let state = KllState::decode(buffer).map_err(|e| format!("decode KllState: {e}"))?;
+        // DataCollector's kllprocessor wraps the state in a
+        // `SketchEnvelope{kll: KllState}` via sketchlib-go's
+        // `SerializePortableFO` + `proto.Marshal`. Try envelope first,
+        // fall back to bare `KllState` for callers (e.g. unit tests)
+        // that encode the state directly. Mirrors the PR #14 fix on
+        // `CountMinSketchAccumulator::from_sketchlib_proto_bytes`.
+        let state = match SketchEnvelope::decode(buffer) {
+            Ok(env) => match env.sketch_state {
+                Some(sketch_envelope::SketchState::Kll(st)) => st,
+                Some(other) => {
+                    return Err(format!(
+                        "SketchEnvelope contains non-KLL sketch: {:?}",
+                        std::mem::discriminant(&other)
+                    )
+                    .into());
+                }
+                None => KllState::decode(buffer).map_err(|e| format!("decode KllState: {e}"))?,
+            },
+            Err(_) => KllState::decode(buffer).map_err(|e| format!("decode KllState: {e}"))?,
+        };
         if state.k < 8 {
             return Err(format!("KllState.k must be >= 8 (got {})", state.k).into());
         }
@@ -607,6 +626,51 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrapped() {
+        // Mirrors what DataCollector's kllprocessor emits: the state
+        // wrapped in a `SketchEnvelope{kll: ...}` via sketchlib-go's
+        // `SerializePortableFO` + `proto.Marshal`.
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope};
+        use prost::Message;
+
+        let items: Vec<f64> = (0..64).map(|i| i as f64).collect();
+        let state = KllState {
+            k: 200,
+            m: 8,
+            num_levels: 1,
+            levels: vec![0, 64],
+            items,
+            coin: None,
+        };
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::Kll(state)),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes)
+            .expect("envelope-wrapped decode should succeed");
+        assert_eq!(acc.inner.count(), 64);
+    }
+
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() {
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, CountMinState, SketchEnvelope};
+        use prost::Message;
+
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::CountMin(
+                CountMinState::default(),
+            )),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes);
+        assert!(result.is_err(), "wrong-sketch envelope should error");
+    }
+
     #[test]
     fn test_from_sketchlib_proto_bytes_rejects_small_k() {
         use asap_sketchlib::proto::sketchlib::KllState;

asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs

@@ -47,11 +47,32 @@ impl DDSketchAccumulator {
     /// DataCollector's `ddsketchprocessor` emits when
     /// `encoding = DD_SKETCH_ENCODING_PROTO`.
     pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
-        use asap_sketchlib::proto::sketchlib::DdSketchState;
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope};
         use prost::Message;
 
-        let state =
-            DdSketchState::decode(buffer).map_err(|e| format!("decode DDSketchState: {e}"))?;
+        // DataCollector's ddsketchprocessor wraps the state in a
+        // `SketchEnvelope{ddsketch: DdSketchState}` via sketchlib-go's
+        // `SerializePortableFO` + `proto.Marshal`. Try envelope first,
+        // fall back to bare `DdSketchState` for callers (e.g. unit
+        // tests) that encode the state directly. Mirrors the PR #14
+        // fix on `CountMinSketchAccumulator::from_sketchlib_proto_bytes`.
+        let state = match SketchEnvelope::decode(buffer) {
+            Ok(env) => match env.sketch_state {
+                Some(sketch_envelope::SketchState::Ddsketch(st)) => st,
+                Some(other) => {
+                    return Err(format!(
+                        "SketchEnvelope contains non-DDSketch sketch: {:?}",
+                        std::mem::discriminant(&other)
+                    )
+                    .into());
+                }
+                None => DdSketchState::decode(buffer)
+                    .map_err(|e| format!("decode DDSketchState: {e}"))?,
+            },
+            Err(_) => {
+                DdSketchState::decode(buffer).map_err(|e| format!("decode DDSketchState: {e}"))?
+            }
+        };
         if !(state.alpha > 0.0 && state.alpha < 1.0) {
             return Err(format!(
                 "DDSketchState alpha {} out of range (expected 0 < alpha < 1)",
@@ -193,6 +214,50 @@ mod tests {
         assert!(result.unwrap_err().to_string().contains("alpha"));
     }
 
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrapped() {
+        // Mirrors what DataCollector's ddsketchprocessor emits: the
+        // state wrapped in a `SketchEnvelope{ddsketch: ...}` via
+        // sketchlib-go's `SerializePortableFO` + `proto.Marshal`.
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope};
+        use prost::Message;
+
+        let state = DdSketchState {
+            alpha: 0.01,
+            store_counts: vec![1, 2, 3, 4],
+            store_offset: -2,
+            count: 10,
+            sum: 50.0,
+            min: 1.0,
+            max: 4.0,
+        };
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes)
+            .expect("envelope-wrapped decode should succeed");
+        assert_eq!(acc.inner.alpha, 0.01);
+        assert_eq!(acc.inner.count, 10);
+    }
+
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() {
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope};
+        use prost::Message;
+
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes);
+        assert!(result.is_err(), "wrong-sketch envelope should error");
+    }
+
     #[test]
     fn test_aggregate_core_merge_aligns_buckets() {
         let a = DDSketchAccumulator {

asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs

@@ -47,11 +47,33 @@ impl HllSketchAccumulator {
     /// that DataCollector's `hllprocessor` emits when
     /// `encoding = HLL_SKETCH_ENCODING_PROTO`.
     pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
-        use asap_sketchlib::proto::sketchlib::{HllVariant as ProtoVariant, HyperLogLogState};
+        use asap_sketchlib::proto::sketchlib::{
+            sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope,
+        };
         use prost::Message;
 
-        let state = HyperLogLogState::decode(buffer)
-            .map_err(|e| format!("decode HyperLogLogState: {e}"))?;
+        // DataCollector's hllprocessor wraps the state in a
+        // `SketchEnvelope{hll: HyperLogLogState}` via sketchlib-go's
+        // `SerializePortableFO` + `proto.Marshal`. Try envelope first,
+        // fall back to bare `HyperLogLogState` for callers (e.g. unit
+        // tests) that encode the state directly. Mirrors the PR #14
+        // fix on `CountMinSketchAccumulator::from_sketchlib_proto_bytes`.
+        let state = match SketchEnvelope::decode(buffer) {
+            Ok(env) => match env.sketch_state {
+                Some(sketch_envelope::SketchState::Hll(st)) => st,
+                Some(other) => {
+                    return Err(format!(
+                        "SketchEnvelope contains non-HLL sketch: {:?}",
+                        std::mem::discriminant(&other)
+                    )
+                    .into());
+                }
+                None => HyperLogLogState::decode(buffer)
+                    .map_err(|e| format!("decode HyperLogLogState: {e}"))?,
+            },
+            Err(_) => HyperLogLogState::decode(buffer)
+                .map_err(|e| format!("decode HyperLogLogState: {e}"))?,
+        };
         if state.precision == 0 || state.precision > 20 {
             return Err(format!(
                 "HyperLogLogState precision {} out of range (expected 1..=20)",
@@ -220,6 +242,51 @@ mod tests {
         assert_eq!(acc.inner.hip_est, 42.0);
     }
 
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrapped() {
+        // Mirrors what DataCollector's hllprocessor emits: the state
+        // wrapped in a `SketchEnvelope{hll: ...}` via sketchlib-go's
+        // `SerializePortableFO` + `proto.Marshal`.
+        use asap_sketchlib::proto::sketchlib::{
+            sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope,
+        };
+        use prost::Message;
+
+        let state = HyperLogLogState {
+            variant: ProtoVariant::Regular as i32,
+            precision: 2,
+            registers: vec![1, 2, 3, 4],
+            hip_kxq0: 0.0,
+            hip_kxq1: 0.0,
+            hip_est: 0.0,
+        };
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::Hll(state)),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes)
+            .expect("envelope-wrapped decode should succeed");
+        assert_eq!(acc.inner.variant, HllVariant::Regular);
+        assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]);
+    }
+
+    #[test]
+    fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() {
+        use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope};
+        use prost::Message;
+
+        let env = SketchEnvelope {
+            sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())),
+            ..Default::default()
+        };
+        let bytes = env.encode_to_vec();
+
+        let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes);
+        assert!(result.is_err(), "wrong-sketch envelope should error");
+    }
+
     #[test]
     fn test_from_sketchlib_proto_bytes_register_length_mismatch() {
         use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant;