zeppelin_embed/meta/
bitmap.rs

1//! Small ownership boundary around the selected compressed-bitmap crate.
2
3use roaring::RoaringBitmap;
4
5/// A compressed set of segment-local document identifiers.
6#[derive(Clone, Debug, Default, PartialEq)]
7pub struct DocBitmap {
8    inner: RoaringBitmap,
9}
10
11impl Eq for DocBitmap {}
12
13impl DocBitmap {
14    pub(crate) fn compact_for_cache(&mut self) {
15        self.inner = self.inner.clone();
16    }
17
18    pub(crate) fn resident_bytes(&self) -> Option<usize> {
19        let statistics = self.inner.statistics();
20        #[cfg(target_pointer_width = "64")]
21        const CONTAINER_BYTES: usize = 32;
22        #[cfg(target_pointer_width = "32")]
23        const CONTAINER_BYTES: usize = 16;
24        let containers = (statistics.n_containers as usize).checked_mul(CONTAINER_BYTES)?;
25        let arrays = (statistics.n_values_array_containers as usize)
26            .checked_mul(std::mem::size_of::<u16>())?;
27        let bitsets = usize::try_from(statistics.n_bytes_bitset_containers).ok()?;
28        containers.checked_add(arrays)?.checked_add(bitsets)
29    }
30
31    /// Creates an empty bitmap.
32    #[must_use]
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    /// Creates the document universe `0..row_count`.
38    #[must_use]
39    pub fn full(row_count: u32) -> Self {
40        Self {
41            inner: (0..row_count).collect(),
42        }
43    }
44
45    /// Creates a bitmap from document identifiers.
46    #[must_use]
47    pub fn from_ids(ids: impl IntoIterator<Item = u32>) -> Self {
48        Self {
49            inner: ids.into_iter().collect(),
50        }
51    }
52
53    /// Inserts one document identifier and reports whether it was new.
54    pub fn insert(&mut self, document: u32) -> bool {
55        self.inner.insert(document)
56    }
57
58    /// Removes one document identifier and reports whether it was present.
59    pub fn remove(&mut self, document: u32) -> bool {
60        self.inner.remove(document)
61    }
62
63    /// Returns whether the document identifier is present.
64    #[must_use]
65    pub fn contains(&self, document: u32) -> bool {
66        self.inner.contains(document)
67    }
68
69    /// Returns the exact bitmap cardinality.
70    #[must_use]
71    pub fn cardinality(&self) -> u64 {
72        self.inner.len()
73    }
74
75    /// Returns whether the bitmap contains no identifiers.
76    #[must_use]
77    pub fn is_empty(&self) -> bool {
78        self.inner.is_empty()
79    }
80
81    /// Iterates identifiers in ascending order.
82    pub fn iter(&self) -> impl Iterator<Item = u32> + '_ {
83        self.inner.iter().inspect(|_| {
84            #[cfg(any(test, feature = "test-seams"))]
85            super::bitmap_observer::row();
86        })
87    }
88
89    /// Finds the first row outside a dense `0..row_count` segment without
90    /// walking its valid prefix. The inclusive lower bound also covers MAX.
91    pub(crate) fn first_at_or_after(&self, row_count: u32) -> Option<u32> {
92        #[cfg(any(test, feature = "test-seams"))]
93        super::bitmap_observer::probe();
94        self.inner.range(row_count..).next()
95    }
96
97    pub(crate) const fn as_roaring(&self) -> &RoaringBitmap {
98        &self.inner
99    }
100
101    /// Mutates this bitmap to its union with `other`.
102    pub fn union_with(&mut self, other: &Self) {
103        self.inner |= &other.inner;
104    }
105
106    /// Mutates this bitmap to its intersection with `other`.
107    pub fn intersect_with(&mut self, other: &Self) {
108        self.inner &= &other.inner;
109    }
110
111    /// Removes every identifier present in `other`.
112    pub fn subtract(&mut self, other: &Self) {
113        self.inner -= &other.inner;
114    }
115
116    /// Returns whether every identifier is also present in `other`.
117    #[must_use]
118    pub fn is_subset(&self, other: &Self) -> bool {
119        self.inner.is_subset(&other.inner)
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::DocBitmap;
126
127    #[test]
128    fn astra_15_empty_zero_and_max_u32_bitmap_bounds_are_checked() {
129        let empty = DocBitmap::new();
130        assert_eq!(empty.first_at_or_after(0), None);
131        assert_eq!(empty.first_at_or_after(u32::MAX), None);
132        let sparse = DocBitmap::from_ids([0, 2, 65_535, 65_536, u32::MAX - 1]);
133        for (bound, expected) in [
134            (0, Some(0)),
135            (1, Some(2)),
136            (3, Some(65_535)),
137            (65_536, Some(65_536)),
138            (65_537, Some(u32::MAX - 1)),
139            (u32::MAX, None),
140        ] {
141            assert_eq!(sparse.first_at_or_after(bound), expected);
142        }
143        let maximum = DocBitmap::from_ids([u32::MAX]);
144        assert_eq!(maximum.first_at_or_after(u32::MAX), Some(u32::MAX));
145        let dense = DocBitmap::full(131_072);
146        assert_eq!(dense.first_at_or_after(65_536), Some(65_536));
147        assert_eq!(dense.first_at_or_after(131_071), Some(131_071));
148        assert_eq!(dense.first_at_or_after(131_072), None);
149    }
150}
zeppelin_embed/meta/bitmap.rs

zeppelin_embed/meta/
bitmap.rs
