attribute_search_engine/index/
btree_range.rs

1use super::{string_to_payload_type, SearchIndex};
2use crate::{
3    Query, Result, SearchEngineError, SupportedQueries, SUPPORTS_EXACT, SUPPORTS_INRANGE,
4    SUPPORTS_MAXIMUM, SUPPORTS_MINIMUM, SUPPORTS_OUTRANGE,
5};
6use std::{
7    collections::{BTreeMap, HashSet},
8    hash::Hash,
9    ops::{Bound, RangeBounds},
10    str::FromStr,
11};
12
13/// SearchIndexBTreeRange is a index backed by a BTreeMap that can match
14/// Exact, InRange, OutRange, Minimum and Maximum queries.
15///
16/// # Example
17/// ```
18/// use attribute_search_engine::{SearchIndex, SearchIndexBTreeRange};
19/// use std::collections::HashSet;
20/// use attribute_search_engine::Query;
21///
22/// let mut index_age = SearchIndexBTreeRange::<usize, i32>::new();
23/// index_age.insert(0, 17);
24/// index_age.insert(1, 42);
25/// index_age.insert(2, 31);
26/// index_age.insert(3, 26);
27///
28/// let result = index_age.search(&Query::Exact("<unused>".into(), "42".into()));
29/// assert_eq!(result, Ok(HashSet::from_iter(vec![1])));
30///
31/// let result = index_age.search(&Query::InRange("<unused>".into(), "20".into(), "40".into()));
32/// assert_eq!(result, Ok(HashSet::from_iter(vec![2, 3])));
33/// ```
34pub struct SearchIndexBTreeRange<P, V> {
35    index: BTreeMap<V, HashSet<P>>,
36}
37
38impl<P, V> Default for SearchIndexBTreeRange<P, V>
39where
40    P: Eq + Hash + Clone + 'static,
41    V: Ord + FromStr + 'static,
42{
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl<P, V> SearchIndexBTreeRange<P, V>
49where
50    P: Eq + Hash + Clone + 'static,
51    V: Ord + FromStr + 'static,
52{
53    /// Creates a new `SearchIndexBTreeRange`.
54    ///
55    /// # Example
56    /// ```rust
57    /// use attribute_search_engine::SearchIndexBTreeRange;
58    ///
59    /// let index = SearchIndexBTreeRange::<usize, i32>::new();
60    /// ```
61    pub fn new() -> Self {
62        Self {
63            index: BTreeMap::new(),
64        }
65    }
66
67    /// Insert a new entry in the index.
68    ///
69    /// # Example
70    /// ```rust
71    /// use attribute_search_engine::SearchIndexBTreeRange;
72    ///
73    /// let mut index = SearchIndexBTreeRange::<usize, i32>::new();
74    ///
75    /// // You insert an entry by giving a row / primary id and an attribute value:
76    /// index.insert(123, 42);
77    /// // The same row / primary id can have multiple attributes assigned:
78    /// index.insert(123, 69);
79    /// // Add as much entries as you want for as many rows you want:
80    /// index.insert(124, 32);
81    /// ```
82    pub fn insert(&mut self, primary_id: P, attribute_value: V) {
83        self.index
84            .entry(attribute_value)
85            .or_default()
86            .insert(primary_id);
87    }
88
89    /// This internal function helps with searching for all kinds of
90    /// ranges and merging the result to a HashSet.
91    fn search_range(&self, range: impl RangeBounds<V>) -> HashSet<P> {
92        let mut result_set = HashSet::<P>::new();
93        for (_, primary_set) in self.index.range(range) {
94            result_set = result_set.union(primary_set).cloned().collect();
95        }
96        result_set
97    }
98}
99
100impl<P, V> SearchIndex<P> for SearchIndexBTreeRange<P, V>
101where
102    P: Eq + Hash + Clone + 'static,
103    V: Ord + FromStr + 'static,
104{
105    fn search(&self, query: &Query) -> Result<HashSet<P>> {
106        match query {
107            Query::Exact(_, value_str) => {
108                let value: V = string_to_payload_type(value_str)?;
109                Ok(self.index.get(&value).cloned().unwrap_or(HashSet::new()))
110            }
111            Query::InRange(_, min_str, max_str) => {
112                let min: V = string_to_payload_type(min_str)?;
113                let max: V = string_to_payload_type(max_str)?;
114                if min > max {
115                    return Ok(HashSet::new());
116                }
117                Ok(self.search_range(min..=max))
118            }
119            Query::Minimum(_, min_str) => {
120                let min: V = string_to_payload_type(min_str)?;
121                Ok(self.search_range(min..))
122            }
123            Query::Maximum(_, max_str) => {
124                let max: V = string_to_payload_type(max_str)?;
125                Ok(self.search_range(..=max))
126            }
127            Query::OutRange(_, start_str, end_str) => {
128                let start: V = string_to_payload_type(start_str)?;
129                let end: V = string_to_payload_type(end_str)?;
130                if start > end {
131                    return Ok(HashSet::new());
132                }
133                Ok(self
134                    .search_range(..start)
135                    .union(&self.search_range((Bound::Excluded(end), Bound::Unbounded)))
136                    .cloned()
137                    .collect())
138            }
139            _ => Err(SearchEngineError::UnsupportedQuery),
140        }
141    }
142
143    fn supported_queries(&self) -> SupportedQueries {
144        SUPPORTS_EXACT | SUPPORTS_INRANGE | SUPPORTS_MINIMUM | SUPPORTS_MAXIMUM | SUPPORTS_OUTRANGE
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn search_index_exact_string() {
154        let mut index = SearchIndexBTreeRange::<usize, String>::new();
155        index.insert(0, "A".into());
156        index.insert(0, "B".into());
157        index.insert(0, "C".into());
158        index.insert(1, "A".into());
159        index.insert(1, "B".into());
160        index.insert(2, "A".into());
161
162        let result = index.search(&Query::Exact("<not used>".into(), "A".into()));
163        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2])));
164
165        let result = index.search(&Query::Exact("<not used>".into(), "B".into()));
166        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1])));
167
168        let result = index.search(&Query::Exact("<not used>".into(), "C".into()));
169        assert_eq!(result, Ok(HashSet::from_iter(vec![0])));
170
171        let result = index.search(&Query::Exact("<not used>".into(), "D".into()));
172        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
173    }
174
175    #[test]
176    fn search_index_exact_number() {
177        let mut index = SearchIndexBTreeRange::<usize, i32>::new();
178        index.insert(0, 0);
179        index.insert(0, 1);
180        index.insert(0, 2);
181        index.insert(1, 0);
182        index.insert(1, 1);
183        index.insert(2, 0);
184
185        let result = index.search(&Query::Exact("<not used>".into(), "0".into()));
186        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2])));
187
188        let result = index.search(&Query::Exact("<not used>".into(), "1".into()));
189        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1])));
190
191        let result = index.search(&Query::Exact("<not used>".into(), "2".into()));
192        assert_eq!(result, Ok(HashSet::from_iter(vec![0])));
193
194        let result = index.search(&Query::Exact("<not used>".into(), "4".into()));
195        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
196    }
197
198    #[test]
199    fn search_index_inrange_number() {
200        let mut index = SearchIndexBTreeRange::<usize, i32>::new();
201        index.insert(0, 00);
202        index.insert(1, 10);
203        index.insert(2, 20);
204        index.insert(3, 30);
205        index.insert(4, 40);
206        index.insert(5, 50);
207
208        let result = index.search(&Query::InRange(
209            "<not used>".into(),
210            "0".into(),
211            "50".into(),
212        ));
213        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3, 4, 5])));
214
215        let result = index.search(&Query::InRange(
216            "<not used>".into(),
217            "10".into(),
218            "40".into(),
219        ));
220        assert_eq!(result, Ok(HashSet::from_iter(vec![1, 2, 3, 4])));
221
222        let result = index.search(&Query::InRange(
223            "<not used>".into(),
224            "-20".into(),
225            "20".into(),
226        ));
227        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2])));
228
229        let result = index.search(&Query::InRange(
230            "<not used>".into(),
231            "10".into(),
232            "10".into(),
233        ));
234        assert_eq!(result, Ok(HashSet::from_iter(vec![1])));
235
236        let result = index.search(&Query::InRange(
237            "<not used>".into(),
238            "50".into(),
239            "10".into(),
240        ));
241        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
242
243        let result = index.search(&Query::InRange(
244            "<not used>".into(),
245            "51".into(),
246            "100".into(),
247        ));
248        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
249    }
250
251    #[test]
252    fn search_index_outrange_number() {
253        let mut index = SearchIndexBTreeRange::<usize, i32>::new();
254        index.insert(0, 00);
255        index.insert(1, 10);
256        index.insert(2, 20);
257        index.insert(3, 30);
258        index.insert(4, 40);
259        index.insert(5, 50);
260
261        let result = index.search(&Query::OutRange(
262            "<not used>".into(),
263            "0".into(),
264            "50".into(),
265        ));
266        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
267
268        let result = index.search(&Query::OutRange(
269            "<not used>".into(),
270            "10".into(),
271            "40".into(),
272        ));
273        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 5])));
274
275        let result = index.search(&Query::OutRange(
276            "<not used>".into(),
277            "-20".into(),
278            "20".into(),
279        ));
280        assert_eq!(result, Ok(HashSet::from_iter(vec![3, 4, 5])));
281
282        let result = index.search(&Query::OutRange(
283            "<not used>".into(),
284            "10".into(),
285            "10".into(),
286        ));
287        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 2, 3, 4, 5])));
288
289        let result = index.search(&Query::OutRange(
290            "<not used>".into(),
291            "50".into(),
292            "10".into(),
293        ));
294        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
295
296        let result = index.search(&Query::OutRange(
297            "<not used>".into(),
298            "51".into(),
299            "100".into(),
300        ));
301        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3, 4, 5])));
302    }
303
304    #[test]
305    fn search_index_minimum_number() {
306        let mut index = SearchIndexBTreeRange::<usize, i32>::new();
307        index.insert(0, 00);
308        index.insert(1, 10);
309        index.insert(2, 20);
310        index.insert(3, 30);
311        index.insert(4, 40);
312        index.insert(5, 50);
313
314        let result = index.search(&Query::Minimum("<not used>".into(), "0".into()));
315        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3, 4, 5])));
316
317        let result = index.search(&Query::Minimum("<not used>".into(), "10".into()));
318        assert_eq!(result, Ok(HashSet::from_iter(vec![1, 2, 3, 4, 5])));
319
320        let result = index.search(&Query::Minimum("<not used>".into(), "-20".into()));
321        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3, 4, 5])));
322
323        let result = index.search(&Query::Minimum("<not used>".into(), "30".into()));
324        assert_eq!(result, Ok(HashSet::from_iter(vec![3, 4, 5])));
325
326        let result = index.search(&Query::Minimum("<not used>".into(), "50".into()));
327        assert_eq!(result, Ok(HashSet::from_iter(vec![5])));
328
329        let result = index.search(&Query::Minimum("<not used>".into(), "51".into()));
330        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
331    }
332
333    #[test]
334    fn search_index_maximum_number() {
335        let mut index = SearchIndexBTreeRange::<usize, i32>::new();
336        index.insert(0, 00);
337        index.insert(1, 10);
338        index.insert(2, 20);
339        index.insert(3, 30);
340        index.insert(4, 40);
341        index.insert(5, 50);
342
343        let result = index.search(&Query::Maximum("<not used>".into(), "50".into()));
344        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3, 4, 5])));
345
346        let result = index.search(&Query::Maximum("<not used>".into(), "40".into()));
347        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3, 4])));
348
349        let result = index.search(&Query::Maximum("<not used>".into(), "30".into()));
350        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3])));
351
352        let result = index.search(&Query::Maximum("<not used>".into(), "0".into()));
353        assert_eq!(result, Ok(HashSet::from_iter(vec![0])));
354
355        let result = index.search(&Query::Maximum("<not used>".into(), "-1".into()));
356        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
357    }
358
359    #[test]
360    fn search_index_unsupported_queries() {
361        let mut index = SearchIndexBTreeRange::<usize, i32>::new();
362        index.insert(0, 0);
363
364        assert_eq!(
365            index.search(&Query::Prefix("<not used>".into(), "0".into())),
366            Err(SearchEngineError::UnsupportedQuery)
367        );
368        assert_eq!(
369            index.search(&Query::Or(vec![])),
370            Err(SearchEngineError::UnsupportedQuery)
371        );
372        assert_eq!(
373            index.search(&Query::And(vec![])),
374            Err(SearchEngineError::UnsupportedQuery)
375        );
376        assert_eq!(
377            index.search(&Query::Exclude(
378                Box::new(Query::Exact("<not used>".into(), "0".into())),
379                vec![]
380            )),
381            Err(SearchEngineError::UnsupportedQuery)
382        );
383    }
384}
attribute_search_engine/index/btree_range.rs

attribute_search_engine/index/
btree_range.rs
