Skip to main content

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 lance_core::deepsize::DeepSizeOf;
18use lance_core::{Error, Result};
19use roaring::RoaringBitmap;
20use serde::{Deserialize, Serialize};
21use std::convert::TryFrom;
22
23pub mod frag_reuse;
24pub mod mem_wal;
25pub mod metrics;
26pub mod optimize;
27pub mod prefilter;
28pub mod progress;
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/// Default version for vector index metadata.
45///
46/// Most vector indices should use this version unless they need to bump for a
47/// format change.
48pub const VECTOR_INDEX_VERSION: u32 = 1;
49/// Version for IVF_RQ indices.
50pub const IVF_RQ_INDEX_VERSION: u32 = 2;
51
52/// The factor of threshold to trigger split / join for vector index.
53///
54/// If the number of rows in the single partition is greater than `MAX_PARTITION_SIZE_FACTOR * target_partition_size`,
55/// the partition will be split.
56/// If the number of rows in the single partition is less than `MIN_PARTITION_SIZE_PERCENT *target_partition_size / 100`,
57/// the partition will be joined.
58pub const MAX_PARTITION_SIZE_FACTOR: usize = 4;
59pub const MIN_PARTITION_SIZE_PERCENT: usize = 25;
60
61pub mod pb {
62    #![allow(clippy::use_self)]
63    include!(concat!(env!("OUT_DIR"), "/lance.index.pb.rs"));
64}
65
66pub mod pbold {
67    #![allow(clippy::use_self)]
68    include!(concat!(env!("OUT_DIR"), "/lance.table.rs"));
69}
70
71/// Protobuf headers for serialized index cache entries (FTS posting lists,
72/// scalar indices, and IVF vector partitions).
73pub mod cache_pb {
74    #![allow(clippy::use_self)]
75    include!(concat!(env!("OUT_DIR"), "/lance.index.cache.rs"));
76}
77
78/// Generic methods common across all types of secondary indices
79///
80#[async_trait]
81pub trait Index: Send + Sync + DeepSizeOf {
82    /// Cast to [Any].
83    fn as_any(&self) -> &dyn Any;
84
85    /// Cast to [Index]
86    fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
87
88    /// Cast to [vector::VectorIndex]
89    fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn vector::VectorIndex>>;
90
91    /// Retrieve index statistics as a JSON Value
92    fn statistics(&self) -> Result<serde_json::Value>;
93
94    /// Prewarm the index.
95    ///
96    /// This will load the index into memory and cache it.
97    async fn prewarm(&self) -> Result<()>;
98
99    /// Get the type of the index
100    fn index_type(&self) -> IndexType;
101
102    /// Read through the index and determine which fragment ids are covered by the index
103    ///
104    /// This is a kind of slow operation.  It's better to use the fragment_bitmap.  This
105    /// only exists for cases where the fragment_bitmap has become corrupted or missing.
106    async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
107}
108
109/// Index Type
110#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)]
111pub enum IndexType {
112    // Preserve 0-100 for simple indices.
113    Scalar = 0, // Legacy scalar index, alias to BTree
114
115    BTree = 1, // BTree
116
117    Bitmap = 2, // Bitmap
118
119    LabelList = 3, // LabelList
120
121    Inverted = 4, // Inverted
122
123    NGram = 5, // NGram
124
125    FragmentReuse = 6,
126
127    MemWal = 7,
128
129    ZoneMap = 8, // ZoneMap
130
131    BloomFilter = 9, // Bloom filter
132
133    RTree = 10, // RTree
134
135    Fm = 11, // FM-Index
136
137    // 100+ and up for vector index.
138    /// Flat vector index.
139    Vector = 100, // Legacy vector index, alias to IvfPq
140    IvfFlat = 101,
141    IvfSq = 102,
142    IvfPq = 103,
143    IvfHnswSq = 104,
144    IvfHnswPq = 105,
145    IvfHnswFlat = 106,
146    IvfRq = 107,
147}
148
149impl std::fmt::Display for IndexType {
150    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
151        match self {
152            Self::Scalar | Self::BTree => write!(f, "BTree"),
153            Self::Bitmap => write!(f, "Bitmap"),
154            Self::LabelList => write!(f, "LabelList"),
155            Self::Inverted => write!(f, "Inverted"),
156            Self::NGram => write!(f, "NGram"),
157            Self::FragmentReuse => write!(f, "FragmentReuse"),
158            Self::MemWal => write!(f, "MemWal"),
159            Self::ZoneMap => write!(f, "ZoneMap"),
160            Self::BloomFilter => write!(f, "BloomFilter"),
161            Self::RTree => write!(f, "RTree"),
162            Self::Fm => write!(f, "Fm"),
163            Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"),
164            Self::IvfFlat => write!(f, "IVF_FLAT"),
165            Self::IvfSq => write!(f, "IVF_SQ"),
166            Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"),
167            Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"),
168            Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"),
169            Self::IvfRq => write!(f, "IVF_RQ"),
170        }
171    }
172}
173
174impl TryFrom<i32> for IndexType {
175    type Error = Error;
176
177    fn try_from(value: i32) -> Result<Self> {
178        match value {
179            v if v == Self::Scalar as i32 => Ok(Self::Scalar),
180            v if v == Self::BTree as i32 => Ok(Self::BTree),
181            v if v == Self::Bitmap as i32 => Ok(Self::Bitmap),
182            v if v == Self::LabelList as i32 => Ok(Self::LabelList),
183            v if v == Self::NGram as i32 => Ok(Self::NGram),
184            v if v == Self::Inverted as i32 => Ok(Self::Inverted),
185            v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse),
186            v if v == Self::MemWal as i32 => Ok(Self::MemWal),
187            v if v == Self::ZoneMap as i32 => Ok(Self::ZoneMap),
188            v if v == Self::BloomFilter as i32 => Ok(Self::BloomFilter),
189            v if v == Self::RTree as i32 => Ok(Self::RTree),
190            v if v == Self::Fm as i32 => Ok(Self::Fm),
191            v if v == Self::Vector as i32 => Ok(Self::Vector),
192            v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat),
193            v if v == Self::IvfSq as i32 => Ok(Self::IvfSq),
194            v if v == Self::IvfPq as i32 => Ok(Self::IvfPq),
195            v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq),
196            v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq),
197            v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat),
198            v if v == Self::IvfRq as i32 => Ok(Self::IvfRq),
199            _ => Err(Error::invalid_input_source(
200                format!("the input value {} is not a valid IndexType", value).into(),
201            )),
202        }
203    }
204}
205
206impl TryFrom<&str> for IndexType {
207    type Error = Error;
208
209    fn try_from(value: &str) -> Result<Self> {
210        match value {
211            "BTree" | "BTREE" => Ok(Self::BTree),
212            "Bitmap" | "BITMAP" => Ok(Self::Bitmap),
213            "LabelList" | "LABELLIST" => Ok(Self::LabelList),
214            "Inverted" | "INVERTED" => Ok(Self::Inverted),
215            "NGram" | "NGRAM" => Ok(Self::NGram),
216            "ZoneMap" | "ZONEMAP" => Ok(Self::ZoneMap),
217            "BloomFilter" | "BLOOMFILTER" | "BLOOM_FILTER" => Ok(Self::BloomFilter),
218            "RTree" | "RTREE" | "R_TREE" => Ok(Self::RTree),
219            "Fm" | "FM" => Ok(Self::Fm),
220            "Vector" | "VECTOR" => Ok(Self::Vector),
221            "IVF_FLAT" => Ok(Self::IvfFlat),
222            "IVF_SQ" => Ok(Self::IvfSq),
223            "IVF_PQ" => Ok(Self::IvfPq),
224            "IVF_RQ" => Ok(Self::IvfRq),
225            "IVF_HNSW_FLAT" => Ok(Self::IvfHnswFlat),
226            "IVF_HNSW_SQ" => Ok(Self::IvfHnswSq),
227            "IVF_HNSW_PQ" => Ok(Self::IvfHnswPq),
228            "FragmentReuse" => Ok(Self::FragmentReuse),
229            "MemWal" => Ok(Self::MemWal),
230            _ => Err(Error::invalid_input(format!(
231                "invalid index type: {}",
232                value
233            ))),
234        }
235    }
236}
237
238impl IndexType {
239    pub fn is_scalar(&self) -> bool {
240        matches!(
241            self,
242            Self::Scalar
243                | Self::BTree
244                | Self::Bitmap
245                | Self::LabelList
246                | Self::Inverted
247                | Self::NGram
248                | Self::ZoneMap
249                | Self::BloomFilter
250                | Self::RTree
251                | Self::Fm,
252        )
253    }
254
255    pub fn is_vector(&self) -> bool {
256        matches!(
257            self,
258            Self::Vector
259                | Self::IvfPq
260                | Self::IvfHnswSq
261                | Self::IvfHnswPq
262                | Self::IvfHnswFlat
263                | Self::IvfFlat
264                | Self::IvfSq
265                | Self::IvfRq
266        )
267    }
268
269    pub fn is_system(&self) -> bool {
270        matches!(self, Self::FragmentReuse | Self::MemWal)
271    }
272
273    /// Returns the current format version of the index type,
274    /// bump this when the index format changes.
275    /// Indices which higher version than these will be ignored for compatibility,
276    /// This would happen when creating index in a newer version of Lance,
277    /// but then opening the index in older version of Lance
278    pub fn version(&self) -> i32 {
279        match self {
280            Self::Scalar => 0,
281            Self::BTree => 0,
282            Self::Bitmap => 0,
283            Self::LabelList => 0,
284            Self::Inverted => 0,
285            Self::NGram => 0,
286            Self::FragmentReuse => 0,
287            Self::MemWal => 0,
288            Self::ZoneMap => 0,
289            Self::BloomFilter => 0,
290            Self::RTree => 0,
291            Self::Fm => 0,
292
293            // IMPORTANT: if any vector index subtype needs a format bump that is
294            // not backward compatible, its new version must be set to
295            // (current max vector index version + 1), even if only one subtype
296            // changed. Compatibility filtering currently cannot distinguish vector
297            // subtypes from details-only metadata, so vector versions effectively
298            // share one global monotonic compatibility level.
299            Self::Vector
300            | Self::IvfFlat
301            | Self::IvfSq
302            | Self::IvfPq
303            | Self::IvfHnswSq
304            | Self::IvfHnswPq
305            | Self::IvfHnswFlat => VECTOR_INDEX_VERSION as i32,
306            Self::IvfRq => IVF_RQ_INDEX_VERSION as i32,
307        }
308    }
309
310    /// Returns the target partition size for the index type.
311    ///
312    /// This is used to compute the number of partitions for the index.
313    /// The partition size is optimized for the best performance of the index.
314    ///
315    /// This is for vector indices only.
316    pub fn target_partition_size(&self) -> usize {
317        match self {
318            Self::Vector => 8192,
319            Self::IvfFlat => 4096,
320            Self::IvfSq => 8192,
321            Self::IvfPq => 8192,
322            Self::IvfRq => 4096,
323            Self::IvfHnswFlat => 1 << 20,
324            Self::IvfHnswSq => 1 << 20,
325            Self::IvfHnswPq => 1 << 20,
326            _ => 8192,
327        }
328    }
329
330    /// Returns the highest supported vector index version in this Lance build.
331    pub fn max_vector_version() -> u32 {
332        [
333            Self::Vector,
334            Self::IvfFlat,
335            Self::IvfSq,
336            Self::IvfPq,
337            Self::IvfHnswSq,
338            Self::IvfHnswPq,
339            Self::IvfHnswFlat,
340            Self::IvfRq,
341        ]
342        .into_iter()
343        .map(|index_type| index_type.version() as u32)
344        .max()
345        .unwrap_or(VECTOR_INDEX_VERSION)
346    }
347}
348
349pub trait IndexParams: Send + Sync {
350    fn as_any(&self) -> &dyn Any;
351
352    fn index_name(&self) -> &str;
353}
354
355#[derive(Serialize, Deserialize, Debug)]
356pub struct IndexMetadata {
357    #[serde(rename = "type")]
358    pub index_type: String,
359    pub distance_type: String,
360}
361
362pub fn is_system_index(index_meta: &lance_table::format::IndexMetadata) -> bool {
363    index_meta.name == FRAG_REUSE_INDEX_NAME || index_meta.name == MEM_WAL_INDEX_NAME
364}
365
366pub fn infer_system_index_type(
367    index_meta: &lance_table::format::IndexMetadata,
368) -> Option<IndexType> {
369    if index_meta.name == FRAG_REUSE_INDEX_NAME {
370        Some(IndexType::FragmentReuse)
371    } else if index_meta.name == MEM_WAL_INDEX_NAME {
372        Some(IndexType::MemWal)
373    } else {
374        None
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381
382    #[test]
383    fn test_ivf_rq_has_dedicated_index_version() {
384        assert!(IndexType::IvfRq.version() > IndexType::IvfPq.version());
385        assert_eq!(IndexType::IvfRq.version() as u32, IVF_RQ_INDEX_VERSION);
386    }
387
388    #[test]
389    fn test_max_vector_version_tracks_highest_supported() {
390        assert_eq!(IndexType::max_vector_version(), IVF_RQ_INDEX_VERSION);
391    }
392
393    #[test]
394    fn test_ivf_rq_target_partition_size() {
395        assert_eq!(IndexType::IvfRq.target_partition_size(), 4096);
396    }
397
398    #[test]
399    fn test_index_type_try_from_i32_covers_all_variants() {
400        let all = [
401            IndexType::Scalar,
402            IndexType::BTree,
403            IndexType::Bitmap,
404            IndexType::LabelList,
405            IndexType::Inverted,
406            IndexType::NGram,
407            IndexType::FragmentReuse,
408            IndexType::MemWal,
409            IndexType::ZoneMap,
410            IndexType::BloomFilter,
411            IndexType::RTree,
412            IndexType::Fm,
413            IndexType::Vector,
414            IndexType::IvfFlat,
415            IndexType::IvfSq,
416            IndexType::IvfPq,
417            IndexType::IvfHnswSq,
418            IndexType::IvfHnswPq,
419            IndexType::IvfHnswFlat,
420            IndexType::IvfRq,
421        ];
422
423        for index_type in all {
424            assert_eq!(
425                IndexType::try_from(index_type as i32).unwrap(),
426                index_type,
427                "IndexType::try_from(i32) should support {:?}",
428                index_type
429            );
430        }
431    }
432
433    #[test]
434    fn test_index_type_try_from_str_covers_all_parseable_variants() {
435        let cases = [
436            ("BTree", IndexType::BTree),
437            ("BTREE", IndexType::BTree),
438            ("Bitmap", IndexType::Bitmap),
439            ("BITMAP", IndexType::Bitmap),
440            ("LabelList", IndexType::LabelList),
441            ("LABELLIST", IndexType::LabelList),
442            ("Inverted", IndexType::Inverted),
443            ("INVERTED", IndexType::Inverted),
444            ("NGram", IndexType::NGram),
445            ("NGRAM", IndexType::NGram),
446            ("ZoneMap", IndexType::ZoneMap),
447            ("ZONEMAP", IndexType::ZoneMap),
448            ("BloomFilter", IndexType::BloomFilter),
449            ("BLOOMFILTER", IndexType::BloomFilter),
450            ("BLOOM_FILTER", IndexType::BloomFilter),
451            ("RTree", IndexType::RTree),
452            ("RTREE", IndexType::RTree),
453            ("R_TREE", IndexType::RTree),
454            ("Fm", IndexType::Fm),
455            ("FM", IndexType::Fm),
456            ("Vector", IndexType::Vector),
457            ("VECTOR", IndexType::Vector),
458            ("IVF_FLAT", IndexType::IvfFlat),
459            ("IVF_SQ", IndexType::IvfSq),
460            ("IVF_PQ", IndexType::IvfPq),
461            ("IVF_RQ", IndexType::IvfRq),
462            ("IVF_HNSW_FLAT", IndexType::IvfHnswFlat),
463            ("IVF_HNSW_SQ", IndexType::IvfHnswSq),
464            ("IVF_HNSW_PQ", IndexType::IvfHnswPq),
465            ("FragmentReuse", IndexType::FragmentReuse),
466            ("MemWal", IndexType::MemWal),
467        ];
468
469        for (text, expected) in cases {
470            assert_eq!(
471                IndexType::try_from(text).unwrap(),
472                expected,
473                "IndexType::try_from(&str) should support '{text}'"
474            );
475        }
476    }
477}