ddsketch/canonical/store/
mod.rs

1//! Sketch storage.
2
3use datadog_protos::sketches::Store as ProtoStore;
4
5use super::error::ProtoConversionError;
6
7mod collapsing_highest;
8pub use self::collapsing_highest::CollapsingHighestDenseStore;
9
10mod collapsing_lowest;
11pub use self::collapsing_lowest::CollapsingLowestDenseStore;
12
13mod dense;
14pub use self::dense::DenseStore;
15
16mod sparse;
17pub use self::sparse::SparseStore;
18
19/// Storage for sketch observations.
20///
21/// Stores manage holding the counts of mapped values, such that they contain a list of bins and the number of
22/// observations currently counted in each bin.
23pub trait Store: Clone + Send + Sync {
24    /// Adds a count to the bin at the given index.
25    fn add(&mut self, index: i32, count: u64);
26
27    /// Returns the total count across all bins.
28    fn total_count(&self) -> u64;
29
30    /// Returns the minimum index with a non-zero count, or `None` if empty.
31    fn min_index(&self) -> Option<i32>;
32
33    /// Returns the maximum index with a non-zero count, or `None` if empty.
34    fn max_index(&self) -> Option<i32>;
35
36    /// Returns the index of the bin containing the given rank.
37    ///
38    /// The rank is 0-indexed, so rank 0 is the first observation.
39    fn key_at_rank(&self, rank: u64) -> Option<i32>;
40
41    /// Merges another store into this one.
42    fn merge(&mut self, other: &Self);
43
44    /// Returns `true` if the store is empty.
45    fn is_empty(&self) -> bool;
46
47    /// Clears all bins from the store.
48    fn clear(&mut self);
49
50    /// Populates this store from a protobuf `Store`.
51    fn merge_from_proto(&mut self, proto: &ProtoStore) -> Result<(), ProtoConversionError>;
52
53    /// Converts this store to a protobuf `Store`.
54    fn to_proto(&self) -> ProtoStore;
55}
56
57/// Validates and converts a protobuf `f64` count to `u64`.
58///
59/// # Errors
60///
61/// If the count is negative, or has a fractional part, an error is returned.
62pub(crate) fn validate_proto_count(index: i32, count: f64) -> Result<u64, ProtoConversionError> {
63    if count < 0.0 {
64        return Err(ProtoConversionError::NegativeBinCount { index, count });
65    }
66    if count.fract() != 0.0 {
67        return Err(ProtoConversionError::NonIntegerBinCount { index, count });
68    }
69    Ok(count as u64)
70}
71
72/// Generates the shared [`Store`] trait conformance suite for a concrete store type.
73///
74/// Every `Store` implementation shares the same observable contract for adding counts, reporting
75/// totals/min/max, ranking, merging, clearing, and round-tripping through protobuf. Rather than hand-duplicate that
76/// suite across each sibling implementation, each implementation's `tests` module invokes this macro with its
77/// concrete type. Implementation-specific behavior (collapsing, private-field layout, the dense-only bin iterator)
78/// is still covered by inline tests in the respective module.
79///
80/// The store type must implement `Default`. The collapsing stores default to 2048 bins, which is far larger than
81/// any index range these cases use, so no collapsing occurs and the shared assertions hold for every
82/// implementation.
83#[cfg(test)]
84macro_rules! store_conformance_tests {
85    ($store:ty) => {
86        #[test]
87        fn add_single_value_sets_count_and_bounds() {
88            let mut store = <$store>::default();
89            store.add(5, 1);
90
91            assert_eq!(store.total_count(), 1);
92            assert_eq!(store.min_index(), Some(5));
93            assert_eq!(store.max_index(), Some(5));
94        }
95
96        #[test]
97        fn add_repeated_at_same_index_accumulates() {
98            let mut store = <$store>::default();
99            store.add(5, 3);
100            store.add(5, 2);
101
102            assert_eq!(store.total_count(), 5);
103            assert_eq!(store.min_index(), Some(5));
104            assert_eq!(store.max_index(), Some(5));
105        }
106
107        #[test]
108        fn add_distinct_indices_tracks_distribution_and_bounds() {
109            let mut store = <$store>::default();
110            store.add(5, 1);
111            store.add(10, 2);
112            store.add(3, 3);
113
114            assert_eq!(store.total_count(), 6);
115            assert_eq!(store.min_index(), Some(3));
116            assert_eq!(store.max_index(), Some(10));
117            assert_eq!(
118                $crate::canonical::store::conformance::distribution(&store),
119                [(3, 3), (5, 1), (10, 2)].into_iter().collect()
120            );
121        }
122
123        #[test]
124        fn add_with_zero_count_is_a_noop() {
125            let mut store = <$store>::default();
126            store.add(5, 0);
127
128            assert!(store.is_empty());
129            assert_eq!(store.total_count(), 0);
130            assert_eq!(store.min_index(), None);
131        }
132
133        #[test]
134        fn key_at_rank_maps_each_rank_to_its_index() {
135            let mut store = <$store>::default();
136            store.add(5, 3);
137            store.add(10, 2);
138
139            assert_eq!(store.key_at_rank(0), Some(5));
140            assert_eq!(store.key_at_rank(2), Some(5));
141            assert_eq!(store.key_at_rank(3), Some(10));
142            assert_eq!(store.key_at_rank(4), Some(10));
143            assert_eq!(store.key_at_rank(5), None);
144        }
145
146        #[test]
147        fn empty_store_reports_no_min_max_or_rank() {
148            let store = <$store>::default();
149
150            assert!(store.is_empty());
151            assert_eq!(store.total_count(), 0);
152            assert_eq!(store.min_index(), None);
153            assert_eq!(store.max_index(), None);
154            assert_eq!(store.key_at_rank(0), None);
155        }
156
157        #[test]
158        fn merge_combines_distributions() {
159            let mut store1 = <$store>::default();
160            store1.add(5, 2);
161            store1.add(10, 1);
162
163            let mut store2 = <$store>::default();
164            store2.add(5, 1);
165            store2.add(15, 3);
166
167            store1.merge(&store2);
168
169            assert_eq!(store1.total_count(), 7);
170            assert_eq!(store1.min_index(), Some(5));
171            assert_eq!(store1.max_index(), Some(15));
172            assert_eq!(
173                $crate::canonical::store::conformance::distribution(&store1),
174                [(5, 3), (10, 1), (15, 3)].into_iter().collect()
175            );
176        }
177
178        #[test]
179        fn merge_from_empty_store_is_a_noop() {
180            let mut store = <$store>::default();
181            store.add(5, 2);
182
183            let empty = <$store>::default();
184            store.merge(&empty);
185
186            assert_eq!(store.total_count(), 2);
187            assert_eq!(
188                $crate::canonical::store::conformance::distribution(&store),
189                [(5, 2)].into_iter().collect()
190            );
191        }
192
193        #[test]
194        fn clear_resets_to_empty() {
195            let mut store = <$store>::default();
196            store.add(5, 2);
197            store.add(10, 1);
198
199            store.clear();
200
201            assert!(store.is_empty());
202            assert_eq!(store.total_count(), 0);
203            assert_eq!(store.min_index(), None);
204        }
205
206        #[test]
207        fn handles_negative_indices() {
208            let mut store = <$store>::default();
209            store.add(-5, 1);
210            store.add(5, 1);
211
212            assert_eq!(store.total_count(), 2);
213            assert_eq!(store.min_index(), Some(-5));
214            assert_eq!(store.max_index(), Some(5));
215            assert_eq!(
216                $crate::canonical::store::conformance::distribution(&store),
217                [(-5, 1), (5, 1)].into_iter().collect()
218            );
219        }
220
221        #[test]
222        fn handles_widely_scattered_indices() {
223            let mut store = <$store>::default();
224            store.add(-1000, 1);
225            store.add(0, 2);
226            store.add(1000, 3);
227
228            assert_eq!(store.total_count(), 6);
229            assert_eq!(store.min_index(), Some(-1000));
230            assert_eq!(store.max_index(), Some(1000));
231            assert_eq!(
232                $crate::canonical::store::conformance::distribution(&store),
233                [(-1000, 1), (0, 2), (1000, 3)].into_iter().collect()
234            );
235        }
236
237        #[test]
238        fn proto_round_trip_preserves_distribution() {
239            let mut store = <$store>::default();
240            store.add(-3, 4);
241            store.add(5, 2);
242            store.add(10, 1);
243
244            let proto = store.to_proto();
245            let mut restored = <$store>::default();
246            restored
247                .merge_from_proto(&proto)
248                .expect("round-tripped proto must be valid");
249
250            assert_eq!(restored.total_count(), store.total_count());
251            assert_eq!(
252                $crate::canonical::store::conformance::distribution(&restored),
253                $crate::canonical::store::conformance::distribution(&store)
254            );
255        }
256    };
257}
258
259#[cfg(test)]
260pub(crate) use store_conformance_tests;
261
262#[cfg(test)]
263pub(crate) mod conformance {
264    use std::collections::BTreeMap;
265
266    use super::Store;
267
268    /// Reconstructs the full `index -> count` distribution of a store using only the public [`Store`] trait API.
269    ///
270    /// This walks every rank in `[0, total_count)` and tallies the index each rank maps to. Doing this through the
271    /// trait (rather than each implementation's private fields) lets the shared conformance suite assert exact bin
272    /// distributions for every `Store`, including after merges and protobuf round-trips.
273    pub(crate) fn distribution<S: Store>(store: &S) -> BTreeMap<i32, u64> {
274        let mut dist = BTreeMap::new();
275        for rank in 0..store.total_count() {
276            let index = store
277                .key_at_rank(rank)
278                .expect("every rank below total_count must map to an index");
279            *dist.entry(index).or_insert(0) += 1;
280        }
281        dist
282    }
283}
ddsketch/canonical/store/mod.rs

ddsketch/canonical/store/
mod.rs
