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;
9mod merger;
10pub mod parser;
11pub mod query;
12mod scorer;
13pub mod tokenizer;
14mod wand;
15
16use std::sync::Arc;
17
18use arrow_schema::{DataType, Field};
19use async_trait::async_trait;
20pub use builder::InvertedIndexBuilder;
21use datafusion::execution::SendableRecordBatchStream;
22pub use index::*;
23use lance_core::{Result, cache::LanceCache};
24pub use scorer::MemBM25Scorer;
25use tantivy::tokenizer::Language;
26pub use tokenizer::*;
27
28use lance_core::Error;
29
30use crate::pbold;
31use crate::progress::IndexBuildProgress;
32use crate::{
33    frag_reuse::FragReuseIndex,
34    scalar::{
35        CreatedIndex, ScalarIndex,
36        expression::{FtsQueryParser, ScalarQueryParser},
37        registry::{ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest},
38    },
39};
40
41use super::IndexStore;
42
43#[derive(Debug, Default)]
44pub struct InvertedIndexPlugin;
45
46impl InvertedIndexPlugin {
47    pub async fn train_inverted_index(
48        data: SendableRecordBatchStream,
49        index_store: &dyn IndexStore,
50        params: InvertedIndexParams,
51        fragment_ids: Option<Vec<u32>>,
52        progress: Arc<dyn IndexBuildProgress>,
53    ) -> Result<CreatedIndex> {
54        let fragment_mask = fragment_ids.as_ref().and_then(|frag_ids| {
55            if !frag_ids.is_empty() {
56                // Create a mask with fragment_id in high 32 bits for distributed indexing
57                // This mask is used to filter partitions belonging to specific fragments
58                // If multiple fragments processed, use first fragment_id <<32 as mask
59                Some((frag_ids[0] as u64) << 32)
60            } else {
61                None
62            }
63        });
64
65        let details = pbold::InvertedIndexDetails::try_from(&params)?;
66        let mut inverted_index =
67            InvertedIndexBuilder::new_with_fragment_mask(params, fragment_mask)
68                .with_progress(progress);
69        inverted_index.update(data, index_store).await?;
70        Ok(CreatedIndex {
71            index_details: prost_types::Any::from_msg(&details).unwrap(),
72            index_version: INVERTED_INDEX_VERSION,
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        INVERTED_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}
lance_index/scalar/inverted.rs

lance_index/scalar/
inverted.rs
