Skip to main content

apache_datasketches/tuple/generic/
sketch.rs

1use super::builder::resize_factor_multiplier;
2use super::summary::{erase, TupleSummary};
3use crate::error::SketchError;
4use crate::tuple::ResizeFactor;
5use apache_datasketches_sys::tuple_generic::ffi as sys;
6use cxx::UniquePtr;
7use std::marker::PhantomData;
8
9/// A mutable, update-only Tuple sketch carrying a user-defined summary `S`
10/// per distinct key. Build one with
11/// [`TupleSketchBuilder`](super::TupleSketchBuilder).
12pub struct TupleSketch<S: TupleSummary> {
13    pub(crate) inner: UniquePtr<sys::TupleGenericSketchShim>,
14    pub(crate) _marker: PhantomData<fn() -> S>,
15}
16
17// Sound because `S: Send` is a supertrait of `TupleSummary` and the sys-crate
18// box is `Box<dyn RawSummaryOps + Send>`. Deliberately not `Sync`.
19unsafe impl<S: TupleSummary> Send for TupleSketch<S> {}
20
21impl<S: TupleSummary> TupleSketch<S> {
22    pub(crate) fn from_parts(lg_k: u8, rf: ResizeFactor, p: f32) -> Result<Self, SketchError> {
23        let inner = sys::new_tuple_generic_sketch(lg_k, resize_factor_multiplier(rf), p)
24            .map_err(|e| SketchError::InvalidConfig(e.what().to_string()))?;
25        Ok(Self {
26            inner,
27            _marker: PhantomData,
28        })
29    }
30
31    /// Adds a `u64` key. `S::create` runs first, entirely in Rust; the
32    /// resulting summary is combined into an existing entry or cloned into a
33    /// new one.
34    pub fn update_u64(&mut self, key: u64, value: &S::Update) {
35        let summary = erase(S::create(value));
36        self.inner.pin_mut().update_u64(key, &summary);
37    }
38
39    /// Adds an `i64` key. See [`Self::update_u64`].
40    pub fn update_i64(&mut self, key: i64, value: &S::Update) {
41        let summary = erase(S::create(value));
42        self.inner.pin_mut().update_i64(key, &summary);
43    }
44
45    /// Adds a `u32` key. See [`Self::update_u64`].
46    pub fn update_u32(&mut self, key: u32, value: &S::Update) {
47        let summary = erase(S::create(value));
48        self.inner.pin_mut().update_u32(key, &summary);
49    }
50
51    /// Adds an `i32` key. See [`Self::update_u64`].
52    pub fn update_i32(&mut self, key: i32, value: &S::Update) {
53        let summary = erase(S::create(value));
54        self.inner.pin_mut().update_i32(key, &summary);
55    }
56
57    /// Adds a `u16` key. See [`Self::update_u64`].
58    pub fn update_u16(&mut self, key: u16, value: &S::Update) {
59        let summary = erase(S::create(value));
60        self.inner.pin_mut().update_u16(key, &summary);
61    }
62
63    /// Adds an `i16` key. See [`Self::update_u64`].
64    pub fn update_i16(&mut self, key: i16, value: &S::Update) {
65        let summary = erase(S::create(value));
66        self.inner.pin_mut().update_i16(key, &summary);
67    }
68
69    /// Adds a `u8` key. See [`Self::update_u64`].
70    pub fn update_u8(&mut self, key: u8, value: &S::Update) {
71        let summary = erase(S::create(value));
72        self.inner.pin_mut().update_u8(key, &summary);
73    }
74
75    /// Adds an `i8` key. See [`Self::update_u64`].
76    pub fn update_i8(&mut self, key: i8, value: &S::Update) {
77        let summary = erase(S::create(value));
78        self.inner.pin_mut().update_i8(key, &summary);
79    }
80
81    /// Adds an `f64` key. See [`Self::update_u64`].
82    pub fn update_f64(&mut self, key: f64, value: &S::Update) {
83        let summary = erase(S::create(value));
84        self.inner.pin_mut().update_f64(key, &summary);
85    }
86
87    /// Adds a string key. See [`Self::update_u64`].
88    pub fn update_str(&mut self, key: &str, value: &S::Update) {
89        let summary = erase(S::create(value));
90        self.inner.pin_mut().update_str(key, &summary);
91    }
92
93    /// Adds an arbitrary byte-slice key. See [`Self::update_u64`].
94    pub fn update_bytes(&mut self, key: &[u8], value: &S::Update) {
95        let summary = erase(S::create(value));
96        self.inner.pin_mut().update_bytes(key, &summary);
97    }
98
99    /// Removes retained entries in excess of the nominal size `k`, lowering
100    /// theta to do so. Note this shifts [`Self::get_estimate`].
101    pub fn trim(&mut self) {
102        self.inner.pin_mut().trim();
103    }
104
105    /// Resets this sketch to its initial, empty state.
106    pub fn reset(&mut self) {
107        self.inner.pin_mut().reset();
108    }
109
110    /// Returns the current estimate of the number of distinct keys added.
111    pub fn get_estimate(&self) -> f64 {
112        self.inner.get_estimate()
113    }
114
115    /// Returns the lower bound of the confidence interval around
116    /// [`Self::get_estimate`], for `num_std_dev` of `1`, `2`, or `3`.
117    pub fn get_lower_bound(&self, num_std_dev: u8) -> Result<f64, SketchError> {
118        self.inner
119            .get_lower_bound(num_std_dev)
120            .map_err(|e| SketchError::InvalidConfig(e.what().to_string()))
121    }
122
123    /// Returns the upper bound of the confidence interval around
124    /// [`Self::get_estimate`]. See [`Self::get_lower_bound`].
125    pub fn get_upper_bound(&self, num_std_dev: u8) -> Result<f64, SketchError> {
126        self.inner
127            .get_upper_bound(num_std_dev)
128            .map_err(|e| SketchError::InvalidConfig(e.what().to_string()))
129    }
130
131    /// Returns `true` if no keys have been added.
132    pub fn is_empty(&self) -> bool {
133        self.inner.is_empty()
134    }
135
136    /// Returns `true` if the sketch has begun sampling.
137    pub fn is_estimation_mode(&self) -> bool {
138        self.inner.is_estimation_mode()
139    }
140
141    /// Returns `true` if retained entries are sorted by hash value.
142    pub fn is_ordered(&self) -> bool {
143        self.inner.is_ordered()
144    }
145
146    /// Returns the current theta threshold.
147    pub fn get_theta(&self) -> f64 {
148        self.inner.get_theta()
149    }
150
151    /// Returns the number of retained entries.
152    pub fn get_num_retained(&self) -> u32 {
153        self.inner.get_num_retained()
154    }
155
156    /// Produces an immutable
157    /// [`CompactTupleSketch`](super::CompactTupleSketch) snapshot. If
158    /// `ordered` is `true`, its entries are sorted by hash value.
159    pub fn compact(&self, ordered: bool) -> super::CompactTupleSketch<S> {
160        super::CompactTupleSketch::from_shim(self.inner.compact(ordered))
161    }
162}