attribute_search_engine/index/prefix/
mod.rs

1mod tree;
2
3use super::SearchIndex;
4use crate::{Query, Result, SearchEngineError, SupportedQueries, SUPPORTS_EXACT, SUPPORTS_PREFIX};
5use std::{collections::HashSet, hash::Hash};
6use tree::HashSetPrefixTree;
7
8/// SearchIndexPrefixTree is a index backed by a prefix tree that can match
9/// Exact and Prefix queries. It can only store String attribute values.
10///
11/// # Example
12/// ```
13/// use attribute_search_engine::{SearchIndex, SearchIndexPrefixTree};
14/// use std::collections::HashSet;
15/// use attribute_search_engine::Query;
16///
17/// let mut index_firstname = SearchIndexPrefixTree::<usize>::new();
18/// index_firstname.insert(0, "Alex".into());
19/// index_firstname.insert(1, "Alexander".into());
20/// index_firstname.insert(2, "Andrea".into());
21/// index_firstname.insert(3, "Ben".into());
22///
23/// let result = index_firstname.search(&Query::Exact("<unused>".into(), "Alex".into()));
24/// assert_eq!(result, Ok(HashSet::from_iter(vec![0])));
25///
26/// let result = index_firstname.search(&Query::Prefix("<unused>".into(), "Alex".into()));
27/// assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1])));
28/// ```
29pub struct SearchIndexPrefixTree<P> {
30    index: HashSetPrefixTree<P>,
31}
32
33impl<P: Eq + Hash + Clone> Default for SearchIndexPrefixTree<P> {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl<P: Eq + Hash + Clone> SearchIndexPrefixTree<P> {
40    /// Creates a new `SearchIndexPrefixTree`.
41    ///
42    /// # Example
43    /// ```rust
44    /// use attribute_search_engine::SearchIndexPrefixTree;
45    ///
46    /// let index = SearchIndexPrefixTree::<usize>::new();
47    /// ```
48    pub fn new() -> Self {
49        Self {
50            index: HashSetPrefixTree::new(),
51        }
52    }
53
54    /// Insert a new entry in the index.
55    ///
56    /// # Example
57    /// ```rust
58    /// use attribute_search_engine::SearchIndexPrefixTree;
59    ///
60    /// let mut index = SearchIndexPrefixTree::<usize>::new();
61    ///
62    /// // You insert an entry by giving a row / primary id and an attribute value:
63    /// index.insert(123, "Hello".into());
64    /// // The same row / primary id can have multiple attributes assigned:
65    /// index.insert(123, "World".into());
66    /// // Add as much entries as you want for as many rows you want:
67    /// index.insert(124, "Rust".into());
68    /// ```
69    pub fn insert(&mut self, primary_id: P, attribute_value: String) {
70        self.index.insert(&attribute_value, primary_id);
71    }
72}
73
74impl<P: Eq + Hash + Clone> SearchIndex<P> for SearchIndexPrefixTree<P> {
75    fn search(&self, query: &Query) -> Result<HashSet<P>> {
76        match query {
77            Query::Exact(_, value) => Ok(self.index.get(value).unwrap_or_default()),
78            Query::Prefix(_, value) => Ok(self.index.get_prefix(value).unwrap_or_default()),
79            _ => Err(SearchEngineError::UnsupportedQuery),
80        }
81    }
82
83    fn supported_queries(&self) -> SupportedQueries {
84        SUPPORTS_EXACT | SUPPORTS_PREFIX
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn search_index_exact_string() {
94        let mut index = SearchIndexPrefixTree::<usize>::new();
95        index.insert(0, "A".into());
96        index.insert(0, "B".into());
97        index.insert(0, "C".into());
98        index.insert(1, "A".into());
99        index.insert(1, "B".into());
100        index.insert(2, "A".into());
101
102        let result = index.search(&Query::Exact("<not used>".into(), "A".into()));
103        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2])));
104
105        let result = index.search(&Query::Exact("<not used>".into(), "B".into()));
106        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1])));
107
108        let result = index.search(&Query::Exact("<not used>".into(), "C".into()));
109        assert_eq!(result, Ok(HashSet::from_iter(vec![0])));
110
111        let result = index.search(&Query::Exact("<not used>".into(), "D".into()));
112        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
113    }
114
115    #[test]
116    fn search_index_prefix_string() {
117        let mut index = SearchIndexPrefixTree::<usize>::new();
118        index.insert(0, "A".into());
119        index.insert(1, "AA".into());
120        index.insert(2, "AB".into());
121        index.insert(3, "ABA".into());
122        index.insert(4, "ABB".into());
123        index.insert(5, "ABC".into());
124        index.insert(6, "B".into());
125        index.insert(7, "BA".into());
126        index.insert(8, "BB".into());
127
128        let result = index.search(&Query::Prefix("<not used>".into(), "A".into()));
129        assert_eq!(result, Ok(HashSet::from_iter(vec![0, 1, 2, 3, 4, 5])));
130
131        let result = index.search(&Query::Prefix("<not used>".into(), "B".into()));
132        assert_eq!(result, Ok(HashSet::from_iter(vec![6, 7, 8])));
133
134        let result = index.search(&Query::Prefix("<not used>".into(), "C".into()));
135        assert_eq!(result, Ok(HashSet::from_iter(vec![])));
136
137        let result = index.search(&Query::Prefix("<not used>".into(), "AB".into()));
138        assert_eq!(result, Ok(HashSet::from_iter(vec![2, 3, 4, 5])));
139
140        let result = index.search(&Query::Prefix("<not used>".into(), "ABC".into()));
141        assert_eq!(result, Ok(HashSet::from_iter(vec![5])));
142    }
143
144    #[test]
145    fn search_index_unsupported_queries() {
146        let mut index = SearchIndexPrefixTree::<usize>::new();
147        index.insert(0, "".into());
148
149        assert_eq!(
150            index.search(&Query::InRange("<not used>".into(), "0".into(), "1".into())),
151            Err(SearchEngineError::UnsupportedQuery)
152        );
153        assert_eq!(
154            index.search(&Query::OutRange(
155                "<not used>".into(),
156                "0".into(),
157                "1".into()
158            )),
159            Err(SearchEngineError::UnsupportedQuery)
160        );
161        assert_eq!(
162            index.search(&Query::Minimum("<not used>".into(), "0".into())),
163            Err(SearchEngineError::UnsupportedQuery)
164        );
165        assert_eq!(
166            index.search(&Query::Maximum("<not used>".into(), "0".into())),
167            Err(SearchEngineError::UnsupportedQuery)
168        );
169        assert_eq!(
170            index.search(&Query::Or(vec![])),
171            Err(SearchEngineError::UnsupportedQuery)
172        );
173        assert_eq!(
174            index.search(&Query::And(vec![])),
175            Err(SearchEngineError::UnsupportedQuery)
176        );
177        assert_eq!(
178            index.search(&Query::Exclude(
179                Box::new(Query::Exact("<not used>".into(), "0".into())),
180                vec![]
181            )),
182            Err(SearchEngineError::UnsupportedQuery)
183        );
184    }
185}