Skip to main content

uqa_analysis/
tokenizer.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Tokenizers for the analysis pipeline. An [`Analyzer`] owns exactly one
8//! tokenizer.
9//!
10//! [`Analyzer`]: crate::analyzer::Analyzer
11
12use std::sync::OnceLock;
13
14use regex::Regex;
15use serde::{Deserialize, Serialize};
16
17use crate::error::{AnalysisError, AnalysisResult};
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20#[serde(tag = "type", rename_all = "snake_case")]
21pub enum Tokenizer {
22    Whitespace,
23    Standard,
24    Letter,
25    NGram { min_gram: usize, max_gram: usize },
26    Pattern { pattern: String },
27    Keyword,
28}
29
30impl Tokenizer {
31    /// Validate configuration without tokenizing input. This is used when an
32    /// analyzer is registered, while [`Self::tokenize`] repeats the checks so
33    /// deserialized legacy values can never bypass them.
34    pub fn validate(&self) -> AnalysisResult<()> {
35        match self {
36            Tokenizer::NGram { min_gram, max_gram } => {
37                validate_gram_bounds("n-gram tokenizer", *min_gram, *max_gram)
38            }
39            Tokenizer::Pattern { pattern } => {
40                Regex::new(pattern)
41                    .map(|_| ())
42                    .map_err(|source| AnalysisError::InvalidRegex {
43                        component: "pattern tokenizer",
44                        pattern: pattern.clone(),
45                        source,
46                    })
47            }
48            _ => Ok(()),
49        }
50    }
51
52    pub fn tokenize(&self, text: &str) -> AnalysisResult<Vec<String>> {
53        let tokens = match self {
54            Tokenizer::Whitespace => text.split_whitespace().map(str::to_owned).collect(),
55            Tokenizer::Standard => standard_word_re()?
56                .find_iter(text)
57                .map(|m| m.as_str().to_owned())
58                .collect(),
59            Tokenizer::Letter => letter_re()?
60                .find_iter(text)
61                .map(|m| m.as_str().to_owned())
62                .collect(),
63            Tokenizer::NGram { min_gram, max_gram } => {
64                validate_gram_bounds("n-gram tokenizer", *min_gram, *max_gram)?;
65                let mut out = Vec::new();
66                for word in text.split_whitespace() {
67                    let chars: Vec<char> = word.chars().collect();
68                    for n in *min_gram..=*max_gram {
69                        if chars.len() < n {
70                            continue;
71                        }
72                        for i in 0..=(chars.len() - n) {
73                            out.push(chars[i..i + n].iter().collect());
74                        }
75                    }
76                }
77                out
78            }
79            Tokenizer::Pattern { pattern } => {
80                let re = Regex::new(pattern).map_err(|source| AnalysisError::InvalidRegex {
81                    component: "pattern tokenizer",
82                    pattern: pattern.clone(),
83                    source,
84                })?;
85                re.split(text)
86                    .filter(|s| !s.is_empty())
87                    .map(str::to_owned)
88                    .collect()
89            }
90            Tokenizer::Keyword => {
91                if text.is_empty() {
92                    Vec::new()
93                } else {
94                    vec![text.to_owned()]
95                }
96            }
97        };
98        Ok(tokens)
99    }
100}
101
102fn validate_gram_bounds(
103    component: &'static str,
104    min_gram: usize,
105    max_gram: usize,
106) -> AnalysisResult<()> {
107    if min_gram == 0 || max_gram < min_gram {
108        return Err(AnalysisError::InvalidGramBounds {
109            component,
110            min_gram,
111            max_gram,
112        });
113    }
114    Ok(())
115}
116
117fn standard_word_re() -> AnalysisResult<&'static Regex> {
118    static RE: OnceLock<Result<Regex, String>> = OnceLock::new();
119    RE.get_or_init(|| Regex::new(r"\w+").map_err(|error| error.to_string()))
120        .as_ref()
121        .map_err(|message| AnalysisError::BuiltInRegex {
122            component: "standard tokenizer",
123            message: message.clone(),
124        })
125}
126
127fn letter_re() -> AnalysisResult<&'static Regex> {
128    static RE: OnceLock<Result<Regex, String>> = OnceLock::new();
129    RE.get_or_init(|| Regex::new(r"[a-zA-Z]+").map_err(|error| error.to_string()))
130        .as_ref()
131        .map_err(|message| AnalysisError::BuiltInRegex {
132            component: "letter tokenizer",
133            message: message.clone(),
134        })
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn whitespace_splits_on_whitespace() {
143        let t = Tokenizer::Whitespace;
144        assert_eq!(
145            t.tokenize("hello  world\n  rust").unwrap(),
146            vec!["hello", "world", "rust"]
147        );
148    }
149
150    #[test]
151    fn standard_extracts_unicode_words() {
152        let t = Tokenizer::Standard;
153        assert_eq!(
154            t.tokenize("Rust 2024! Carácter.").unwrap(),
155            vec!["Rust", "2024", "Carácter"]
156        );
157    }
158
159    #[test]
160    fn letter_extracts_ascii_letters_only() {
161        let t = Tokenizer::Letter;
162        assert_eq!(t.tokenize("abc123 xyz").unwrap(), vec!["abc", "xyz"]);
163    }
164
165    #[test]
166    fn ngram_emits_substrings_per_word() {
167        let t = Tokenizer::NGram {
168            min_gram: 2,
169            max_gram: 3,
170        };
171        // "ab" word: 2-grams [ab]
172        // "abc" word: 2-grams [ab, bc], 3-grams [abc]
173        assert_eq!(t.tokenize("ab abc").unwrap(), vec!["ab", "ab", "bc", "abc"]);
174    }
175
176    #[test]
177    fn pattern_splits_on_regex() {
178        let t = Tokenizer::Pattern {
179            pattern: r"\W+".to_string(),
180        };
181        assert_eq!(t.tokenize("hello, world!").unwrap(), vec!["hello", "world"]);
182    }
183
184    #[test]
185    fn keyword_emits_whole_input() {
186        let t = Tokenizer::Keyword;
187        assert_eq!(t.tokenize("a b c").unwrap(), vec!["a b c"]);
188        assert!(t.tokenize("").unwrap().is_empty());
189    }
190
191    #[test]
192    fn round_trips_via_serde_json() {
193        let t = Tokenizer::NGram {
194            min_gram: 2,
195            max_gram: 4,
196        };
197        let s = serde_json::to_string(&t).unwrap();
198        let back: Tokenizer = serde_json::from_str(&s).unwrap();
199        assert_eq!(
200            back.tokenize("foobar").unwrap(),
201            t.tokenize("foobar").unwrap()
202        );
203    }
204}