1use std::collections::{BTreeMap, HashSet};
9
10use crate::sketches::{EquiDepthHistogram, HllSketch, Sketch};
11use crate::{ColumnStats, Error, Result};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
15#[non_exhaustive]
16pub struct PortableSketchBlob {
17 kind: String,
18 fields: Vec<i32>,
19 snapshot_id: Option<i64>,
20 sequence_number: Option<i64>,
21 payload: Vec<u8>,
22 properties: BTreeMap<String, String>,
23}
24
25impl PortableSketchBlob {
26 pub fn new(kind: impl Into<String>, fields: Vec<i32>, payload: Vec<u8>) -> Self {
28 Self {
29 kind: kind.into(),
30 fields,
31 snapshot_id: None,
32 sequence_number: None,
33 payload,
34 properties: BTreeMap::new(),
35 }
36 }
37
38 pub fn with_snapshot_metadata(mut self, snapshot_id: i64, sequence_number: i64) -> Self {
40 self.snapshot_id = Some(snapshot_id);
41 self.sequence_number = Some(sequence_number);
42 self
43 }
44
45 pub fn with_properties(mut self, properties: BTreeMap<String, String>) -> Self {
47 self.properties = properties;
48 self
49 }
50
51 pub fn kind(&self) -> &str {
52 &self.kind
53 }
54
55 pub fn fields(&self) -> &[i32] {
56 &self.fields
57 }
58
59 pub const fn snapshot_id(&self) -> Option<i64> {
60 self.snapshot_id
61 }
62
63 pub const fn sequence_number(&self) -> Option<i64> {
64 self.sequence_number
65 }
66
67 pub fn payload(&self) -> &[u8] {
68 &self.payload
69 }
70
71 pub fn properties(&self) -> &BTreeMap<String, String> {
72 &self.properties
73 }
74}
75
76#[derive(Debug, Clone, Default, PartialEq, Eq)]
78#[non_exhaustive]
79pub struct PortableStatsSnapshot {
80 snapshot_id: Option<i64>,
81 blobs: Vec<PortableSketchBlob>,
82}
83
84impl PortableStatsSnapshot {
85 pub fn new(snapshot_id: Option<i64>, blobs: Vec<PortableSketchBlob>) -> Self {
86 Self { snapshot_id, blobs }
87 }
88
89 pub const fn snapshot_id(&self) -> Option<i64> {
90 self.snapshot_id
91 }
92
93 pub fn blobs(&self) -> &[PortableSketchBlob] {
94 &self.blobs
95 }
96
97 pub fn blobs_for_field(&self, field_id: i32) -> impl Iterator<Item = &PortableSketchBlob> {
99 self.blobs
100 .iter()
101 .filter(move |blob| blob.fields.contains(&field_id))
102 }
103
104 pub fn validate(&self) -> Result<()> {
109 let mut seen = HashSet::new();
110 for blob in &self.blobs {
111 if !is_supported_column_kind(blob.kind()) {
112 continue;
113 }
114 let [field_id] = blob.fields() else {
115 return Err(Error::InvalidPuffin(format!(
116 "{} requires exactly one field id; got {:?}",
117 blob.kind(),
118 blob.fields()
119 )));
120 };
121 if *field_id < 0 {
122 return Err(Error::InvalidPuffin(format!(
123 "{} carries negative field id {field_id}",
124 blob.kind()
125 )));
126 }
127 if !seen.insert((*field_id, blob.kind().to_owned())) {
128 return Err(Error::InvalidPuffin(format!(
129 "duplicate {} blob for field {field_id}",
130 blob.kind()
131 )));
132 }
133 decode_known(blob)?;
134 }
135 Ok(())
136 }
137
138 pub fn decode_column(&self, field_id: i32) -> Result<Option<DecodedColumnStats>> {
142 if field_id < 0 {
143 return Err(Error::InvalidPuffin(format!(
144 "field id must be non-negative; got {field_id}"
145 )));
146 }
147 self.validate()?;
148
149 let mut hll = None;
150 let mut histogram = None;
151 for blob in self.blobs.iter().filter(|blob| blob.fields() == [field_id]) {
152 match blob.kind() {
153 HllSketch::KIND => hll = Some(HllSketch::from_bytes(blob.payload())?),
154 EquiDepthHistogram::KIND => {
155 histogram = Some(EquiDepthHistogram::from_bytes(blob.payload())?)
156 }
157 _ => {}
158 }
159 }
160
161 if hll.is_none() && histogram.is_none() {
162 return Ok(None);
163 }
164 let stats = hll.as_ref().map_or_else(ColumnStats::new, |sketch| {
165 ColumnStats::new().with_distinct_count(sketch.estimate())
166 });
167 Ok(Some(DecodedColumnStats {
168 field_id,
169 stats,
170 hll,
171 histogram,
172 }))
173 }
174}
175
176#[derive(Debug, Clone)]
178#[non_exhaustive]
179pub struct DecodedColumnStats {
180 field_id: i32,
181 stats: ColumnStats,
182 hll: Option<HllSketch>,
183 histogram: Option<EquiDepthHistogram>,
184}
185
186impl DecodedColumnStats {
187 pub const fn field_id(&self) -> i32 {
188 self.field_id
189 }
190
191 pub const fn column_stats(&self) -> &ColumnStats {
196 &self.stats
197 }
198
199 pub const fn hll(&self) -> Option<&HllSketch> {
200 self.hll.as_ref()
201 }
202
203 pub const fn histogram(&self) -> Option<&EquiDepthHistogram> {
204 self.histogram.as_ref()
205 }
206}
207
208pub fn is_supported_column_kind(kind: &str) -> bool {
210 matches!(kind, HllSketch::KIND | EquiDepthHistogram::KIND)
211}
212
213fn decode_known(blob: &PortableSketchBlob) -> Result<()> {
214 match blob.kind() {
215 HllSketch::KIND => HllSketch::from_bytes(blob.payload()).map(|_| ()),
216 EquiDepthHistogram::KIND => EquiDepthHistogram::from_bytes(blob.payload()).map(|_| ()),
217 _ => Ok(()),
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 fn fixture() -> (PortableStatsSnapshot, Vec<u8>, Vec<u8>) {
226 let mut hll = HllSketch::new(10).unwrap();
227 for value in 0..256_u32 {
228 hll.add(&value.to_le_bytes());
229 }
230 let histogram = EquiDepthHistogram::from_values(
231 &(0..100).map(|value| value as f64).collect::<Vec<_>>(),
232 10,
233 )
234 .unwrap();
235 let hll_bytes = hll.to_bytes().unwrap();
236 let histogram_bytes = histogram.to_bytes().unwrap();
237 let snapshot = PortableStatsSnapshot::new(
238 Some(42),
239 vec![
240 PortableSketchBlob::new(HllSketch::KIND, vec![17], hll_bytes.clone()),
241 PortableSketchBlob::new(
242 EquiDepthHistogram::KIND,
243 vec![17],
244 histogram_bytes.clone(),
245 ),
246 PortableSketchBlob::new("vendor.future-v1", vec![17], b"opaque".to_vec()),
247 ],
248 );
249 (snapshot, hll_bytes, histogram_bytes)
250 }
251
252 #[test]
253 fn decodes_supported_kinds_and_preserves_raw_payloads() {
254 let (snapshot, hll_bytes, histogram_bytes) = fixture();
255 snapshot.validate().unwrap();
256 let decoded = snapshot.decode_column(17).unwrap().unwrap();
257
258 assert_eq!(decoded.field_id(), 17);
259 assert!(decoded.column_stats().distinct_count.is_some());
260 assert_eq!(decoded.histogram().unwrap().total(), 100);
261 assert_eq!(decoded.histogram().unwrap().estimate_range(0.0, 49.0), 50);
262 assert_eq!(
263 snapshot
264 .blobs()
265 .iter()
266 .find(|blob| blob.kind() == HllSketch::KIND)
267 .unwrap()
268 .payload(),
269 hll_bytes
270 );
271 assert_eq!(
272 snapshot
273 .blobs()
274 .iter()
275 .find(|blob| blob.kind() == EquiDepthHistogram::KIND)
276 .unwrap()
277 .payload(),
278 histogram_bytes
279 );
280 }
281
282 #[test]
283 fn unknown_kind_is_preserved_but_not_projected() {
284 let snapshot = PortableStatsSnapshot::new(
285 None,
286 vec![PortableSketchBlob::new(
287 "samkhya.hll-v2",
288 vec![17],
289 b"future".to_vec(),
290 )],
291 );
292 assert!(snapshot.decode_column(17).unwrap().is_none());
293 assert_eq!(snapshot.blobs()[0].payload(), b"future");
294 }
295
296 #[test]
297 fn corrupt_known_payload_fails_closed() {
298 let snapshot = PortableStatsSnapshot::new(
299 None,
300 vec![PortableSketchBlob::new(
301 HllSketch::KIND,
302 vec![17],
303 b"corrupt".to_vec(),
304 )],
305 );
306 assert!(snapshot.validate().is_err());
307 assert!(snapshot.decode_column(17).is_err());
308 }
309
310 #[test]
311 fn duplicate_known_kind_is_rejected() {
312 let (mut snapshot, _, _) = fixture();
313 let duplicate = snapshot
314 .blobs()
315 .iter()
316 .find(|blob| blob.kind() == HllSketch::KIND)
317 .unwrap()
318 .clone();
319 snapshot.blobs.push(duplicate);
320 assert!(snapshot.validate().is_err());
321 }
322}