Skip to main content

apache_datasketches/tuple/generic/
union.rs

1use super::builder::resize_factor_multiplier;
2use super::{CompactTupleSketch, TupleInput, TupleSummary};
3use crate::error::SketchError;
4use crate::tuple::ResizeFactor;
5use apache_datasketches_sys::tuple_generic_input::TupleGenericInputRef;
6use apache_datasketches_sys::tuple_generic_union::ffi as sys;
7use cxx::UniquePtr;
8use std::marker::PhantomData;
9
10/// Builder for [`TupleUnion`]. `lg_k` defaults to `12`, `resize_factor` to
11/// [`ResizeFactor::X8`], `p` to `1.0`.
12pub struct TupleUnionBuilder<S: TupleSummary> {
13    lg_k: u8,
14    resize_factor: ResizeFactor,
15    p: f32,
16    _marker: PhantomData<fn() -> S>,
17}
18
19// Hand-written rather than `#[derive(..)]`, for the same reason as
20// `TupleSketchBuilder` (see the note in builder.rs): a derive would add an
21// `S: Debug`/`S: Clone`/`S: Copy` bound to each impl even though every field
22// here is unconditionally `Debug + Clone + Copy`. `TupleSummary` requires
23// neither `Debug` nor `Copy`, so deriving would make this builder silently
24// non-`Debug` and non-`Copy` for most summaries -- and inconsistent with its
25// sketch-side counterpart.
26impl<S: TupleSummary> std::fmt::Debug for TupleUnionBuilder<S> {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("TupleUnionBuilder")
29            .field("lg_k", &self.lg_k)
30            .field("resize_factor", &self.resize_factor)
31            .field("p", &self.p)
32            .finish()
33    }
34}
35
36impl<S: TupleSummary> Clone for TupleUnionBuilder<S> {
37    fn clone(&self) -> Self {
38        *self
39    }
40}
41
42impl<S: TupleSummary> Copy for TupleUnionBuilder<S> {}
43
44impl<S: TupleSummary> Default for TupleUnionBuilder<S> {
45    fn default() -> Self {
46        Self {
47            lg_k: 12,
48            resize_factor: ResizeFactor::X8,
49            p: 1.0,
50            _marker: PhantomData,
51        }
52    }
53}
54
55impl<S: TupleSummary> TupleUnionBuilder<S> {
56    /// Creates a builder with default settings.
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    /// Sets the base-2 log of the target number of retained entries.
62    pub fn lg_k(mut self, lg_k: u8) -> Self {
63        self.lg_k = lg_k;
64        self
65    }
66
67    /// Sets the hash table's growth [`ResizeFactor`].
68    pub fn resize_factor(mut self, resize_factor: ResizeFactor) -> Self {
69        self.resize_factor = resize_factor;
70        self
71    }
72
73    /// Sets the sampling probability.
74    pub fn p(mut self, p: f32) -> Self {
75        self.p = p;
76        self
77    }
78
79    /// Builds the union. Returns [`SketchError::InvalidConfig`] if `lg_k` is
80    /// out of range or `p` is outside `(0, 1]`.
81    pub fn build(self) -> Result<TupleUnion<S>, SketchError> {
82        let inner = sys::new_tuple_generic_union(
83            self.lg_k,
84            resize_factor_multiplier(self.resize_factor),
85            self.p,
86        )
87        .map_err(|e| SketchError::InvalidConfig(e.what().to_string()))?;
88        Ok(TupleUnion {
89            inner,
90            _marker: PhantomData,
91        })
92    }
93}
94
95/// A streaming union over generic Tuple sketches. Summaries for a key present
96/// in more than one input are merged with
97/// [`TupleSummary::union_combine`](super::TupleSummary::union_combine).
98pub struct TupleUnion<S: TupleSummary> {
99    inner: UniquePtr<sys::TupleGenericUnionShim>,
100    _marker: PhantomData<fn() -> S>,
101}
102
103// Sound because `S: Send` is a supertrait of `TupleSummary` and every summary
104// the union owns is a `Box<dyn RawSummaryOps + Send>`. Deliberately not
105// `Sync`, matching `TupleSketch<S>` and `CompactTupleSketch<S>`.
106unsafe impl<S: TupleSummary> Send for TupleUnion<S> {}
107
108impl<S: TupleSummary> TupleUnion<S> {
109    /// Merges the given sketch into the running result.
110    ///
111    /// Infallible: unlike ArrayOfDoubles there is no `num_values` to agree
112    /// on — the type system already guarantees both operands carry `S`.
113    pub fn update(&mut self, input: &impl TupleInput<S>) {
114        match input.as_input() {
115            TupleGenericInputRef::Sketch(s) => self.inner.pin_mut().update_with_sketch(s),
116            TupleGenericInputRef::Compact(c) => self.inner.pin_mut().update_with_compact(c),
117        }
118    }
119
120    /// Returns the union's current result. If `ordered` is `true`, entries
121    /// are sorted by hash value.
122    pub fn get_result(&self, ordered: bool) -> CompactTupleSketch<S> {
123        CompactTupleSketch::from_shim(self.inner.get_result(ordered))
124    }
125
126    /// Resets this union to its initial, empty state.
127    pub fn reset(&mut self) {
128        self.inner.pin_mut().reset();
129    }
130}