lb_tantivy/query/bitset/
mod.rs

1use common::{BitSet, TinySet};
2
3use crate::docset::{DocSet, TERMINATED};
4use crate::DocId;
5
6/// A `BitSetDocSet` makes it possible to iterate through a bitset as if it was a `DocSet`.
7///
8/// # Implementation detail
9///
10/// Skipping is relatively fast here as we can directly point to the
11/// right tiny bitset bucket.
12///
13/// TODO: Consider implementing a `BitTreeSet` in order to advance faster
14/// when the bitset is sparse
15pub struct BitSetDocSet {
16    docs: BitSet,
17    cursor_bucket: u32, //< index associated with the current tiny bitset
18    cursor_tinybitset: TinySet,
19    doc: u32,
20}
21
22impl BitSetDocSet {
23    fn go_to_bucket(&mut self, bucket_addr: u32) {
24        self.cursor_bucket = bucket_addr;
25        self.cursor_tinybitset = self.docs.tinyset(bucket_addr);
26    }
27}
28
29impl From<BitSet> for BitSetDocSet {
30    fn from(docs: BitSet) -> BitSetDocSet {
31        let first_tiny_bitset = if docs.max_value() == 0 {
32            TinySet::empty()
33        } else {
34            docs.tinyset(0)
35        };
36        let mut docset = BitSetDocSet {
37            docs,
38            cursor_bucket: 0,
39            cursor_tinybitset: first_tiny_bitset,
40            doc: 0u32,
41        };
42        docset.advance();
43        docset
44    }
45}
46
47impl DocSet for BitSetDocSet {
48    #[inline]
49    fn advance(&mut self) -> DocId {
50        if let Some(lower) = self.cursor_tinybitset.pop_lowest() {
51            self.doc = (self.cursor_bucket * 64u32) | lower;
52            return self.doc;
53        }
54        if let Some(cursor_bucket) = self.docs.first_non_empty_bucket(self.cursor_bucket + 1) {
55            self.go_to_bucket(cursor_bucket);
56            let lower = self.cursor_tinybitset.pop_lowest().unwrap();
57            self.doc = (cursor_bucket * 64u32) | lower;
58            self.doc
59        } else {
60            self.doc = TERMINATED;
61            TERMINATED
62        }
63    }
64
65    fn seek(&mut self, target: DocId) -> DocId {
66        if target >= self.docs.max_value() {
67            self.doc = TERMINATED;
68            return TERMINATED;
69        }
70        let target_bucket = target / 64u32;
71        if target_bucket > self.cursor_bucket {
72            self.go_to_bucket(target_bucket);
73            let greater_filter: TinySet = TinySet::range_greater_or_equal(target);
74            self.cursor_tinybitset = self.cursor_tinybitset.intersect(greater_filter);
75            self.advance()
76        } else {
77            let mut doc = self.doc();
78            while doc < target {
79                doc = self.advance();
80            }
81            doc
82        }
83    }
84
85    /// Returns the current document
86    fn doc(&self) -> DocId {
87        self.doc
88    }
89
90    /// Returns the number of values set in the underlying bitset.
91    fn size_hint(&self) -> u32 {
92        self.docs.len() as u32
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use std::collections::BTreeSet;
99
100    use common::BitSet;
101
102    use super::BitSetDocSet;
103    use crate::docset::{DocSet, TERMINATED};
104    use crate::tests::generate_nonunique_unsorted;
105    use crate::DocId;
106
107    fn create_docbitset(docs: &[DocId], max_doc: DocId) -> BitSetDocSet {
108        let mut docset = BitSet::with_max_value(max_doc);
109        for &doc in docs {
110            docset.insert(doc);
111        }
112        BitSetDocSet::from(docset)
113    }
114
115    #[test]
116    fn test_bitset_large() {
117        let arr = generate_nonunique_unsorted(100_000, 5_000);
118        let mut btreeset: BTreeSet<u32> = BTreeSet::new();
119        let mut bitset = BitSet::with_max_value(100_000);
120        for el in arr {
121            btreeset.insert(el);
122            bitset.insert(el);
123        }
124        for i in 0..100_000 {
125            assert_eq!(btreeset.contains(&i), bitset.contains(i));
126        }
127        assert_eq!(btreeset.len(), bitset.len());
128        let mut bitset_docset = BitSetDocSet::from(bitset);
129        let mut remaining = true;
130        for el in btreeset.into_iter() {
131            assert!(remaining);
132            assert_eq!(bitset_docset.doc(), el);
133            remaining = bitset_docset.advance() != TERMINATED;
134        }
135        assert!(!remaining);
136    }
137
138    #[test]
139    fn test_empty() {
140        let bitset = BitSet::with_max_value(1000);
141        let mut empty = BitSetDocSet::from(bitset);
142        assert_eq!(empty.advance(), TERMINATED)
143    }
144
145    #[test]
146    fn test_seek_terminated() {
147        let bitset = BitSet::with_max_value(1000);
148        let mut empty = BitSetDocSet::from(bitset);
149        assert_eq!(empty.seek(TERMINATED), TERMINATED)
150    }
151
152    fn test_go_through_sequential(docs: &[DocId]) {
153        let mut docset = create_docbitset(docs, 1_000u32);
154        for &doc in docs {
155            assert_eq!(doc, docset.doc());
156            docset.advance();
157        }
158        assert_eq!(docset.advance(), TERMINATED);
159    }
160
161    #[test]
162    fn test_docbitset_sequential() {
163        test_go_through_sequential(&[1, 2, 3]);
164        test_go_through_sequential(&[1, 2, 3, 4, 5, 63, 64, 65]);
165        test_go_through_sequential(&[63, 64, 65]);
166        test_go_through_sequential(&[1, 2, 3, 4, 95, 96, 97, 98, 99]);
167    }
168
169    #[test]
170    fn test_docbitset_skip() {
171        {
172            let mut docset = create_docbitset(&[1, 5, 6, 7, 5112], 10_000);
173            assert_eq!(docset.seek(7), 7);
174            assert_eq!(docset.doc(), 7);
175            assert_eq!(docset.advance(), 5112);
176            assert_eq!(docset.doc(), 5112);
177            assert_eq!(docset.advance(), TERMINATED);
178        }
179        {
180            let mut docset = create_docbitset(&[1, 5, 6, 7, 5112], 10_000);
181            assert_eq!(docset.seek(3), 5);
182            assert_eq!(docset.doc(), 5);
183            assert_eq!(docset.advance(), 6);
184        }
185        {
186            let mut docset = create_docbitset(&[5112], 10_000);
187            assert_eq!(docset.seek(5112), 5112);
188            assert_eq!(docset.doc(), 5112);
189            assert_eq!(docset.advance(), TERMINATED);
190        }
191        {
192            let mut docset = create_docbitset(&[5112], 10_000);
193            assert_eq!(docset.seek(5113), TERMINATED);
194            assert_eq!(docset.advance(), TERMINATED);
195        }
196        {
197            let mut docset = create_docbitset(&[5112], 10_000);
198            assert_eq!(docset.seek(5111), 5112);
199            assert_eq!(docset.doc(), 5112);
200            assert_eq!(docset.advance(), TERMINATED);
201        }
202        {
203            let mut docset = create_docbitset(&[1, 5, 6, 7, 5112, 5500, 6666], 10_000);
204            assert_eq!(docset.seek(5112), 5112);
205            assert_eq!(docset.doc(), 5112);
206            assert_eq!(docset.advance(), 5500);
207            assert_eq!(docset.doc(), 5500);
208            assert_eq!(docset.advance(), 6666);
209            assert_eq!(docset.doc(), 6666);
210            assert_eq!(docset.advance(), TERMINATED);
211        }
212        {
213            let mut docset = create_docbitset(&[1, 5, 6, 7, 5112, 5500, 6666], 10_000);
214            assert_eq!(docset.seek(5111), 5112);
215            assert_eq!(docset.doc(), 5112);
216            assert_eq!(docset.advance(), 5500);
217            assert_eq!(docset.doc(), 5500);
218            assert_eq!(docset.advance(), 6666);
219            assert_eq!(docset.doc(), 6666);
220            assert_eq!(docset.advance(), TERMINATED);
221        }
222        {
223            let mut docset = create_docbitset(&[1, 5, 6, 7, 5112, 5513, 6666], 10_000);
224            assert_eq!(docset.seek(5111), 5112);
225            assert_eq!(docset.doc(), 5112);
226            assert_eq!(docset.advance(), 5513);
227            assert_eq!(docset.doc(), 5513);
228            assert_eq!(docset.advance(), 6666);
229            assert_eq!(docset.doc(), 6666);
230            assert_eq!(docset.advance(), TERMINATED);
231        }
232    }
233}
234
235#[cfg(all(test, feature = "unstable"))]
236mod bench {
237
238    use super::{BitSet, BitSetDocSet};
239    use crate::docset::TERMINATED;
240    use crate::{test, tests, DocSet};
241
242    #[bench]
243    fn bench_bitset_1pct_insert(b: &mut test::Bencher) {
244        let els = tests::generate_nonunique_unsorted(1_000_000u32, 10_000);
245        b.iter(|| {
246            let mut bitset = BitSet::with_max_value(1_000_000);
247            for el in els.iter().cloned() {
248                bitset.insert(el);
249            }
250        });
251    }
252
253    #[bench]
254    fn bench_bitset_1pct_clone(b: &mut test::Bencher) {
255        let els = tests::generate_nonunique_unsorted(1_000_000u32, 10_000);
256        let mut bitset = BitSet::with_max_value(1_000_000);
257        for el in els {
258            bitset.insert(el);
259        }
260        b.iter(|| bitset.clone());
261    }
262
263    #[bench]
264    fn bench_bitset_1pct_clone_iterate(b: &mut test::Bencher) {
265        let els = tests::sample(1_000_000u32, 0.01);
266        let mut bitset = BitSet::with_max_value(1_000_000);
267        for el in els {
268            bitset.insert(el);
269        }
270        b.iter(|| {
271            let mut docset = BitSetDocSet::from(bitset.clone());
272            while docset.advance() != TERMINATED {}
273        });
274    }
275}
lb_tantivy/query/bitset/mod.rs

lb_tantivy/query/bitset/
mod.rs
