apache_datasketches/theta/
union.rs

1use super::input::ThetaInput;
2use super::{CompactThetaSketch, ResizeFactor};
3use crate::error::SketchError;
4use apache_datasketches_sys::theta_input::ThetaInputRef;
5use apache_datasketches_sys::theta_union::ffi as sys;
6use cxx::UniquePtr;
7
8/// Builder for [`ThetaUnion`], mirroring upstream's `theta_union::builder`.
9/// `lg_k` defaults to `12`, `resize_factor` to [`ResizeFactor::X8`], `p` to
10/// `1.0` (no sampling). As with [`super::ThetaSketchBuilder`], the seed is
11/// never exposed.
12pub struct ThetaUnionBuilder {
13    lg_k: u8,
14    resize_factor: ResizeFactor,
15    p: f32,
16}
17
18impl Default for ThetaUnionBuilder {
19    fn default() -> Self {
20        Self {
21            lg_k: 12,
22            resize_factor: ResizeFactor::default(),
23            p: 1.0,
24        }
25    }
26}
27
28impl ThetaUnionBuilder {
29    /// Creates a new builder with default settings (`lg_k = 12`,
30    /// `resize_factor = X8`, `p = 1.0`).
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Sets the base-2 log of the target number of retained entries in the
36    /// union's result.
37    pub fn lg_k(mut self, lg_k: u8) -> Self {
38        self.lg_k = lg_k;
39        self
40    }
41
42    /// Sets the hash table's growth [`ResizeFactor`].
43    pub fn resize_factor(mut self, resize_factor: ResizeFactor) -> Self {
44        self.resize_factor = resize_factor;
45        self
46    }
47
48    /// Sets the sampling probability. `1.0` (the default) disables
49    /// sampling.
50    pub fn p(mut self, p: f32) -> Self {
51        self.p = p;
52        self
53    }
54
55    /// Builds the union. Returns [`SketchError::InvalidConfig`] if `lg_k`
56    /// is out of range.
57    pub fn build(self) -> Result<ThetaUnion, SketchError> {
58        let inner = sys::new_theta_union(self.lg_k, self.resize_factor.into(), self.p)
59            .map_err(|e| SketchError::InvalidConfig(e.what().to_string()))?;
60        Ok(ThetaUnion { inner })
61    }
62}
63
64/// A streaming union accumulator over theta sketches. Accepts any of
65/// [`super::ThetaSketch`], [`CompactThetaSketch`], or
66/// [`super::WrappedCompactThetaSketch`] via the sealed [`ThetaInput`] trait.
67pub struct ThetaUnion {
68    inner: UniquePtr<sys::ThetaUnionShim>,
69}
70
71unsafe impl Send for ThetaUnion {}
72
73impl ThetaUnion {
74    /// Merges the given sketch into this union's running result.
75    pub fn update(&mut self, input: &impl ThetaInput) {
76        match input.as_theta_input() {
77            ThetaInputRef::Sketch(s) => self.inner.pin_mut().update_with_sketch(s),
78            ThetaInputRef::Compact(c) => self.inner.pin_mut().update_with_compact(c),
79            ThetaInputRef::Wrapped(w) => self.inner.pin_mut().update_with_wrapped(w),
80        }
81    }
82
83    /// Returns the union's current result as a
84    /// [`CompactThetaSketch`]. If `ordered` is `true`, the result's
85    /// entries are sorted by hash value.
86    pub fn get_result(&self, ordered: bool) -> CompactThetaSketch {
87        CompactThetaSketch::from_shim(self.inner.get_result(ordered))
88    }
89
90    /// Resets this union to its initial, empty state.
91    pub fn reset(&mut self) {
92        self.inner.pin_mut().reset();
93    }
94}
apache_datasketches/theta/union.rs

apache_datasketches/theta/
union.rs
