From 4d0865888494983438ff5e78ee6d422922097cea Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Fri, 17 Apr 2026 17:42:43 -0400 Subject: [PATCH] fix(ingest): envelope-aware decoders for KLL / HLL / DDSketch / CountSketch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #14 fixed `CountMinSketchAccumulator::from_sketchlib_proto_bytes` to try `SketchEnvelope::decode` first and fall back to bare `CountMinState`. DataCollector's sketchlib-go processors wrap every sketch state in a `SketchEnvelope` via `SerializePortableFO` + `proto.Marshal`, so the bare-proto decoder was producing "invalid wire type" errors and silently falling through to ยง5.2 for every real-world sketch data point. The same bug existed for the other four sketch accumulators. This PR applies the envelope-first-then-bare-fallback pattern verbatim to: * `CountSketchAccumulator::from_sketchlib_proto_bytes` * `DatasketchesKLLAccumulator::from_sketchlib_proto_bytes` * `HllSketchAccumulator::from_sketchlib_proto_bytes` * `DDSketchAccumulator::from_sketchlib_proto_bytes` On envelope decode success, dispatch on the oneof variant and extract the expected state. On non-matching variant, return a clear `SketchEnvelope contains non- sketch` error instead of silently producing garbage. On envelope decode failure or empty oneof, fall back to bare proto decode so unit tests (which encode states directly without the envelope) keep working. Tested end-to-end via 8 new unit tests: two per accumulator (envelope-wrapped happy path + wrong-sketch-type rejection). All 10 envelope tests pass (8 new + 2 pre-existing CMS), 578 lib tests total (up from 570); clippy + fmt clean. Before the fix, wiring any of these four processors in a real deployment would fail with a decode error on every data point. The CMS fix shipped in PR #14; this PR completes the pattern across the remaining sketch types. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../count_sketch_accumulator.rs | 78 ++++++++++++++++++- .../datasketches_kll_accumulator.rs | 68 +++++++++++++++- .../dd_sketch_accumulator.rs | 71 ++++++++++++++++- .../hll_sketch_accumulator.rs | 73 ++++++++++++++++- 4 files changed, 279 insertions(+), 11 deletions(-) diff --git a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs index 7e5d0896..93ad50cb 100644 --- a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs +++ b/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>`. pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { - 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; diff --git a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs index 086ae876..009383d8 100644 --- a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs +++ b/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> { - 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 = (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; diff --git a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs index 48f0dc36..8ebd1351 100644 --- a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs +++ b/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> { - 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 { diff --git a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs index f1eff42f..1e271c17 100644 --- a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs +++ b/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> { - 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;