1use 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";
36pub 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#[async_trait]
51pub trait Index: Send + Sync + DeepSizeOf {
52 fn as_any(&self) -> &dyn Any;
54
55 fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
57
58 fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn vector::VectorIndex>>;
60
61 fn statistics(&self) -> Result<serde_json::Value>;
63
64 async fn prewarm(&self) -> Result<()>;
68
69 fn index_type(&self) -> IndexType;
71
72 async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
77}
78
79#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)]
81pub enum IndexType {
82 Scalar = 0, BTree = 1, Bitmap = 2, LabelList = 3, Inverted = 4, NGram = 5, FragmentReuse = 6,
96
97 MemWal = 7,
98
99 Vector = 100, 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 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 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}