Skip to main content

samkhya_core/
portable.rs

1//! Engine-neutral, decoded view of versioned sketch payloads.
2//!
3//! Puffin discovers and transports opaque bytes; this module is the shared
4//! handoff that lets every engine decode those bytes with identical rules.
5//! Unknown blob kinds remain available in the raw bundle and are ignored by
6//! typed projection, while malformed payloads for known kinds fail closed.
7
8use std::collections::{BTreeMap, HashSet};
9
10use crate::sketches::{EquiDepthHistogram, HllSketch, Sketch};
11use crate::{ColumnStats, Error, Result};
12
13/// One decompressed Puffin blob with its portability metadata intact.
14#[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    /// Construct a blob when snapshot metadata is unavailable.
27    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    /// Attach the snapshot identity carried by Puffin blob metadata.
39    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    /// Attach arbitrary Puffin blob properties.
46    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/// Decompressed sketch blobs associated with one table snapshot.
77#[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    /// Return every raw blob whose field list includes `field_id`.
98    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    /// Validate known single-column kinds and decode their payloads.
105    ///
106    /// Unknown kinds are intentionally skipped according to the Puffin reader
107    /// contract. Duplicate `(field, kind)` entries are rejected as ambiguous.
108    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    /// Decode the supported sketches for one Iceberg field id.
139    ///
140    /// `Ok(None)` means no supported HLL or equi-depth histogram is present.
141    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/// Typed adapter view of the supported sketches for one column.
177#[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    /// Canonical scalar statistics suitable for native planner adapters.
192    ///
193    /// The HLL contributes `distinct_count`. A histogram remains available
194    /// through [`Self::histogram`] but does not imply a full table row count.
195    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
208/// Whether `kind` has a typed, single-column projection in this release.
209pub 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}