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 const VECTOR_INDEX_VERSION: u32 = 1;
45
46pub mod pb {
47 #![allow(clippy::use_self)]
48 include!(concat!(env!("OUT_DIR"), "/lance.index.pb.rs"));
49}
50
51#[async_trait]
54pub trait Index: Send + Sync + DeepSizeOf {
55 fn as_any(&self) -> &dyn Any;
57
58 fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
60
61 fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn vector::VectorIndex>>;
63
64 fn statistics(&self) -> Result<serde_json::Value>;
66
67 async fn prewarm(&self) -> Result<()>;
71
72 fn index_type(&self) -> IndexType;
74
75 async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
80}
81
82#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)]
84pub enum IndexType {
85 Scalar = 0, BTree = 1, Bitmap = 2, LabelList = 3, Inverted = 4, NGram = 5, FragmentReuse = 6,
99
100 MemWal = 7,
101
102 ZoneMap = 8, Vector = 100, IvfFlat = 101,
108 IvfSq = 102,
109 IvfPq = 103,
110 IvfHnswSq = 104,
111 IvfHnswPq = 105,
112 IvfHnswFlat = 106,
113}
114
115impl std::fmt::Display for IndexType {
116 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
117 match self {
118 Self::Scalar | Self::BTree => write!(f, "BTree"),
119 Self::Bitmap => write!(f, "Bitmap"),
120 Self::LabelList => write!(f, "LabelList"),
121 Self::Inverted => write!(f, "Inverted"),
122 Self::NGram => write!(f, "NGram"),
123 Self::FragmentReuse => write!(f, "FragmentReuse"),
124 Self::MemWal => write!(f, "MemWal"),
125 Self::ZoneMap => write!(f, "ZoneMap"),
126 Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"),
127 Self::IvfFlat => write!(f, "IVF_FLAT"),
128 Self::IvfSq => write!(f, "IVF_SQ"),
129 Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"),
130 Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"),
131 Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"),
132 }
133 }
134}
135
136impl TryFrom<i32> for IndexType {
137 type Error = Error;
138
139 fn try_from(value: i32) -> Result<Self> {
140 match value {
141 v if v == Self::Scalar as i32 => Ok(Self::Scalar),
142 v if v == Self::BTree as i32 => Ok(Self::BTree),
143 v if v == Self::Bitmap as i32 => Ok(Self::Bitmap),
144 v if v == Self::LabelList as i32 => Ok(Self::LabelList),
145 v if v == Self::NGram as i32 => Ok(Self::NGram),
146 v if v == Self::Inverted as i32 => Ok(Self::Inverted),
147 v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse),
148 v if v == Self::MemWal as i32 => Ok(Self::MemWal),
149 v if v == Self::ZoneMap as i32 => Ok(Self::ZoneMap),
150 v if v == Self::Vector as i32 => Ok(Self::Vector),
151 v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat),
152 v if v == Self::IvfSq as i32 => Ok(Self::IvfSq),
153 v if v == Self::IvfPq as i32 => Ok(Self::IvfPq),
154 v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq),
155 v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq),
156 v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat),
157 _ => Err(Error::InvalidInput {
158 source: format!("the input value {} is not a valid IndexType", value).into(),
159 location: location!(),
160 }),
161 }
162 }
163}
164
165impl IndexType {
166 pub fn is_scalar(&self) -> bool {
167 matches!(
168 self,
169 Self::Scalar
170 | Self::BTree
171 | Self::Bitmap
172 | Self::LabelList
173 | Self::Inverted
174 | Self::NGram
175 | Self::ZoneMap
176 )
177 }
178
179 pub fn is_vector(&self) -> bool {
180 matches!(
181 self,
182 Self::Vector
183 | Self::IvfPq
184 | Self::IvfHnswSq
185 | Self::IvfHnswPq
186 | Self::IvfHnswFlat
187 | Self::IvfFlat
188 | Self::IvfSq
189 )
190 }
191
192 pub fn is_system(&self) -> bool {
193 matches!(self, Self::FragmentReuse | Self::MemWal)
194 }
195}
196
197pub trait IndexParams: Send + Sync {
198 fn as_any(&self) -> &dyn Any;
199
200 fn index_name(&self) -> &str;
201}
202
203#[derive(Serialize, Deserialize, Debug)]
204pub struct IndexMetadata {
205 #[serde(rename = "type")]
206 pub index_type: String,
207 pub distance_type: String,
208}
209
210pub fn is_system_index(index_meta: &lance_table::format::Index) -> bool {
211 index_meta.name == FRAG_REUSE_INDEX_NAME || index_meta.name == MEM_WAL_INDEX_NAME
212}
213
214pub fn infer_system_index_type(index_meta: &lance_table::format::Index) -> Option<IndexType> {
215 if index_meta.name == FRAG_REUSE_INDEX_NAME {
216 Some(IndexType::FragmentReuse)
217 } else if index_meta.name == MEM_WAL_INDEX_NAME {
218 Some(IndexType::MemWal)
219 } else {
220 None
221 }
222}