lance_index/scalar/inverted/
tokenizer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use lance_core::{Error, Result};
5use serde::{Deserialize, Serialize};
6use std::{env, path::PathBuf};
7
8#[cfg(feature = "tokenizer-jieba")]
9mod jieba;
10
11pub mod lance_tokenizer;
12#[cfg(feature = "tokenizer-lindera")]
13mod lindera;
14
15#[cfg(feature = "tokenizer-jieba")]
16use jieba::JiebaTokenizerBuilder;
17
18#[cfg(feature = "tokenizer-lindera")]
19use lindera::LinderaTokenizerBuilder;
20
21use crate::pbold;
22use crate::scalar::inverted::tokenizer::lance_tokenizer::{
23    JsonTokenizer, LanceTokenizer, TextTokenizer,
24};
25
26/// Tokenizer configs
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
28pub struct InvertedIndexParams {
29    /// lance tokenizer takes care of different data types, such as text, json, etc.
30    /// - 'text': parsing input documents into tokens
31    /// - 'json': parsing input json string into tokens
32    /// - none: auto type inference
33    pub(crate) lance_tokenizer: Option<String>,
34    /// base tokenizer:
35    /// - `simple`: splits tokens on whitespace and punctuation
36    /// - `whitespace`: splits tokens on whitespace
37    /// - `raw`: no tokenization
38    /// - `lindera/*`: Lindera tokenizer
39    /// - `jieba/*`: Jieba tokenizer
40    ///
41    /// `simple` is recommended for most cases and the default value
42    pub(crate) base_tokenizer: String,
43
44    /// language for stemming and stop words
45    /// this is only used when `stem` or `remove_stop_words` is true
46    pub(crate) language: tantivy::tokenizer::Language,
47
48    /// If true, store the position of the term in the document
49    /// This can significantly increase the size of the index
50    /// If false, only store the frequency of the term in the document
51    /// Default is false
52    #[serde(default)]
53    pub(crate) with_position: bool,
54
55    /// maximum token length
56    /// - `None`: no limit
57    /// - `Some(n)`: remove tokens longer than `n`
58    pub(crate) max_token_length: Option<usize>,
59
60    /// whether lower case tokens
61    #[serde(default = "bool_true")]
62    pub(crate) lower_case: bool,
63
64    /// whether apply stemming
65    #[serde(default = "bool_true")]
66    pub(crate) stem: bool,
67
68    /// whether remove stop words
69    #[serde(default = "bool_true")]
70    pub(crate) remove_stop_words: bool,
71
72    /// use customized stop words.
73    /// - `None`: use built-in stop words based on language
74    /// - `Some(words)`: use customized stop words
75    pub(crate) custom_stop_words: Option<Vec<String>>,
76
77    /// ascii folding
78    #[serde(default = "bool_true")]
79    pub(crate) ascii_folding: bool,
80
81    /// min ngram length
82    #[serde(default = "default_min_ngram_length")]
83    pub(crate) min_ngram_length: u32,
84
85    /// max ngram length
86    #[serde(default = "default_max_ngram_length")]
87    pub(crate) max_ngram_length: u32,
88
89    /// whether prefix only
90    #[serde(default)]
91    pub(crate) prefix_only: bool,
92
93    /// If true, skip the partition merge stage after indexing.
94    /// This can be useful for distributed indexing where merge is handled separately.
95    #[serde(default)]
96    pub(crate) skip_merge: bool,
97}
98
99impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails {
100    type Error = Error;
101
102    fn try_from(params: &InvertedIndexParams) -> Result<Self> {
103        Ok(Self {
104            base_tokenizer: Some(params.base_tokenizer.clone()),
105            language: serde_json::to_string(&params.language)?,
106            with_position: params.with_position,
107            max_token_length: params.max_token_length.map(|l| l as u32),
108            lower_case: params.lower_case,
109            stem: params.stem,
110            remove_stop_words: params.remove_stop_words,
111            ascii_folding: params.ascii_folding,
112            min_ngram_length: params.min_ngram_length,
113            max_ngram_length: params.max_ngram_length,
114            prefix_only: params.prefix_only,
115        })
116    }
117}
118
119impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams {
120    type Error = Error;
121
122    fn try_from(details: &pbold::InvertedIndexDetails) -> Result<Self> {
123        let defaults = Self::default();
124        Ok(Self {
125            lance_tokenizer: defaults.lance_tokenizer,
126            base_tokenizer: details
127                .base_tokenizer
128                .as_ref()
129                .cloned()
130                .unwrap_or(defaults.base_tokenizer),
131            language: serde_json::from_str(details.language.as_str())?,
132            with_position: details.with_position,
133            max_token_length: details.max_token_length.map(|l| l as usize),
134            lower_case: details.lower_case,
135            stem: details.stem,
136            remove_stop_words: details.remove_stop_words,
137            custom_stop_words: defaults.custom_stop_words,
138            ascii_folding: details.ascii_folding,
139            min_ngram_length: details.min_ngram_length,
140            max_ngram_length: details.max_ngram_length,
141            prefix_only: details.prefix_only,
142            skip_merge: defaults.skip_merge,
143        })
144    }
145}
146
147fn bool_true() -> bool {
148    true
149}
150
151fn default_min_ngram_length() -> u32 {
152    3
153}
154
155fn default_max_ngram_length() -> u32 {
156    3
157}
158
159impl Default for InvertedIndexParams {
160    fn default() -> Self {
161        Self::new("simple".to_owned(), tantivy::tokenizer::Language::English)
162    }
163}
164
165impl InvertedIndexParams {
166    /// Create a new `InvertedIndexParams` with the given base tokenizer and language.
167    ///
168    /// The `base_tokenizer` can be one of the following:
169    /// - `simple`: splits tokens on whitespace and punctuation, default
170    /// - `whitespace`: splits tokens on whitespace
171    /// - `raw`: no tokenization
172    /// - `ngram`: N-Gram tokenizer
173    /// - `lindera/*`: Lindera tokenizer
174    /// - `jieba/*`: Jieba tokenizer
175    ///
176    /// The `language` is used for stemming and removing stop words,
177    /// this is not used for `lindera/*` and `jieba/*` tokenizers.
178    /// Default to `English`.
179    pub fn new(base_tokenizer: String, language: tantivy::tokenizer::Language) -> Self {
180        Self {
181            lance_tokenizer: None,
182            base_tokenizer,
183            language,
184            with_position: false,
185            max_token_length: Some(40),
186            lower_case: true,
187            stem: true,
188            remove_stop_words: true,
189            custom_stop_words: None,
190            ascii_folding: true,
191            min_ngram_length: default_min_ngram_length(),
192            max_ngram_length: default_max_ngram_length(),
193            prefix_only: false,
194            skip_merge: false,
195        }
196    }
197
198    pub fn lance_tokenizer(mut self, lance_tokenizer: String) -> Self {
199        self.lance_tokenizer = Some(lance_tokenizer);
200        self
201    }
202
203    pub fn base_tokenizer(mut self, base_tokenizer: String) -> Self {
204        self.base_tokenizer = base_tokenizer;
205        self
206    }
207
208    pub fn language(mut self, language: &str) -> Result<Self> {
209        // need to convert to valid JSON string
210        let language = serde_json::from_str(format!("\"{}\"", language).as_str())?;
211        self.language = language;
212        Ok(self)
213    }
214
215    /// Set whether to store the position of the term in the document.
216    /// This can significantly increase the size of the index.
217    /// If false, only store the frequency of the term in the document.
218    /// This doesn't work with `ngram` tokenizer.
219    /// Default to `false`.
220    pub fn with_position(mut self, with_position: bool) -> Self {
221        self.with_position = with_position;
222        self
223    }
224
225    /// Get whether positions are stored in this index.
226    pub fn has_positions(&self) -> bool {
227        self.with_position
228    }
229
230    pub fn max_token_length(mut self, max_token_length: Option<usize>) -> Self {
231        self.max_token_length = max_token_length;
232        self
233    }
234
235    pub fn lower_case(mut self, lower_case: bool) -> Self {
236        self.lower_case = lower_case;
237        self
238    }
239
240    pub fn stem(mut self, stem: bool) -> Self {
241        self.stem = stem;
242        self
243    }
244
245    pub fn remove_stop_words(mut self, remove_stop_words: bool) -> Self {
246        self.remove_stop_words = remove_stop_words;
247        self
248    }
249
250    pub fn custom_stop_words(mut self, custom_stop_words: Option<Vec<String>>) -> Self {
251        self.custom_stop_words = custom_stop_words;
252        self
253    }
254
255    pub fn ascii_folding(mut self, ascii_folding: bool) -> Self {
256        self.ascii_folding = ascii_folding;
257        self
258    }
259
260    /// Set the minimum N-Gram length, only works when `base_tokenizer` is `ngram`.
261    /// Must be greater than 0 and not greater than `max_ngram_length`.
262    /// Default to 3.
263    pub fn ngram_min_length(mut self, min_length: u32) -> Self {
264        self.min_ngram_length = min_length;
265        self
266    }
267
268    /// Set the maximum N-Gram length, only works when `base_tokenizer` is `ngram`.
269    /// Must be greater than 0 and not less than `min_ngram_length`.
270    /// Default to 3.
271    pub fn ngram_max_length(mut self, max_length: u32) -> Self {
272        self.max_ngram_length = max_length;
273        self
274    }
275
276    /// Set whether only prefix N-Gram is generated, only works when `base_tokenizer` is `ngram`.
277    /// Default to `false`.
278    pub fn ngram_prefix_only(mut self, prefix_only: bool) -> Self {
279        self.prefix_only = prefix_only;
280        self
281    }
282
283    /// Skip merging partitions after indexing.
284    pub fn skip_merge(mut self, skip_merge: bool) -> Self {
285        self.skip_merge = skip_merge;
286        self
287    }
288
289    pub fn build(&self) -> Result<Box<dyn LanceTokenizer>> {
290        let mut builder = self.build_base_tokenizer()?;
291        if let Some(max_token_length) = self.max_token_length {
292            builder = builder.filter_dynamic(tantivy::tokenizer::RemoveLongFilter::limit(
293                max_token_length,
294            ));
295        }
296        if self.lower_case {
297            builder = builder.filter_dynamic(tantivy::tokenizer::LowerCaser);
298        }
299        if self.stem {
300            builder = builder.filter_dynamic(tantivy::tokenizer::Stemmer::new(self.language));
301        }
302        if self.remove_stop_words {
303            let stop_word_filter = match &self.custom_stop_words {
304                Some(words) => tantivy::tokenizer::StopWordFilter::remove(words.iter().cloned()),
305                None => {
306                    tantivy::tokenizer::StopWordFilter::new(self.language).ok_or_else(|| {
307                        Error::invalid_input(format!(
308                            "removing stop words for language {:?} is not supported yet",
309                            self.language
310                        ))
311                    })?
312                }
313            };
314            builder = builder.filter_dynamic(stop_word_filter);
315        }
316        if self.ascii_folding {
317            builder = builder.filter_dynamic(tantivy::tokenizer::AsciiFoldingFilter);
318        }
319        let tokenizer = builder.build();
320
321        match self.lance_tokenizer {
322            Some(ref t) if t == "text" => Ok(Box::new(TextTokenizer::new(tokenizer))),
323            Some(ref t) if t == "json" => Ok(Box::new(JsonTokenizer::new(tokenizer))),
324            None => Ok(Box::new(TextTokenizer::new(tokenizer))),
325            _ => Err(Error::invalid_input(format!(
326                "unknown lance tokenizer {}",
327                self.lance_tokenizer.as_ref().unwrap()
328            ))),
329        }
330    }
331
332    fn build_base_tokenizer(&self) -> Result<tantivy::tokenizer::TextAnalyzerBuilder> {
333        match self.base_tokenizer.as_str() {
334            "simple" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
335                tantivy::tokenizer::SimpleTokenizer::default(),
336            )
337            .dynamic()),
338            "whitespace" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
339                tantivy::tokenizer::WhitespaceTokenizer::default(),
340            )
341            .dynamic()),
342            "raw" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
343                tantivy::tokenizer::RawTokenizer::default(),
344            )
345            .dynamic()),
346            "ngram" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
347                tantivy::tokenizer::NgramTokenizer::new(
348                    self.min_ngram_length as usize,
349                    self.max_ngram_length as usize,
350                    self.prefix_only,
351                )
352                .map_err(|e| Error::invalid_input(e.to_string()))?,
353            )
354            .dynamic()),
355            #[cfg(feature = "tokenizer-lindera")]
356            s if s.starts_with("lindera/") => {
357                let Some(home) = language_model_home() else {
358                    return Err(Error::invalid_input(format!(
359                        "unknown base tokenizer {}",
360                        self.base_tokenizer
361                    )));
362                };
363                lindera::LinderaBuilder::load(&home.join(s))?.build()
364            }
365            #[cfg(feature = "tokenizer-jieba")]
366            s if s.starts_with("jieba/") || s == "jieba" => {
367                let s = if s == "jieba" { "jieba/default" } else { s };
368                let Some(home) = language_model_home() else {
369                    return Err(Error::invalid_input(format!(
370                        "unknown base tokenizer {}",
371                        self.base_tokenizer
372                    )));
373                };
374                jieba::JiebaBuilder::load(&home.join(s))?.build()
375            }
376            _ => Err(Error::invalid_input(format!(
377                "unknown base tokenizer {}",
378                self.base_tokenizer
379            ))),
380        }
381    }
382}
383
384pub const LANCE_LANGUAGE_MODEL_HOME_ENV_KEY: &str = "LANCE_LANGUAGE_MODEL_HOME";
385
386pub const LANCE_LANGUAGE_MODEL_DEFAULT_DIRECTORY: &str = "lance/language_models";
387
388pub fn language_model_home() -> Option<PathBuf> {
389    match env::var(LANCE_LANGUAGE_MODEL_HOME_ENV_KEY) {
390        Ok(p) => Some(PathBuf::from(p)),
391        Err(_) => dirs::data_local_dir().map(|p| p.join(LANCE_LANGUAGE_MODEL_DEFAULT_DIRECTORY)),
392    }
393}
lance_index/scalar/inverted/tokenizer.rs

lance_index/scalar/inverted/
tokenizer.rs
