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/lancedb/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 scalar;
30pub mod traits;
31pub mod vector;
32
33pub use crate::traits::*;
34
35pub const INDEX_FILE_NAME: &str = "index.idx";
36/// The name of the auxiliary index file.
37///
38/// This file is used to store additional information about the index, to improve performance.
39/// - For 'IVF_HNSW' index, it stores the partitioned PQ Storage.
40pub const INDEX_AUXILIARY_FILE_NAME: &str = "auxiliary.idx";
41pub const INDEX_METADATA_SCHEMA_KEY: &str = "lance:index";
42
43pub mod pb {
44    #![allow(clippy::use_self)]
45    include!(concat!(env!("OUT_DIR"), "/lance.index.pb.rs"));
46}
47
48/// Generic methods common across all types of secondary indices
49///
50#[async_trait]
51pub trait Index: Send + Sync + DeepSizeOf {
52    /// Cast to [Any].
53    fn as_any(&self) -> &dyn Any;
54
55    /// Cast to [Index]
56    fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
57
58    /// Cast to [vector::VectorIndex]
59    fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn vector::VectorIndex>>;
60
61    /// Retrieve index statistics as a JSON Value
62    fn statistics(&self) -> Result<serde_json::Value>;
63
64    /// Prewarm the index.
65    ///
66    /// This will load the index into memory and cache it.
67    async fn prewarm(&self) -> Result<()>;
68
69    /// Get the type of the index
70    fn index_type(&self) -> IndexType;
71
72    /// Read through the index and determine which fragment ids are covered by the index
73    ///
74    /// This is a kind of slow operation.  It's better to use the fragment_bitmap.  This
75    /// only exists for cases where the fragment_bitmap has become corrupted or missing.
76    async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
77}
78
79/// Index Type
80#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)]
81pub enum IndexType {
82    // Preserve 0-100 for simple indices.
83    Scalar = 0, // Legacy scalar index, alias to BTree
84
85    BTree = 1, // BTree
86
87    Bitmap = 2, // Bitmap
88
89    LabelList = 3, // LabelList
90
91    Inverted = 4, // Inverted
92
93    NGram = 5, // NGram
94
95    FragmentReuse = 6,
96
97    MemWal = 7,
98
99    // 100+ and up for vector index.
100    /// Flat vector index.
101    Vector = 100, // Legacy vector index, alias to IvfPq
102    IvfFlat = 101,
103    IvfSq = 102,
104    IvfPq = 103,
105    IvfHnswSq = 104,
106    IvfHnswPq = 105,
107    IvfHnswFlat = 106,
108}
109
110impl std::fmt::Display for IndexType {
111    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
112        match self {
113            Self::Scalar | Self::BTree => write!(f, "BTree"),
114            Self::Bitmap => write!(f, "Bitmap"),
115            Self::LabelList => write!(f, "LabelList"),
116            Self::Inverted => write!(f, "Inverted"),
117            Self::NGram => write!(f, "NGram"),
118            Self::FragmentReuse => write!(f, "FragmentReuse"),
119            Self::MemWal => write!(f, "MemWal"),
120            Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"),
121            Self::IvfFlat => write!(f, "IVF_FLAT"),
122            Self::IvfSq => write!(f, "IVF_SQ"),
123            Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"),
124            Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"),
125            Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"),
126        }
127    }
128}
129
130impl TryFrom<i32> for IndexType {
131    type Error = Error;
132
133    fn try_from(value: i32) -> Result<Self> {
134        match value {
135            v if v == Self::Scalar as i32 => Ok(Self::Scalar),
136            v if v == Self::BTree as i32 => Ok(Self::BTree),
137            v if v == Self::Bitmap as i32 => Ok(Self::Bitmap),
138            v if v == Self::LabelList as i32 => Ok(Self::LabelList),
139            v if v == Self::NGram as i32 => Ok(Self::NGram),
140            v if v == Self::Inverted as i32 => Ok(Self::Inverted),
141            v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse),
142            v if v == Self::MemWal as i32 => Ok(Self::MemWal),
143            v if v == Self::Vector as i32 => Ok(Self::Vector),
144            v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat),
145            v if v == Self::IvfSq as i32 => Ok(Self::IvfSq),
146            v if v == Self::IvfPq as i32 => Ok(Self::IvfPq),
147            v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq),
148            v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq),
149            v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat),
150            _ => Err(Error::InvalidInput {
151                source: format!("the input value {} is not a valid IndexType", value).into(),
152                location: location!(),
153            }),
154        }
155    }
156}
157
158impl IndexType {
159    pub fn is_scalar(&self) -> bool {
160        matches!(
161            self,
162            Self::Scalar
163                | Self::BTree
164                | Self::Bitmap
165                | Self::LabelList
166                | Self::Inverted
167                | Self::NGram
168        )
169    }
170
171    pub fn is_vector(&self) -> bool {
172        matches!(
173            self,
174            Self::Vector
175                | Self::IvfPq
176                | Self::IvfHnswSq
177                | Self::IvfHnswPq
178                | Self::IvfHnswFlat
179                | Self::IvfFlat
180                | Self::IvfSq
181        )
182    }
183
184    pub fn is_system(&self) -> bool {
185        matches!(self, Self::FragmentReuse | Self::MemWal)
186    }
187
188    /// Returns the current format version of the index type,
189    /// bump this when the index format changes.
190    /// Indices which higher version than these will be ignored for compatibility,
191    /// This would happen when creating index in a newer version of Lance,
192    /// but then opening the index in older version of Lance
193    pub fn version(&self) -> i32 {
194        match self {
195            Self::Scalar => 0,
196            Self::BTree => 0,
197            Self::Bitmap => 0,
198            Self::LabelList => 0,
199            Self::Inverted => 0,
200            Self::NGram => 0,
201            Self::FragmentReuse => 0,
202            Self::MemWal => 0,
203
204            // for now all vector indices are built by the same builder,
205            // so they share the same version.
206            Self::Vector
207            | Self::IvfFlat
208            | Self::IvfSq
209            | Self::IvfPq
210            | Self::IvfHnswSq
211            | Self::IvfHnswPq
212            | Self::IvfHnswFlat => 1,
213        }
214    }
215}
216
217pub trait IndexParams: Send + Sync {
218    fn as_any(&self) -> &dyn Any;
219
220    fn index_type(&self) -> IndexType;
221
222    fn index_name(&self) -> &str;
223}
224
225#[derive(Serialize, Deserialize, Debug)]
226pub struct IndexMetadata {
227    #[serde(rename = "type")]
228    pub index_type: String,
229    pub distance_type: String,
230}
231
232pub fn is_system_index(index_meta: &lance_table::format::Index) -> bool {
233    index_meta.name == FRAG_REUSE_INDEX_NAME || index_meta.name == MEM_WAL_INDEX_NAME
234}
235
236pub fn infer_system_index_type(index_meta: &lance_table::format::Index) -> Option<IndexType> {
237    if index_meta.name == FRAG_REUSE_INDEX_NAME {
238        Some(IndexType::FragmentReuse)
239    } else if index_meta.name == MEM_WAL_INDEX_NAME {
240        Some(IndexType::MemWal)
241    } else {
242        None
243    }
244}