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    /// Total memory limit in MiB for the build stage.
94    ///
95    /// This is split evenly across FTS workers at build time. By default Lance
96    /// uses roughly `num_cpus / 2` workers, unless `LANCE_FTS_NUM_SHARDS` is set.
97    /// If unset, each worker defaults to a 2 GiB build-time memory limit.
98    ///
99    /// This is a build-time only parameter and is not persisted with the index.
100    #[serde(
101        rename = "memory_limit",
102        skip_serializing,
103        default,
104        alias = "worker_memory_limit_mb"
105    )]
106    pub(crate) memory_limit_mb: Option<u64>,
107
108    /// Number of workers to use for FTS build.
109    ///
110    /// This is a build-time only parameter and is not persisted with the index.
111    /// By default Lance uses roughly `num_cpus / 2` workers.
112    /// The effective worker count is clamped to `[1, num_cpus - 2]`.
113    #[serde(rename = "num_workers", skip_serializing, default)]
114    pub(crate) num_workers: Option<usize>,
115}
116
117impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails {
118    type Error = Error;
119
120    fn try_from(params: &InvertedIndexParams) -> Result<Self> {
121        Ok(Self {
122            base_tokenizer: Some(params.base_tokenizer.clone()),
123            language: serde_json::to_string(&params.language)?,
124            with_position: params.with_position,
125            max_token_length: params.max_token_length.map(|l| l as u32),
126            lower_case: params.lower_case,
127            stem: params.stem,
128            remove_stop_words: params.remove_stop_words,
129            ascii_folding: params.ascii_folding,
130            min_ngram_length: params.min_ngram_length,
131            max_ngram_length: params.max_ngram_length,
132            prefix_only: params.prefix_only,
133        })
134    }
135}
136
137impl TryFrom<&pbold::InvertedIndexDetails> for InvertedIndexParams {
138    type Error = Error;
139
140    fn try_from(details: &pbold::InvertedIndexDetails) -> Result<Self> {
141        let defaults = Self::default();
142        Ok(Self {
143            lance_tokenizer: defaults.lance_tokenizer,
144            base_tokenizer: details
145                .base_tokenizer
146                .as_ref()
147                .cloned()
148                .unwrap_or(defaults.base_tokenizer),
149            language: serde_json::from_str(details.language.as_str())?,
150            with_position: details.with_position,
151            max_token_length: details.max_token_length.map(|l| l as usize),
152            lower_case: details.lower_case,
153            stem: details.stem,
154            remove_stop_words: details.remove_stop_words,
155            custom_stop_words: defaults.custom_stop_words,
156            ascii_folding: details.ascii_folding,
157            min_ngram_length: details.min_ngram_length,
158            max_ngram_length: details.max_ngram_length,
159            prefix_only: details.prefix_only,
160            memory_limit_mb: defaults.memory_limit_mb,
161            num_workers: defaults.num_workers,
162        })
163    }
164}
165
166fn bool_true() -> bool {
167    true
168}
169
170fn default_min_ngram_length() -> u32 {
171    3
172}
173
174fn default_max_ngram_length() -> u32 {
175    3
176}
177
178impl Default for InvertedIndexParams {
179    fn default() -> Self {
180        Self::new("simple".to_owned(), tantivy::tokenizer::Language::English)
181    }
182}
183
184impl InvertedIndexParams {
185    /// Create a new `InvertedIndexParams` with the given base tokenizer and language.
186    ///
187    /// The `base_tokenizer` can be one of the following:
188    /// - `simple`: splits tokens on whitespace and punctuation, default
189    /// - `whitespace`: splits tokens on whitespace
190    /// - `raw`: no tokenization
191    /// - `ngram`: N-Gram tokenizer
192    /// - `lindera/*`: Lindera tokenizer
193    /// - `jieba/*`: Jieba tokenizer
194    ///
195    /// The `language` is used for stemming and removing stop words,
196    /// this is not used for `lindera/*` and `jieba/*` tokenizers.
197    /// Default to `English`.
198    pub fn new(base_tokenizer: String, language: tantivy::tokenizer::Language) -> Self {
199        Self {
200            lance_tokenizer: None,
201            base_tokenizer,
202            language,
203            with_position: false,
204            max_token_length: Some(40),
205            lower_case: true,
206            stem: true,
207            remove_stop_words: true,
208            custom_stop_words: None,
209            ascii_folding: true,
210            min_ngram_length: default_min_ngram_length(),
211            max_ngram_length: default_max_ngram_length(),
212            prefix_only: false,
213            memory_limit_mb: None,
214            num_workers: None,
215        }
216    }
217
218    pub fn lance_tokenizer(mut self, lance_tokenizer: String) -> Self {
219        self.lance_tokenizer = Some(lance_tokenizer);
220        self
221    }
222
223    pub fn base_tokenizer(mut self, base_tokenizer: String) -> Self {
224        self.base_tokenizer = base_tokenizer;
225        self
226    }
227
228    pub fn language(mut self, language: &str) -> Result<Self> {
229        // need to convert to valid JSON string
230        let language = serde_json::from_str(format!("\"{}\"", language).as_str())?;
231        self.language = language;
232        Ok(self)
233    }
234
235    /// Set whether to store the position of the term in the document.
236    /// This can significantly increase the size of the index.
237    /// If false, only store the frequency of the term in the document.
238    /// This doesn't work with `ngram` tokenizer.
239    /// Default to `false`.
240    pub fn with_position(mut self, with_position: bool) -> Self {
241        self.with_position = with_position;
242        self
243    }
244
245    /// Get whether positions are stored in this index.
246    pub fn has_positions(&self) -> bool {
247        self.with_position
248    }
249
250    pub fn max_token_length(mut self, max_token_length: Option<usize>) -> Self {
251        self.max_token_length = max_token_length;
252        self
253    }
254
255    pub fn lower_case(mut self, lower_case: bool) -> Self {
256        self.lower_case = lower_case;
257        self
258    }
259
260    pub fn stem(mut self, stem: bool) -> Self {
261        self.stem = stem;
262        self
263    }
264
265    pub fn remove_stop_words(mut self, remove_stop_words: bool) -> Self {
266        self.remove_stop_words = remove_stop_words;
267        self
268    }
269
270    pub fn custom_stop_words(mut self, custom_stop_words: Option<Vec<String>>) -> Self {
271        self.custom_stop_words = custom_stop_words;
272        self
273    }
274
275    pub fn ascii_folding(mut self, ascii_folding: bool) -> Self {
276        self.ascii_folding = ascii_folding;
277        self
278    }
279
280    /// Set the minimum N-Gram length, only works when `base_tokenizer` is `ngram`.
281    /// Must be greater than 0 and not greater than `max_ngram_length`.
282    /// Default to 3.
283    pub fn ngram_min_length(mut self, min_length: u32) -> Self {
284        self.min_ngram_length = min_length;
285        self
286    }
287
288    /// Set the maximum N-Gram length, only works when `base_tokenizer` is `ngram`.
289    /// Must be greater than 0 and not less than `min_ngram_length`.
290    /// Default to 3.
291    pub fn ngram_max_length(mut self, max_length: u32) -> Self {
292        self.max_ngram_length = max_length;
293        self
294    }
295
296    /// Set whether only prefix N-Gram is generated, only works when `base_tokenizer` is `ngram`.
297    /// Default to `false`.
298    pub fn ngram_prefix_only(mut self, prefix_only: bool) -> Self {
299        self.prefix_only = prefix_only;
300        self
301    }
302
303    pub fn memory_limit_mb(mut self, memory_limit_mb: u64) -> Self {
304        self.memory_limit_mb = Some(memory_limit_mb);
305        self
306    }
307
308    /// Set the number of workers to use for this build.
309    ///
310    /// By default Lance uses roughly `num_cpus / 2` workers.
311    /// The effective worker count is clamped to `[1, num_cpus - 2]`.
312    pub fn num_workers(mut self, num_workers: usize) -> Self {
313        self.num_workers = Some(num_workers);
314        self
315    }
316
317    /// Serialize params for the build/training path, including build-only fields.
318    pub fn to_training_json(&self) -> serde_json::Result<serde_json::Value> {
319        let mut value = serde_json::to_value(self)?;
320        let object = value
321            .as_object_mut()
322            .expect("inverted index params should serialize to a JSON object");
323        if let Some(memory_limit_mb) = self.memory_limit_mb {
324            object.insert(
325                "memory_limit".to_string(),
326                serde_json::Value::from(memory_limit_mb),
327            );
328        }
329        if let Some(num_workers) = self.num_workers {
330            object.insert(
331                "num_workers".to_string(),
332                serde_json::Value::from(num_workers),
333            );
334        }
335        Ok(value)
336    }
337
338    pub fn build(&self) -> Result<Box<dyn LanceTokenizer>> {
339        let mut builder = self.build_base_tokenizer()?;
340        if let Some(max_token_length) = self.max_token_length {
341            builder = builder.filter_dynamic(tantivy::tokenizer::RemoveLongFilter::limit(
342                max_token_length,
343            ));
344        }
345        if self.lower_case {
346            builder = builder.filter_dynamic(tantivy::tokenizer::LowerCaser);
347        }
348        if self.stem {
349            builder = builder.filter_dynamic(tantivy::tokenizer::Stemmer::new(self.language));
350        }
351        if self.remove_stop_words {
352            let stop_word_filter = match &self.custom_stop_words {
353                Some(words) => tantivy::tokenizer::StopWordFilter::remove(words.iter().cloned()),
354                None => {
355                    tantivy::tokenizer::StopWordFilter::new(self.language).ok_or_else(|| {
356                        Error::invalid_input(format!(
357                            "removing stop words for language {:?} is not supported yet",
358                            self.language
359                        ))
360                    })?
361                }
362            };
363            builder = builder.filter_dynamic(stop_word_filter);
364        }
365        if self.ascii_folding {
366            builder = builder.filter_dynamic(tantivy::tokenizer::AsciiFoldingFilter);
367        }
368        let tokenizer = builder.build();
369
370        match self.lance_tokenizer {
371            Some(ref t) if t == "text" => Ok(Box::new(TextTokenizer::new(tokenizer))),
372            Some(ref t) if t == "json" => Ok(Box::new(JsonTokenizer::new(tokenizer))),
373            None => Ok(Box::new(TextTokenizer::new(tokenizer))),
374            _ => Err(Error::invalid_input(format!(
375                "unknown lance tokenizer {}",
376                self.lance_tokenizer.as_ref().unwrap()
377            ))),
378        }
379    }
380
381    fn build_base_tokenizer(&self) -> Result<tantivy::tokenizer::TextAnalyzerBuilder> {
382        match self.base_tokenizer.as_str() {
383            "simple" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
384                tantivy::tokenizer::SimpleTokenizer::default(),
385            )
386            .dynamic()),
387            "whitespace" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
388                tantivy::tokenizer::WhitespaceTokenizer::default(),
389            )
390            .dynamic()),
391            "raw" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
392                tantivy::tokenizer::RawTokenizer::default(),
393            )
394            .dynamic()),
395            "ngram" => Ok(tantivy::tokenizer::TextAnalyzer::builder(
396                tantivy::tokenizer::NgramTokenizer::new(
397                    self.min_ngram_length as usize,
398                    self.max_ngram_length as usize,
399                    self.prefix_only,
400                )
401                .map_err(|e| Error::invalid_input(e.to_string()))?,
402            )
403            .dynamic()),
404            #[cfg(feature = "tokenizer-lindera")]
405            s if s.starts_with("lindera/") => {
406                let Some(home) = language_model_home() else {
407                    return Err(Error::invalid_input(format!(
408                        "unknown base tokenizer {}",
409                        self.base_tokenizer
410                    )));
411                };
412                lindera::LinderaBuilder::load(&home.join(s))?.build()
413            }
414            #[cfg(feature = "tokenizer-jieba")]
415            s if s.starts_with("jieba/") || s == "jieba" => {
416                let s = if s == "jieba" { "jieba/default" } else { s };
417                let Some(home) = language_model_home() else {
418                    return Err(Error::invalid_input(format!(
419                        "unknown base tokenizer {}",
420                        self.base_tokenizer
421                    )));
422                };
423                jieba::JiebaBuilder::load(&home.join(s))?.build()
424            }
425            _ => Err(Error::invalid_input(format!(
426                "unknown base tokenizer {}",
427                self.base_tokenizer
428            ))),
429        }
430    }
431}
432
433pub const LANCE_LANGUAGE_MODEL_HOME_ENV_KEY: &str = "LANCE_LANGUAGE_MODEL_HOME";
434
435pub const LANCE_LANGUAGE_MODEL_DEFAULT_DIRECTORY: &str = "lance/language_models";
436
437pub fn language_model_home() -> Option<PathBuf> {
438    match env::var(LANCE_LANGUAGE_MODEL_HOME_ENV_KEY) {
439        Ok(p) => Some(PathBuf::from(p)),
440        Err(_) => dirs::data_local_dir().map(|p| p.join(LANCE_LANGUAGE_MODEL_DEFAULT_DIRECTORY)),
441    }
442}
443
444#[cfg(test)]
445mod tests {
446    use super::InvertedIndexParams;
447
448    #[test]
449    fn test_build_only_fields_are_not_serialized() {
450        let params = InvertedIndexParams::default()
451            .memory_limit_mb(4096)
452            .num_workers(7);
453        let json = serde_json::to_value(&params).unwrap();
454        assert!(json.get("memory_limit").is_none());
455        assert!(json.get("num_workers").is_none());
456    }
457
458    #[test]
459    fn test_memory_limit_serde_accepts_legacy_worker_field_name() {
460        let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap();
461        let obj = json.as_object_mut().unwrap();
462        obj.remove("memory_limit");
463        obj.insert(
464            "worker_memory_limit_mb".to_string(),
465            serde_json::Value::from(2048),
466        );
467        let params: InvertedIndexParams = serde_json::from_value(json).unwrap();
468        assert_eq!(params.memory_limit_mb, Some(2048));
469    }
470
471    #[test]
472    fn test_build_only_fields_deserialize_from_public_names() {
473        let mut json = serde_json::to_value(InvertedIndexParams::default()).unwrap();
474        let obj = json.as_object_mut().unwrap();
475        obj.insert("memory_limit".to_string(), serde_json::Value::from(4096));
476        obj.insert("num_workers".to_string(), serde_json::Value::from(3));
477
478        let params: InvertedIndexParams = serde_json::from_value(json).unwrap();
479        assert_eq!(params.memory_limit_mb, Some(4096));
480        assert_eq!(params.num_workers, Some(3));
481    }
482
483    #[test]
484    fn test_training_json_serializes_build_only_fields() {
485        let params = InvertedIndexParams::default()
486            .memory_limit_mb(4096)
487            .num_workers(3);
488        let json = params.to_training_json().unwrap();
489        assert_eq!(
490            json.get("memory_limit"),
491            Some(&serde_json::Value::from(4096))
492        );
493        assert_eq!(json.get("num_workers"), Some(&serde_json::Value::from(3)));
494    }
495}
lance_index/scalar/inverted/tokenizer.rs

lance_index/scalar/inverted/
tokenizer.rs
