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 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";
37pub const INDEX_AUXILIARY_FILE_NAME: &str = "auxiliary.idx";
42pub const INDEX_METADATA_SCHEMA_KEY: &str = "lance:index";
43
44pub const VECTOR_INDEX_VERSION: u32 = 1;
49pub const IVF_RQ_INDEX_VERSION: u32 = 2;
51
52pub 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
71pub mod cache_pb {
74 #![allow(clippy::use_self)]
75 include!(concat!(env!("OUT_DIR"), "/lance.index.cache.rs"));
76}
77
78#[async_trait]
81pub trait Index: Send + Sync + DeepSizeOf {
82 fn as_any(&self) -> &dyn Any;
84
85 fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
87
88 fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn vector::VectorIndex>>;
90
91 fn statistics(&self) -> Result<serde_json::Value>;
93
94 async fn prewarm(&self) -> Result<()>;
98
99 fn index_type(&self) -> IndexType;
101
102 async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
107}
108
109#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)]
111pub enum IndexType {
112 Scalar = 0, BTree = 1, Bitmap = 2, LabelList = 3, Inverted = 4, NGram = 5, FragmentReuse = 6,
126
127 MemWal = 7,
128
129 ZoneMap = 8, BloomFilter = 9, RTree = 10, Fm = 11, Vector = 100, 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 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 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 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 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}