Skip to main content

apache_datasketches/theta/
builder.rs

1use apache_datasketches_sys::theta_sketch::ffi as sys;
2
3/// Controls how aggressively a theta sketch's internal hash table grows.
4/// Mirrors upstream's `datasketches::resize_factor`. Default is `X8`,
5/// matching `theta_constants::DEFAULT_RESIZE_FACTOR`.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
7pub enum ResizeFactor {
8    /// Grow by 1x (i.e. never resize past the initial allocation).
9    X1,
10    /// Grow by 2x each time the hash table fills.
11    X2,
12    /// Grow by 4x each time the hash table fills.
13    X4,
14    /// Grow by 8x each time the hash table fills. The default.
15    #[default]
16    X8,
17}
18
19impl From<ResizeFactor> for sys::ResizeFactor {
20    fn from(rf: ResizeFactor) -> Self {
21        match rf {
22            ResizeFactor::X1 => sys::ResizeFactor::X1,
23            ResizeFactor::X2 => sys::ResizeFactor::X2,
24            ResizeFactor::X4 => sys::ResizeFactor::X4,
25            ResizeFactor::X8 => sys::ResizeFactor::X8,
26        }
27    }
28}
29
30/// Builder for [`crate::theta::ThetaSketch`], mirroring upstream's
31/// `update_theta_sketch::builder`. `lg_k` defaults to `12`
32/// (`theta_constants::DEFAULT_LG_K`), `resize_factor` to [`ResizeFactor::X8`],
33/// `p` to `1.0` (no sampling). The seed is never exposed — every sketch built
34/// by this crate always uses upstream's `DEFAULT_SEED`.
35#[derive(Debug, Clone, Copy)]
36pub struct ThetaSketchBuilder {
37    lg_k: u8,
38    resize_factor: ResizeFactor,
39    p: f32,
40}
41
42impl Default for ThetaSketchBuilder {
43    fn default() -> Self {
44        Self {
45            lg_k: 12,
46            resize_factor: ResizeFactor::default(),
47            p: 1.0,
48        }
49    }
50}
51
52impl ThetaSketchBuilder {
53    /// Creates a new builder with default settings (`lg_k = 12`,
54    /// `resize_factor = X8`, `p = 1.0`).
55    pub fn new() -> Self {
56        Self::default()
57    }
58
59    /// Sets the base-2 log of the target number of retained entries.
60    pub fn lg_k(mut self, lg_k: u8) -> Self {
61        self.lg_k = lg_k;
62        self
63    }
64
65    /// Sets the hash table's growth [`ResizeFactor`].
66    pub fn resize_factor(mut self, resize_factor: ResizeFactor) -> Self {
67        self.resize_factor = resize_factor;
68        self
69    }
70
71    /// Sets the sampling probability. `1.0` (the default) disables
72    /// sampling; values below `1.0` put the sketch into estimation mode
73    /// from the start.
74    pub fn p(mut self, p: f32) -> Self {
75        self.p = p;
76        self
77    }
78
79    /// Builds the sketch. Returns
80    /// [`SketchError::InvalidConfig`](crate::SketchError::InvalidConfig) if
81    /// `lg_k` is out of range.
82    pub fn build(self) -> Result<super::ThetaSketch, crate::error::SketchError> {
83        super::ThetaSketch::from_parts(self.lg_k, self.resize_factor, self.p)
84    }
85}