apache_datasketches/tuple/
union.rs

1use super::input::ArrayOfDoublesInput;
2use super::{CompactArrayOfDoublesSketch, ResizeFactor};
3use crate::error::SketchError;
4use apache_datasketches_sys::array_of_doubles_input::ArrayOfDoublesInputRef;
5use apache_datasketches_sys::array_of_doubles_union::ffi as sys;
6use cxx::UniquePtr;
7
8/// Builder for [`ArrayOfDoublesUnion`], mirroring upstream's
9/// `array_of_doubles_union::builder`. `lg_k` defaults to `12`,
10/// `resize_factor` to [`ResizeFactor::X8`], `p` to `1.0` (no sampling), and
11/// `num_values` to `1`. As with
12/// [`ArrayOfDoublesSketchBuilder`](super::ArrayOfDoublesSketchBuilder), the
13/// seed is never exposed.
14#[derive(Debug, Clone, Copy)]
15pub struct ArrayOfDoublesUnionBuilder {
16    lg_k: u8,
17    resize_factor: ResizeFactor,
18    p: f32,
19    num_values: u8,
20}
21
22impl Default for ArrayOfDoublesUnionBuilder {
23    fn default() -> Self {
24        Self {
25            lg_k: 12,
26            resize_factor: ResizeFactor::default(),
27            p: 1.0,
28            num_values: 1,
29        }
30    }
31}
32
33impl ArrayOfDoublesUnionBuilder {
34    /// Creates a new builder with default settings (`lg_k = 12`,
35    /// `resize_factor = X8`, `p = 1.0`, `num_values = 1`).
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Sets the base-2 log of the target number of retained entries in the
41    /// union's result.
42    pub fn lg_k(mut self, lg_k: u8) -> Self {
43        self.lg_k = lg_k;
44        self
45    }
46
47    /// Sets the hash table's growth [`ResizeFactor`].
48    pub fn resize_factor(mut self, resize_factor: ResizeFactor) -> Self {
49        self.resize_factor = resize_factor;
50        self
51    }
52
53    /// Sets the sampling probability. `1.0` (the default) disables sampling.
54    pub fn p(mut self, p: f32) -> Self {
55        self.p = p;
56        self
57    }
58
59    /// Sets the fixed number of `f64` values per entry. Must be at least `1`,
60    /// and must match every sketch later passed to
61    /// [`ArrayOfDoublesUnion::update`].
62    pub fn num_values(mut self, num_values: u8) -> Self {
63        self.num_values = num_values;
64        self
65    }
66
67    /// Builds the union. Returns [`SketchError::InvalidConfig`] if `lg_k` is
68    /// out of range, `p` is outside `(0, 1]`, or `num_values` is `0`.
69    pub fn build(self) -> Result<ArrayOfDoublesUnion, SketchError> {
70        if self.num_values == 0 {
71            return Err(SketchError::InvalidConfig(
72                "num_values must be at least 1".to_string(),
73            ));
74        }
75        let inner = sys::new_array_of_doubles_union(
76            self.lg_k,
77            self.resize_factor.into(),
78            self.p,
79            self.num_values,
80        )
81        .map_err(|e| SketchError::InvalidConfig(e.what().to_string()))?;
82        Ok(ArrayOfDoublesUnion {
83            inner,
84            num_values: self.num_values,
85        })
86    }
87}
88
89/// A streaming union accumulator over ArrayOfDoubles sketches. Values are
90/// summed per index when the same key appears in more than one input, using
91/// upstream's `default_array_of_doubles_union_policy`.
92///
93/// Accepts either [`super::ArrayOfDoublesSketch`] or
94/// [`CompactArrayOfDoublesSketch`] via the sealed [`ArrayOfDoublesInput`]
95/// trait.
96pub struct ArrayOfDoublesUnion {
97    inner: UniquePtr<sys::ArrayOfDoublesUnionShim>,
98    num_values: u8,
99}
100
101unsafe impl Send for ArrayOfDoublesUnion {}
102
103impl ArrayOfDoublesUnion {
104    /// Merges the given sketch into this union's running result.
105    ///
106    /// Returns [`SketchError::InvalidConfig`] if the sketch's `num_values`
107    /// differs from this union's. Upstream does not validate this itself —
108    /// merging mismatched array widths would read and write past the shorter
109    /// array's bounds rather than error — so the check happens here, before
110    /// the sketch crosses the FFI boundary.
111    pub fn update(&mut self, input: &impl ArrayOfDoublesInput) -> Result<(), SketchError> {
112        let actual = input.get_num_values();
113        if actual != self.num_values {
114            return Err(SketchError::InvalidConfig(format!(
115                "num_values mismatch: union has {}, input has {actual}",
116                self.num_values
117            )));
118        }
119        match input.as_input() {
120            ArrayOfDoublesInputRef::Sketch(s) => self.inner.pin_mut().update_with_sketch(s),
121            ArrayOfDoublesInputRef::Compact(c) => self.inner.pin_mut().update_with_compact(c),
122        }
123        Ok(())
124    }
125
126    /// Returns the union's current result as a
127    /// [`CompactArrayOfDoublesSketch`]. If `ordered` is `true`, the result's
128    /// entries are sorted by hash value.
129    pub fn get_result(&self, ordered: bool) -> CompactArrayOfDoublesSketch {
130        CompactArrayOfDoublesSketch::from_shim(self.inner.get_result(ordered))
131    }
132
133    /// Resets this union to its initial, empty state. `num_values` is
134    /// preserved.
135    pub fn reset(&mut self) {
136        self.inner.pin_mut().reset();
137    }
138
139    /// Returns the fixed number of `f64` values per entry this union was
140    /// built with. Every input passed to [`Self::update`] must match it.
141    pub fn get_num_values(&self) -> u8 {
142        self.num_values
143    }
144}
apache_datasketches/tuple/union.rs

apache_datasketches/tuple/
union.rs
