lance_index/scalar/
inverted.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4pub mod builder;
5mod encoding;
6mod index;
7mod iter;
8pub mod json;
9pub mod parser;
10pub mod query;
11mod scorer;
12pub mod tokenizer;
13mod wand;
14
15use std::sync::Arc;
16
17use arrow_schema::{DataType, Field};
18use async_trait::async_trait;
19pub use builder::InvertedIndexBuilder;
20use datafusion::execution::SendableRecordBatchStream;
21pub use index::*;
22use lance_core::{Result, cache::LanceCache};
23pub use scorer::MemBM25Scorer;
24use tantivy::tokenizer::Language;
25pub use tokenizer::*;
26
27use lance_core::Error;
28
29use crate::pbold;
30use crate::progress::IndexBuildProgress;
31use crate::{
32    frag_reuse::FragReuseIndex,
33    scalar::{
34        CreatedIndex, ScalarIndex,
35        expression::{FtsQueryParser, ScalarQueryParser},
36        registry::{ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest},
37    },
38};
39
40use super::IndexStore;
41
42#[derive(Debug, Default)]
43pub struct InvertedIndexPlugin;
44
45impl InvertedIndexPlugin {
46    pub async fn train_inverted_index(
47        data: SendableRecordBatchStream,
48        index_store: &dyn IndexStore,
49        params: InvertedIndexParams,
50        fragment_ids: Option<Vec<u32>>,
51        progress: Arc<dyn IndexBuildProgress>,
52    ) -> Result<CreatedIndex> {
53        let fragment_mask = fragment_ids.as_ref().and_then(|frag_ids| {
54            if !frag_ids.is_empty() {
55                // Create a mask with fragment_id in high 32 bits for distributed indexing
56                // This mask is used to filter partitions belonging to specific fragments
57                // If multiple fragments processed, use first fragment_id <<32 as mask
58                Some((frag_ids[0] as u64) << 32)
59            } else {
60                None
61            }
62        });
63
64        let details = pbold::InvertedIndexDetails::try_from(&params)?;
65        let mut inverted_index =
66            InvertedIndexBuilder::new_with_fragment_mask(params, fragment_mask)
67                .with_progress(progress);
68        inverted_index.update(data, index_store, None).await?;
69        Ok(CreatedIndex {
70            index_details: prost_types::Any::from_msg(&details).unwrap(),
71            index_version: current_fts_format_version().index_version(),
72            files: Some(index_store.list_files_with_sizes().await?),
73        })
74    }
75
76    /// Return true if the query can be used to speed up contains_tokens queries
77    fn can_accelerate_queries(details: &pbold::InvertedIndexDetails) -> bool {
78        details.base_tokenizer == Some("simple".to_string())
79            && details.max_token_length.is_none()
80            && details.language == serde_json::to_string(&Language::English).unwrap()
81            && !details.stem
82    }
83}
84
85struct InvertedIndexTrainingRequest {
86    parameters: InvertedIndexParams,
87    criteria: TrainingCriteria,
88}
89
90impl InvertedIndexTrainingRequest {
91    pub fn new(parameters: InvertedIndexParams) -> Self {
92        Self {
93            parameters,
94            criteria: TrainingCriteria::new(TrainingOrdering::None).with_row_id(),
95        }
96    }
97}
98
99impl TrainingRequest for InvertedIndexTrainingRequest {
100    fn as_any(&self) -> &dyn std::any::Any {
101        self
102    }
103
104    fn criteria(&self) -> &TrainingCriteria {
105        &self.criteria
106    }
107}
108
109#[async_trait]
110impl ScalarIndexPlugin for InvertedIndexPlugin {
111    fn name(&self) -> &str {
112        "Inverted"
113    }
114
115    fn new_training_request(
116        &self,
117        params: &str,
118        field: &Field,
119    ) -> Result<Box<dyn TrainingRequest>> {
120        match field.data_type() {
121            DataType::Utf8 | DataType::LargeUtf8 | DataType::LargeBinary => (),
122            DataType::List(f) if matches!(f.data_type(), DataType::Utf8 | DataType::LargeUtf8) => (),
123            DataType::LargeList(f) if matches!(f.data_type(), DataType::Utf8 | DataType::LargeUtf8) => (),
124
125            _ => return Err(Error::invalid_input_source(format!(
126                "A inverted index can only be created on a Utf8 or LargeUtf8 field/list or LargeBinary field. Column has type {:?}",
127                field.data_type()
128            )
129                .into()))
130        }
131
132        let params = serde_json::from_str::<InvertedIndexParams>(params)?;
133        Ok(Box::new(InvertedIndexTrainingRequest::new(params)))
134    }
135
136    fn provides_exact_answer(&self) -> bool {
137        false
138    }
139
140    fn version(&self) -> u32 {
141        max_supported_fts_format_version().index_version()
142    }
143
144    fn new_query_parser(
145        &self,
146        index_name: String,
147        _index_details: &prost_types::Any,
148    ) -> Option<Box<dyn ScalarQueryParser>> {
149        let Ok(index_details) = _index_details.to_msg::<pbold::InvertedIndexDetails>() else {
150            return None;
151        };
152
153        if Self::can_accelerate_queries(&index_details) {
154            Some(Box::new(FtsQueryParser::new(index_name)))
155        } else {
156            None
157        }
158    }
159
160    /// Train a new index
161    ///
162    /// The provided data must fulfill all the criteria returned by `training_criteria`.
163    /// It is the caller's responsibility to ensure this.
164    ///
165    /// Returns index details that describe the index.  These details can potentially be
166    /// useful for planning (although this will currently require inside information on
167    /// the index type) and they will need to be provided when loading the index.
168    ///
169    /// It is the caller's responsibility to store these details somewhere.
170    async fn train_index(
171        &self,
172        data: SendableRecordBatchStream,
173        index_store: &dyn IndexStore,
174        request: Box<dyn TrainingRequest>,
175        fragment_ids: Option<Vec<u32>>,
176        progress: Arc<dyn IndexBuildProgress>,
177    ) -> Result<CreatedIndex> {
178        let request = (request as Box<dyn std::any::Any>)
179            .downcast::<InvertedIndexTrainingRequest>()
180            .map_err(|_| {
181                Error::invalid_input_source(
182                    "must provide training request created by new_training_request".into(),
183                )
184            })?;
185        Self::train_inverted_index(
186            data,
187            index_store,
188            request.parameters.clone(),
189            fragment_ids,
190            progress,
191        )
192        .await
193    }
194
195    /// Load an index from storage
196    ///
197    /// The index details should match the details that were returned when the index was
198    /// originally trained.
199    async fn load_index(
200        &self,
201        index_store: Arc<dyn IndexStore>,
202        _index_details: &prost_types::Any,
203        frag_reuse_index: Option<Arc<FragReuseIndex>>,
204        cache: &LanceCache,
205    ) -> Result<Arc<dyn ScalarIndex>> {
206        Ok(
207            InvertedIndex::load(index_store, frag_reuse_index, cache).await?
208                as Arc<dyn ScalarIndex>,
209        )
210    }
211
212    fn details_as_json(&self, details: &prost_types::Any) -> Result<serde_json::Value> {
213        let index_details = details.to_msg::<pbold::InvertedIndexDetails>()?;
214        let index_params = InvertedIndexParams::try_from(&index_details)?;
215        Ok(serde_json::json!(&index_params))
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn test_plugin_version_tracks_max_supported_format() {
225        let plugin = InvertedIndexPlugin;
226        assert_eq!(
227            plugin.version(),
228            max_supported_fts_format_version().index_version()
229        );
230    }
231}
lance_index/scalar/inverted.rs

lance_index/scalar/
inverted.rs
