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;