lance_index/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Lance secondary index library
5//!
6//! <section class="warning">
7//! This is internal crate used by <a href="https://github.com/lance-format/lance">the lance project</a>.
8//! <br/>
9//! API stability is not guaranteed.
10//! </section>
11
12use std::{any::Any, sync::Arc};
13
14use crate::frag_reuse::FRAG_REUSE_INDEX_NAME;
15use crate::mem_wal::MEM_WAL_INDEX_NAME;
16use async_trait::async_trait;
17use deepsize::DeepSizeOf;
18use lance_core::{Error, Result};
19use roaring::RoaringBitmap;
20use serde::{Deserialize, Serialize};
21use snafu::location;
22use std::convert::TryFrom;
23
24pub mod frag_reuse;
25pub mod mem_wal;
26pub mod metrics;
27pub mod optimize;
28pub mod prefilter;
29pub mod registry;
30pub mod scalar;
31pub mod traits;
32pub mod vector;
33
34pub use crate::traits::*;
35
36pub const INDEX_FILE_NAME: &str = "index.idx";
37/// The name of the auxiliary index file.
38///
39/// This file is used to store additional information about the index, to improve performance.
40/// - For 'IVF_HNSW' index, it stores the partitioned PQ Storage.
41pub const INDEX_AUXILIARY_FILE_NAME: &str = "auxiliary.idx";
42pub const INDEX_METADATA_SCHEMA_KEY: &str = "lance:index";
43
44// Currently all vector indexes are version 1
45pub const VECTOR_INDEX_VERSION: u32 = 1;
46
47/// The factor of threshold to trigger split / join for vector index.
48///
49/// If the number of rows in the single partition is greater than `MAX_PARTITION_SIZE_FACTOR * target_partition_size`,
50/// the partition will be split.
51/// If the number of rows in the single partition is less than `MIN_PARTITION_SIZE_PERCENT *target_partition_size / 100`,
52/// the partition will be joined.
53pub const MAX_PARTITION_SIZE_FACTOR: usize = 4;
54pub const MIN_PARTITION_SIZE_PERCENT: usize = 25;
55
56pub mod pb {
57    #![allow(clippy::use_self)]
58    include!(concat!(env!("OUT_DIR"), "/lance.index.pb.rs"));
59}
60
61pub mod pbold {
62    #![allow(clippy::use_self)]
63    include!(concat!(env!("OUT_DIR"), "/lance.table.rs"));
64}
65
66/// Generic methods common across all types of secondary indices
67///
68#[async_trait]
69pub trait Index: Send + Sync + DeepSizeOf {
70    /// Cast to [Any].
71    fn as_any(&self) -> &dyn Any;
72
73    /// Cast to [Index]
74    fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
75
76    /// Cast to [vector::VectorIndex]
77    fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn vector::VectorIndex>>;
78
79    /// Retrieve index statistics as a JSON Value
80    fn statistics(&self) -> Result<serde_json::Value>;
81
82    /// Prewarm the index.
83    ///
84    /// This will load the index into memory and cache it.
85    async fn prewarm(&self) -> Result<()>;
86
87    /// Get the type of the index
88    fn index_type(&self) -> IndexType;
89
90    /// Read through the index and determine which fragment ids are covered by the index
91    ///
92    /// This is a kind of slow operation.  It's better to use the fragment_bitmap.  This
93    /// only exists for cases where the fragment_bitmap has become corrupted or missing.
94    async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
95}
96
97/// Index Type
98#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)]
99pub enum IndexType {
100    // Preserve 0-100 for simple indices.
101    Scalar = 0, // Legacy scalar index, alias to BTree
102
103    BTree = 1, // BTree
104
105    Bitmap = 2, // Bitmap
106
107    LabelList = 3, // LabelList
108
109    Inverted = 4, // Inverted
110
111    NGram = 5, // NGram
112
113    FragmentReuse = 6,
114
115    MemWal = 7,
116
117    ZoneMap = 8, // ZoneMap
118
119    BloomFilter = 9, // Bloom filter
120
121    // 100+ and up for vector index.
122    /// Flat vector index.
123    Vector = 100, // Legacy vector index, alias to IvfPq
124    IvfFlat = 101,
125    IvfSq = 102,
126    IvfPq = 103,
127    IvfHnswSq = 104,
128    IvfHnswPq = 105,
129    IvfHnswFlat = 106,
130    IvfRq = 107,
131}
132
133impl std::fmt::Display for IndexType {
134    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
135        match self {
136            Self::Scalar | Self::BTree => write!(f, "BTree"),
137            Self::Bitmap => write!(f, "Bitmap"),
138            Self::LabelList => write!(f, "LabelList"),
139            Self::Inverted => write!(f, "Inverted"),
140            Self::NGram => write!(f, "NGram"),
141            Self::FragmentReuse => write!(f, "FragmentReuse"),
142            Self::MemWal => write!(f, "MemWal"),
143            Self::ZoneMap => write!(f, "ZoneMap"),
144            Self::BloomFilter => write!(f, "BloomFilter"),
145            Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"),
146            Self::IvfFlat => write!(f, "IVF_FLAT"),
147            Self::IvfSq => write!(f, "IVF_SQ"),
148            Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"),
149            Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"),
150            Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"),
151            Self::IvfRq => write!(f, "IVF_RQ"),
152        }
153    }
154}
155
156impl TryFrom<i32> for IndexType {
157    type Error = Error;
158
159    fn try_from(value: i32) -> Result<Self> {
160        match value {
161            v if v == Self::Scalar as i32 => Ok(Self::Scalar),
162            v if v == Self::BTree as i32 => Ok(Self::BTree),
163            v if v == Self::Bitmap as i32 => Ok(Self::Bitmap),
164            v if v == Self::LabelList as i32 => Ok(Self::LabelList),
165            v if v == Self::NGram as i32 => Ok(Self::NGram),
166            v if v == Self::Inverted as i32 => Ok(Self::Inverted),
167            v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse),
168            v if v == Self::MemWal as i32 => Ok(Self::MemWal),
169            v if v == Self::ZoneMap as i32 => Ok(Self::ZoneMap),
170            v if v == Self::BloomFilter as i32 => Ok(Self::BloomFilter),
171            v if v == Self::Vector as i32 => Ok(Self::Vector),
172            v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat),
173            v if v == Self::IvfSq as i32 => Ok(Self::IvfSq),
174            v if v == Self::IvfPq as i32 => Ok(Self::IvfPq),
175            v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq),
176            v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq),
177            v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat),
178            _ => Err(Error::InvalidInput {
179                source: format!("the input value {} is not a valid IndexType", value).into(),
180                location: location!(),
181            }),
182        }
183    }
184}
185
186impl TryFrom<&str> for IndexType {
187    type Error = Error;
188
189    fn try_from(value: &str) -> Result<Self> {
190        match value {
191            "BTree" => Ok(Self::BTree),
192            "Bitmap" => Ok(Self::Bitmap),
193            "LabelList" => Ok(Self::LabelList),
194            "Inverted" => Ok(Self::Inverted),
195            "NGram" => Ok(Self::NGram),
196            "FragmentReuse" => Ok(Self::FragmentReuse),
197            "MemWal" => Ok(Self::MemWal),
198            "ZoneMap" => Ok(Self::ZoneMap),
199            "Vector" => Ok(Self::Vector),
200            "IVF_FLAT" => Ok(Self::IvfFlat),
201            "IVF_SQ" => Ok(Self::IvfSq),
202            "IVF_PQ" => Ok(Self::IvfPq),
203            "IVF_RQ" => Ok(Self::IvfRq),
204            "IVF_HNSW_FLAT" => Ok(Self::IvfHnswFlat),
205            "IVF_HNSW_SQ" => Ok(Self::IvfHnswSq),
206            "IVF_HNSW_PQ" => Ok(Self::IvfHnswPq),
207            _ => Err(Error::invalid_input(
208                format!("invalid index type: {}", value),
209                location!(),
210            )),
211        }
212    }
213}
214
215impl IndexType {
216    pub fn is_scalar(&self) -> bool {
217        matches!(
218            self,
219            Self::Scalar
220                | Self::BTree
221                | Self::Bitmap
222                | Self::LabelList
223                | Self::Inverted
224                | Self::NGram
225                | Self::ZoneMap
226                | Self::BloomFilter
227        )
228    }
229
230    pub fn is_vector(&self) -> bool {
231        matches!(
232            self,
233            Self::Vector
234                | Self::IvfPq
235                | Self::IvfHnswSq
236                | Self::IvfHnswPq
237                | Self::IvfHnswFlat
238                | Self::IvfFlat
239                | Self::IvfSq
240                | Self::IvfRq
241        )
242    }
243
244    pub fn is_system(&self) -> bool {
245        matches!(self, Self::FragmentReuse | Self::MemWal)
246    }
247
248    /// Returns the current format version of the index type,
249    /// bump this when the index format changes.
250    /// Indices which higher version than these will be ignored for compatibility,
251    /// This would happen when creating index in a newer version of Lance,
252    /// but then opening the index in older version of Lance
253    pub fn version(&self) -> i32 {
254        match self {
255            Self::Scalar => 0,
256            Self::BTree => 0,
257            Self::Bitmap => 0,
258            Self::LabelList => 0,
259            Self::Inverted => 0,
260            Self::NGram => 0,
261            Self::FragmentReuse => 0,
262            Self::MemWal => 0,
263            Self::ZoneMap => 0,
264            Self::BloomFilter => 0,
265
266            // for now all vector indices are built by the same builder,
267            // so they share the same version.
268            Self::Vector
269            | Self::IvfFlat
270            | Self::IvfSq
271            | Self::IvfPq
272            | Self::IvfHnswSq
273            | Self::IvfHnswPq
274            | Self::IvfHnswFlat
275            | Self::IvfRq => 1,
276        }
277    }
278
279    /// Returns the target partition size for the index type.
280    ///
281    /// This is used to compute the number of partitions for the index.
282    /// The partition size is optimized for the best performance of the index.
283    ///
284    /// This is for vector indices only.
285    pub fn target_partition_size(&self) -> usize {
286        match self {
287            Self::Vector => 8192,
288            Self::IvfFlat => 4096,
289            Self::IvfSq => 8192,
290            Self::IvfPq => 8192,
291            Self::IvfHnswFlat => 1 << 20,
292            Self::IvfHnswSq => 1 << 20,
293            Self::IvfHnswPq => 1 << 20,
294            _ => 8192,
295        }
296    }
297}
298
299pub trait IndexParams: Send + Sync {
300    fn as_any(&self) -> &dyn Any;
301
302    fn index_name(&self) -> &str;
303}
304
305#[derive(Serialize, Deserialize, Debug)]
306pub struct IndexMetadata {
307    #[serde(rename = "type")]
308    pub index_type: String,
309    pub distance_type: String,
310}
311
312pub fn is_system_index(index_meta: &lance_table::format::IndexMetadata) -> bool {
313    index_meta.name == FRAG_REUSE_INDEX_NAME || index_meta.name == MEM_WAL_INDEX_NAME
314}
315
316pub fn infer_system_index_type(
317    index_meta: &lance_table::format::IndexMetadata,
318) -> Option<IndexType> {
319    if index_meta.name == FRAG_REUSE_INDEX_NAME {
320        Some(IndexType::FragmentReuse)
321    } else if index_meta.name == MEM_WAL_INDEX_NAME {
322        Some(IndexType::MemWal)
323    } else {
324        None
325    }
326}