Skip to main content

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 crate::Bytes;
16use crate::bloom_filter::element_hash;
17use crate::fulltext_index::error::Result;
18
19lazy_static::lazy_static! {
20    static ref JIEBA: jieba_rs::Jieba = jieba_rs::Jieba::new();
21}
22
23/// A-Z, a-z, 0-9, and '_' are true
24const VALID_ASCII_TOKEN: [bool; 256] = [
25    false, false, false, false, false, false, false, false, false, false, false, false, false,
26    false, false, false, false, false, false, false, false, false, false, false, false, false,
27    false, false, false, false, false, false, false, false, false, false, false, false, false,
28    false, false, false, false, false, false, false, false, false, true, true, true, true, true,
29    true, true, true, true, true, false, false, false, false, false, false, false, true, true,
30    true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true,
31    true, true, true, true, true, true, true, true, false, false, false, false, true, false, true,
32    true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true,
33    true, true, true, true, true, true, true, true, true, false, false, false, false, false, false,
34    false, false, false, false, false, false, false, false, false, false, false, false, false,
35    false, false, false, false, false, false, false, false, false, false, false, false, false,
36    false, false, false, false, false, false, false, false, false, false, false, false, false,
37    false, false, false, false, false, false, false, false, false, false, false, false, false,
38    false, false, false, false, false, false, false, false, false, false, false, false, false,
39    false, false, false, false, false, false, false, false, false, false, false, false, false,
40    false, false, false, false, false, false, false, false, false, false, false, false, false,
41    false, false, false, false, false, false, false, false, false, false, false, false, false,
42    false, false, false, false, false, false, false, false, false, false, false, false, false,
43    false, false, false, false, false, false, false, false, false, false,
44];
45
46/// `Tokenizer` tokenizes a text into a list of tokens.
47pub trait Tokenizer: Send {
48    fn tokenize<'a>(&self, text: &'a str) -> Vec<&'a str>;
49}
50
51/// `EnglishTokenizer` tokenizes an English text.
52///
53/// It splits the text by non-alphabetic characters.
54#[derive(Debug, Default)]
55pub struct EnglishTokenizer;
56
57impl Tokenizer for EnglishTokenizer {
58    fn tokenize<'a>(&self, text: &'a str) -> Vec<&'a str> {
59        if text.is_ascii() {
60            let mut tokens = Vec::new();
61            let mut start = 0;
62            for (i, &byte) in text.as_bytes().iter().enumerate() {
63                if !VALID_ASCII_TOKEN[byte as usize] {
64                    if start < i {
65                        tokens.push(&text[start..i]);
66                    }
67                    start = i + 1;
68                }
69            }
70
71            if start < text.len() {
72                tokens.push(&text[start..]);
73            }
74
75            tokens
76        } else {
77            text.split(|c: char| !c.is_alphanumeric() && c != '_')
78                .filter(|s| !s.is_empty())
79                .collect()
80        }
81    }
82}
83
84/// `ChineseTokenizer` tokenizes a Chinese text.
85///
86/// It uses Jieba search-mode tokenization to improve recall for Chinese fulltext search.
87/// Enabling HMM also helps merge some unknown fragments into larger tokens, which can reduce
88/// token cardinality versus a fully fragmented output.
89#[derive(Debug, Default)]
90pub struct ChineseTokenizer;
91
92impl Tokenizer for ChineseTokenizer {
93    fn tokenize<'a>(&self, text: &'a str) -> Vec<&'a str> {
94        if text.is_ascii() {
95            EnglishTokenizer {}.tokenize(text)
96        } else {
97            // Search-mode tokenization emits finer-grained searchable terms, while HMM helps
98            // merge some unknown fragments and avoid excessive token fragmentation.
99            let mut tokens = JIEBA
100                .cut_for_search(text, true)
101                .into_iter()
102                .map(|token| token.word)
103                .filter(|token| is_indexable_token(token))
104                .collect::<Vec<_>>();
105
106            let english = EnglishTokenizer {};
107            tokens.extend(
108                english
109                    .tokenize(text)
110                    .into_iter()
111                    .filter(|token| is_ascii_underscore_token(token)),
112            );
113
114            tokens
115        }
116    }
117}
118
119fn is_indexable_token(token: &str) -> bool {
120    token.chars().any(|c| c.is_alphanumeric() || c == '_')
121}
122
123fn is_ascii_underscore_token(token: &str) -> bool {
124    token.is_ascii() && token.chars().any(|c| c == '_')
125}
126
127/// `Analyzer` analyzes a text into a list of tokens.
128///
129/// It uses a `Tokenizer` to tokenize the text and optionally lowercases the tokens.
130pub struct Analyzer {
131    tokenizer: Box<dyn Tokenizer>,
132    case_sensitive: bool,
133}
134
135impl Analyzer {
136    /// Creates a new `Analyzer` with the given `Tokenizer` and case sensitivity.
137    pub fn new(tokenizer: Box<dyn Tokenizer>, case_sensitive: bool) -> Self {
138        Self {
139            tokenizer,
140            case_sensitive,
141        }
142    }
143
144    /// Returns the bloom filter hash of each token in the given text.
145    ///
146    /// Equivalent to hashing every token returned by [`Analyzer::analyze_text`] with
147    /// [`element_hash`]. Only case-insensitive non-ASCII tokens allocate, in `to_lowercase`;
148    /// case-insensitive ASCII tokens are lowercased in `buf`.
149    pub fn analyze_text_hashes<'a>(
150        &self,
151        text: &'a str,
152        buf: &'a mut Vec<u8>,
153    ) -> impl Iterator<Item = u64> + use<'a> {
154        let case_sensitive = self.case_sensitive;
155        self.tokenizer.tokenize(text).into_iter().map(move |token| {
156            if case_sensitive {
157                element_hash(token.as_bytes())
158            } else if token.is_ascii() {
159                buf.clear();
160                buf.extend(token.bytes().map(|b| b.to_ascii_lowercase()));
161                element_hash(buf)
162            } else {
163                element_hash(token.to_lowercase().as_bytes())
164            }
165        })
166    }
167
168    /// Analyzes the given text into a list of tokens.
169    pub fn analyze_text(&self, text: &str) -> Result<Vec<Bytes>> {
170        let res = self
171            .tokenizer
172            .tokenize(text)
173            .iter()
174            .map(|s| {
175                if self.case_sensitive {
176                    s.as_bytes().to_vec()
177                } else {
178                    s.to_lowercase().as_bytes().to_vec()
179                }
180            })
181            .collect();
182        Ok(res)
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    #[test]
191    fn test_analyze_text_hashes_matches_analyze_text() {
192        let text = "Hello, WORLD ship_Ship 清洁表面 ÄÖÜ straße İstanbul x";
193        for (tokenizer, case_sensitive) in [
194            (Box::new(EnglishTokenizer) as Box<dyn Tokenizer>, false),
195            (Box::new(EnglishTokenizer), true),
196            (Box::new(ChineseTokenizer), false),
197        ] {
198            let analyzer = Analyzer::new(tokenizer, case_sensitive);
199            let expected = analyzer
200                .analyze_text(text)
201                .unwrap()
202                .iter()
203                .map(|t| element_hash(t))
204                .collect::<Vec<_>>();
205            let hashes = analyzer
206                .analyze_text_hashes(text, &mut Vec::new())
207                .collect::<Vec<_>>();
208            assert_eq!(expected, hashes);
209        }
210    }
211
212    #[test]
213    fn test_english_tokenizer() {
214        let tokenizer = EnglishTokenizer;
215        let text = "Hello, world!!! This is a----++   test012_345+67890 ship_ship ship__ship _ __ __IDENTIFIER__ _ship ship_";
216        let tokens = tokenizer.tokenize(text);
217        assert_eq!(
218            tokens,
219            vec![
220                "Hello",
221                "world",
222                "This",
223                "is",
224                "a",
225                "test012_345",
226                "67890",
227                "ship_ship",
228                "ship__ship",
229                "_",
230                "__",
231                "__IDENTIFIER__",
232                "_ship",
233                "ship_"
234            ]
235        );
236    }
237
238    #[test]
239    fn test_english_tokenizer_with_utf8() {
240        let tokenizer = EnglishTokenizer;
241        let text = "💸unfold the 纸巾😣and gently 清洁表😭面";
242        let tokens = tokenizer.tokenize(text);
243        assert_eq!(
244            tokens,
245            // Don't care what happens to non-ASCII characters.
246            // It's kind of a misconfiguration to use EnglishTokenizer on non-ASCII text.
247            vec!["unfold", "the", "纸巾", "and", "gently", "清洁表", "面"]
248        );
249    }
250
251    #[test]
252    fn test_chinese_tokenizer() {
253        let tokenizer = ChineseTokenizer;
254        let text = "我喜欢苹果";
255        let tokens = tokenizer.tokenize(text);
256        assert_eq!(tokens, vec!["我", "喜欢", "苹果"]);
257    }
258
259    #[test]
260    fn test_chinese_tokenizer_issue_7943_sample() {
261        let tokenizer = ChineseTokenizer;
262        let text = "[2026/04/09/ 13:56:11.031]2026-04-09 13:56:11.031 - [ trace_id=340a6a44b0bd8e37bb7697ss7da61ff0 span_id=085ff5ttf1e0a23b trace_flags=01] - [http-nio-8081-exec-16] INFO c.h.p.xx.web.service.impl.CCCXForwardKKKServiceImpl.pushout(188) - 登录手机号18888888888的动态key:829889AC8 ship_ship ship__ship _ __ __IDENTIFIER__ _ship ship_ EOF";
263        let tokens = tokenizer.tokenize(text);
264
265        assert_eq!(
266            tokens,
267            vec![
268                "2026",
269                "04",
270                "09",
271                "13",
272                "56",
273                "11.031",
274                "2026-04",
275                "09",
276                "13",
277                "56",
278                "11.031",
279                "trace",
280                "_",
281                "id",
282                "340a6a44b0bd8e37bb7697ss7da61ff0",
283                "span",
284                "_",
285                "id",
286                "085ff5ttf1e0a23b",
287                "trace",
288                "_",
289                "flags",
290                "01",
291                "http",
292                "nio-8081",
293                "exec-16",
294                "INFO",
295                "c",
296                "h",
297                "p",
298                "xx",
299                "web",
300                "service",
301                "impl",
302                "CCCXForwardKKKServiceImpl",
303                "pushout",
304                "188",
305                "登录",
306                "手机",
307                "手机号",
308                "18888888888",
309                "的",
310                "动态",
311                "key",
312                "829889AC8",
313                "ship",
314                "_",
315                "ship",
316                "ship",
317                "__",
318                "ship",
319                "_",
320                "__",
321                "__",
322                "IDENTIFIER",
323                "__",
324                "_",
325                "ship",
326                "ship",
327                "_",
328                "EOF",
329                "trace_id",
330                "span_id",
331                "trace_flags",
332                "ship_ship",
333                "ship__ship",
334                "_",
335                "__",
336                "__IDENTIFIER__",
337                "_ship",
338                "ship_"
339            ]
340        );
341    }
342
343    #[test]
344    fn test_chinese_tokenizer_keeps_ascii_underscore_compounds() {
345        let tokenizer = ChineseTokenizer;
346        let text = "trace_id=abc 登录手机号 dynamic_key=xyz";
347
348        let tokens = tokenizer.tokenize(text);
349
350        assert!(tokens.contains(&"trace_id"));
351        assert!(tokens.contains(&"dynamic_key"));
352        assert!(tokens.contains(&"登录"));
353        assert!(tokens.contains(&"手机号"));
354    }
355
356    #[test]
357    fn test_chinese_tokenizer_skips_non_ascii_underscore_tokens() {
358        let tokenizer = ChineseTokenizer;
359        let text = "登录_id trace_id 手机号_trace";
360
361        let tokens = tokenizer.tokenize(text);
362
363        assert_eq!(
364            tokens,
365            [
366                "登录",
367                "_",
368                "id",
369                "trace",
370                "_",
371                "id",
372                "手机",
373                "手机号",
374                "_",
375                "trace",
376                "trace_id"
377            ]
378        );
379    }
380
381    #[test]
382    fn test_chinese_tokenizer_aggressive_tokenization_probe() {
383        let tokenizer = ChineseTokenizer;
384        let text = "哈基米哦南北绿豆,噢马自立曼波。登录手机号。中国农业银行。装电视台,中国中央广播电视台。压不缩,笑不活。";
385
386        let default_tokens = tokenizer.tokenize(text);
387        let cut_hmm_false = JIEBA
388            .cut(text, false)
389            .into_iter()
390            .map(|token| token.word)
391            .collect::<Vec<_>>();
392        let cut_hmm_true = JIEBA
393            .cut(text, true)
394            .into_iter()
395            .map(|token| token.word)
396            .collect::<Vec<_>>();
397        let cut_for_search_hmm_false = JIEBA
398            .cut_for_search(text, false)
399            .into_iter()
400            .map(|token| token.word)
401            .collect::<Vec<_>>();
402        let cut_for_search_hmm_true = JIEBA
403            .cut_for_search(text, true)
404            .into_iter()
405            .map(|token| token.word)
406            .collect::<Vec<_>>();
407
408        assert_eq!(
409            default_tokens,
410            [
411                "哈基米",
412                "哦",
413                "南北",
414                "绿豆",
415                "噢",
416                "马",
417                "自立",
418                "曼波",
419                "登录",
420                "手机",
421                "手机号",
422                "中国",
423                "农业",
424                "银行",
425                "中国农业银行",
426                "装",
427                "电视",
428                "电视台",
429                "中国",
430                "中央",
431                "广播",
432                "电视",
433                "电视台",
434                "不缩",
435                "压不缩",
436                "笑",
437                "不活",
438            ]
439        );
440        assert_eq!(
441            cut_hmm_false,
442            [
443                "哈",
444                "基",
445                "米",
446                "哦",
447                "南北",
448                "绿豆",
449                ",",
450                "噢",
451                "马",
452                "自立",
453                "曼",
454                "波",
455                "。",
456                "登录",
457                "手机号",
458                "。",
459                "中国农业银行",
460                "。",
461                "装",
462                "电视台",
463                ",",
464                "中国",
465                "中央",
466                "广播",
467                "电视台",
468                "。",
469                "压",
470                "不",
471                "缩",
472                ",",
473                "笑",
474                "不",
475                "活",
476                "。"
477            ]
478        );
479        assert_eq!(
480            cut_hmm_true,
481            [
482                "哈基米",
483                "哦",
484                "南北",
485                "绿豆",
486                ",",
487                "噢",
488                "马",
489                "自立",
490                "曼波",
491                "。",
492                "登录",
493                "手机号",
494                "。",
495                "中国农业银行",
496                "。",
497                "装",
498                "电视台",
499                ",",
500                "中国",
501                "中央",
502                "广播",
503                "电视台",
504                "。",
505                "压不缩",
506                ",",
507                "笑",
508                "不活",
509                "。"
510            ]
511        );
512        assert_eq!(
513            cut_for_search_hmm_false,
514            [
515                "哈",
516                "基",
517                "米",
518                "哦",
519                "南北",
520                "绿豆",
521                ",",
522                "噢",
523                "马",
524                "自立",
525                "曼",
526                "波",
527                "。",
528                "登录",
529                "手机",
530                "手机号",
531                "。",
532                "中国",
533                "农业",
534                "银行",
535                "中国农业银行",
536                "。",
537                "装",
538                "电视",
539                "电视台",
540                ",",
541                "中国",
542                "中央",
543                "广播",
544                "电视",
545                "电视台",
546                "。",
547                "压",
548                "不",
549                "缩",
550                ",",
551                "笑",
552                "不",
553                "活",
554                "。"
555            ]
556        );
557
558        assert_eq!(
559            cut_for_search_hmm_true,
560            [
561                "哈基米",
562                "哦",
563                "南北",
564                "绿豆",
565                ",",
566                "噢",
567                "马",
568                "自立",
569                "曼波",
570                "。",
571                "登录",
572                "手机",
573                "手机号",
574                "。",
575                "中国",
576                "农业",
577                "银行",
578                "中国农业银行",
579                "。",
580                "装",
581                "电视",
582                "电视台",
583                ",",
584                "中国",
585                "中央",
586                "广播",
587                "电视",
588                "电视台",
589                "。",
590                "不缩",
591                "压不缩",
592                ",",
593                "笑",
594                "不活",
595                "。"
596            ]
597        );
598    }
599
600    #[test]
601    fn test_valid_ascii_token_lookup_table() {
602        // Test all ASCII values in a single loop
603        for c in 0u8..=255u8 {
604            let is_valid = VALID_ASCII_TOKEN[c as usize];
605            let should_be_valid = (c as char).is_ascii_alphanumeric() || c == b'_';
606
607            assert_eq!(
608                is_valid,
609                should_be_valid,
610                "Character '{}' (byte {}) validity mismatch: expected {}, got {}",
611                if c.is_ascii() && !c.is_ascii_control() {
612                    c as char
613                } else {
614                    '?'
615                },
616                c,
617                should_be_valid,
618                is_valid
619            );
620        }
621    }
622
623    #[test]
624    fn test_analyzer() {
625        let tokenizer = EnglishTokenizer;
626        let analyzer = Analyzer::new(Box::new(tokenizer), false);
627        let text = "Hello, world! This is a test.";
628        let tokens = analyzer.analyze_text(text).unwrap();
629        assert_eq!(
630            tokens,
631            vec![
632                b"hello".to_vec(),
633                b"world".to_vec(),
634                b"this".to_vec(),
635                b"is".to_vec(),
636                b"a".to_vec(),
637                b"test".to_vec()
638            ]
639        );
640    }
641}