index/fulltext_index/
tokenizer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use common_base::term_token::{self, has_irregular_lowercase};
16
17use crate::Bytes;
18use crate::bloom_filter::element_hash;
19use crate::fulltext_index::error::Result;
20
21/// `Tokenizer` tokenizes a text into a list of tokens.
22pub trait Tokenizer: Send {
23    fn tokenize<'a>(&self, text: &'a str) -> Vec<&'a str>;
24}
25
26/// `ScriptTokenizer` produces the tokens described in [`term_token::tokenize`].
27#[derive(Debug, Default)]
28pub struct ScriptTokenizer;
29
30impl Tokenizer for ScriptTokenizer {
31    fn tokenize<'a>(&self, text: &'a str) -> Vec<&'a str> {
32        term_token::tokenize(text)
33    }
34}
35
36/// `Analyzer` analyzes a text into a list of tokens.
37///
38/// It uses a `Tokenizer` to tokenize the text and optionally lowercases the tokens.
39/// When case-insensitive, it also emits the tokens of the whole lowercased text if
40/// they can differ, so that predicates on `lower(col)` can probe the index.
41pub struct Analyzer {
42    tokenizer: Box<dyn Tokenizer>,
43    case_sensitive: bool,
44}
45
46impl Analyzer {
47    /// Creates a new `Analyzer` with the given `Tokenizer` and case sensitivity.
48    pub fn new(tokenizer: Box<dyn Tokenizer>, case_sensitive: bool) -> Self {
49        Self {
50            tokenizer,
51            case_sensitive,
52        }
53    }
54
55    /// Returns the bloom filter hash of each token in the given text.
56    ///
57    /// Equivalent to hashing every token returned by [`Analyzer::analyze_text`] with
58    /// [`element_hash`]. Only case-insensitive non-ASCII tokens allocate, in `to_lowercase`;
59    /// case-insensitive ASCII tokens are lowercased in `buf`.
60    pub fn analyze_text_hashes<'a>(
61        &self,
62        text: &'a str,
63        buf: &'a mut Vec<u8>,
64    ) -> impl Iterator<Item = u64> + use<'a> {
65        let case_sensitive = self.case_sensitive;
66        let lowercased_text_hashes = self
67            .irregular_lowercased_text(text)
68            .map(|lowered| {
69                self.tokenizer
70                    .tokenize(&lowered)
71                    .into_iter()
72                    .map(|token| element_hash(token.as_bytes()))
73                    .collect::<Vec<_>>()
74            })
75            .unwrap_or_default();
76        self.tokenizer
77            .tokenize(text)
78            .into_iter()
79            .map(move |token| {
80                if case_sensitive {
81                    element_hash(token.as_bytes())
82                } else if token.is_ascii() {
83                    buf.clear();
84                    buf.extend(token.bytes().map(|b| b.to_ascii_lowercase()));
85                    element_hash(buf)
86                } else {
87                    element_hash(token.to_lowercase().as_bytes())
88                }
89            })
90            .chain(lowercased_text_hashes)
91    }
92
93    /// Analyzes the given text into a list of tokens.
94    pub fn analyze_text(&self, text: &str) -> Result<Vec<Bytes>> {
95        let mut res = self
96            .tokenizer
97            .tokenize(text)
98            .iter()
99            .map(|s| {
100                if self.case_sensitive {
101                    s.as_bytes().to_vec()
102                } else {
103                    s.to_lowercase().as_bytes().to_vec()
104                }
105            })
106            .collect::<Vec<_>>();
107        if let Some(lowered) = self.irregular_lowercased_text(text) {
108            res.extend(
109                self.tokenizer
110                    .tokenize(&lowered)
111                    .into_iter()
112                    .map(|token| token.as_bytes().to_vec()),
113            );
114        }
115        Ok(res)
116    }
117
118    /// Returns the lowercased text whose tokens must be indexed too, if any.
119    fn irregular_lowercased_text(&self, text: &str) -> Option<String> {
120        (!self.case_sensitive && has_irregular_lowercase(text)).then(|| text.to_lowercase())
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn test_analyze_text_hashes_matches_analyze_text() {
130        let text = "Hello, WORLD ship_Ship 清洁表面 ÄÖÜ straße İstanbul x";
131        for case_sensitive in [false, true] {
132            let analyzer = Analyzer::new(Box::new(ScriptTokenizer), case_sensitive);
133            let expected = analyzer
134                .analyze_text(text)
135                .unwrap()
136                .iter()
137                .map(|t| element_hash(t))
138                .collect::<Vec<_>>();
139            let hashes = analyzer
140                .analyze_text_hashes(text, &mut Vec::new())
141                .collect::<Vec<_>>();
142            assert_eq!(expected, hashes);
143        }
144    }
145
146    #[test]
147    fn test_analyzer() {
148        let analyzer = Analyzer::new(Box::new(ScriptTokenizer), false);
149        let tokens = analyzer.analyze_text("Hello, World_Two").unwrap();
150        assert_eq!(
151            tokens,
152            vec![b"hello".to_vec(), b"world".to_vec(), b"two".to_vec()]
153        );
154    }
155}
index/fulltext_index/tokenizer.rs

index/fulltext_index/
tokenizer.rs
