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 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;
46
47pub const MAX_PARTITION_SIZE_FACTOR: usize = 4;
54pub const MIN_PARTITION_SIZE_PERCENT: usize = 25;
55
56pub mod pb {
57 #![allow(clippy::use_self)]
58 include!(concat!(env!("OUT_DIR"), "/lance.index.pb.rs"));
59}
60
61pub mod pbold {
62 #![allow(clippy::use_self)]
63 include!(concat!(env!("OUT_DIR"), "/lance.table.rs"));
64}
65
66#[async_trait]
69pub trait Index: Send + Sync + DeepSizeOf {
70 fn as_any(&self) -> &dyn Any;
72
73 fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
75
76 fn as_vector_index(self: Arc<Self>) -> Result<Arc<dyn vector::VectorIndex>>;
78
79 fn statistics(&self) -> Result<serde_json::Value>;
81
82 async fn prewarm(&self) -> Result<()>;
86
87 fn index_type(&self) -> IndexType;
89
90 async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
95}
96
97#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)]
99pub enum IndexType {
100 Scalar = 0, BTree = 1, Bitmap = 2, LabelList = 3, Inverted = 4, NGram = 5, FragmentReuse = 6,
114
115 MemWal = 7,
116
117 ZoneMap = 8, BloomFilter = 9, Vector = 100, IvfFlat = 101,
125 IvfSq = 102,
126 IvfPq = 103,
127 IvfHnswSq = 104,
128 IvfHnswPq = 105,
129 IvfHnswFlat = 106,
130 IvfRq = 107,
131}
132
133impl std::fmt::Display for IndexType {
134 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
135 match self {
136 Self::Scalar | Self::BTree => write!(f, "BTree"),
137 Self::Bitmap => write!(f, "Bitmap"),
138 Self::LabelList => write!(f, "LabelList"),
139 Self::Inverted => write!(f, "Inverted"),
140 Self::NGram => write!(f, "NGram"),
141 Self::FragmentReuse => write!(f, "FragmentReuse"),
142 Self::MemWal => write!(f, "MemWal"),
143 Self::ZoneMap => write!(f, "ZoneMap"),
144 Self::BloomFilter => write!(f, "BloomFilter"),
145 Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"),
146 Self::IvfFlat => write!(f, "IVF_FLAT"),
147 Self::IvfSq => write!(f, "IVF_SQ"),
148 Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"),
149 Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"),
150 Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"),
151 Self::IvfRq => write!(f, "IVF_RQ"),
152 }
153 }
154}
155
156impl TryFrom<i32> for IndexType {
157 type Error = Error;
158
159 fn try_from(value: i32) -> Result<Self> {
160 match value {
161 v if v == Self::Scalar as i32 => Ok(Self::Scalar),
162 v if v == Self::BTree as i32 => Ok(Self::BTree),
163 v if v == Self::Bitmap as i32 => Ok(Self::Bitmap),
164 v if v == Self::LabelList as i32 => Ok(Self::LabelList),
165 v if v == Self::NGram as i32 => Ok(Self::NGram),
166 v if v == Self::Inverted as i32 => Ok(Self::Inverted),
167 v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse),
168 v if v == Self::MemWal as i32 => Ok(Self::MemWal),
169 v if v == Self::ZoneMap as i32 => Ok(Self::ZoneMap),
170 v if v == Self::BloomFilter as i32 => Ok(Self::BloomFilter),
171 v if v == Self::Vector as i32 => Ok(Self::Vector),
172 v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat),
173 v if v == Self::IvfSq as i32 => Ok(Self::IvfSq),
174 v if v == Self::IvfPq as i32 => Ok(Self::IvfPq),
175 v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq),
176 v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq),
177 v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat),
178 _ => Err(Error::InvalidInput {
179 source: format!("the input value {} is not a valid IndexType", value).into(),
180 location: location!(),
181 }),
182 }
183 }
184}
185
186impl TryFrom<&str> for IndexType {
187 type Error = Error;
188
189 fn try_from(value: &str) -> Result<Self> {
190 match value {
191 "BTree" => Ok(Self::BTree),
192 "Bitmap" => Ok(Self::Bitmap),
193 "LabelList" => Ok(Self::LabelList),
194 "Inverted" => Ok(Self::Inverted),
195 "NGram" => Ok(Self::NGram),
196 "FragmentReuse" => Ok(Self::FragmentReuse),
197 "MemWal" => Ok(Self::MemWal),
198 "ZoneMap" => Ok(Self::ZoneMap),
199 "Vector" => Ok(Self::Vector),
200 "IVF_FLAT" => Ok(Self::IvfFlat),
201 "IVF_SQ" => Ok(Self::IvfSq),
202 "IVF_PQ" => Ok(Self::IvfPq),
203 "IVF_RQ" => Ok(Self::IvfRq),
204 "IVF_HNSW_FLAT" => Ok(Self::IvfHnswFlat),
205 "IVF_HNSW_SQ" => Ok(Self::IvfHnswSq),
206 "IVF_HNSW_PQ" => Ok(Self::IvfHnswPq),
207 _ => Err(Error::invalid_input(
208 format!("invalid index type: {}", value),
209 location!(),
210 )),
211 }
212 }
213}
214
215impl IndexType {
216 pub fn is_scalar(&self) -> bool {
217 matches!(
218 self,
219 Self::Scalar
220 | Self::BTree
221 | Self::Bitmap
222 | Self::LabelList
223 | Self::Inverted
224 | Self::NGram
225 | Self::ZoneMap
226 | Self::BloomFilter
227 )
228 }
229
230 pub fn is_vector(&self) -> bool {
231 matches!(
232 self,
233 Self::Vector
234 | Self::IvfPq
235 | Self::IvfHnswSq
236 | Self::IvfHnswPq
237 | Self::IvfHnswFlat
238 | Self::IvfFlat
239 | Self::IvfSq
240 | Self::IvfRq
241 )
242 }
243
244 pub fn is_system(&self) -> bool {
245 matches!(self, Self::FragmentReuse | Self::MemWal)
246 }
247
248 pub fn version(&self) -> i32 {
254 match self {
255 Self::Scalar => 0,
256 Self::BTree => 0,
257 Self::Bitmap => 0,
258 Self::LabelList => 0,
259 Self::Inverted => 0,
260 Self::NGram => 0,
261 Self::FragmentReuse => 0,
262 Self::MemWal => 0,
263 Self::ZoneMap => 0,
264 Self::BloomFilter => 0,
265
266 Self::Vector
269 | Self::IvfFlat
270 | Self::IvfSq
271 | Self::IvfPq
272 | Self::IvfHnswSq
273 | Self::IvfHnswPq
274 | Self::IvfHnswFlat
275 | Self::IvfRq => 1,
276 }
277 }
278
279 pub fn target_partition_size(&self) -> usize {
286 match self {
287 Self::Vector => 8192,
288 Self::IvfFlat => 4096,
289 Self::IvfSq => 8192,
290 Self::IvfPq => 8192,
291 Self::IvfHnswFlat => 1 << 20,
292 Self::IvfHnswSq => 1 << 20,
293 Self::IvfHnswPq => 1 << 20,
294 _ => 8192,
295 }
296 }
297}
298
299pub trait IndexParams: Send + Sync {
300 fn as_any(&self) -> &dyn Any;
301
302 fn index_name(&self) -> &str;
303}
304
305#[derive(Serialize, Deserialize, Debug)]
306pub struct IndexMetadata {
307 #[serde(rename = "type")]
308 pub index_type: String,
309 pub distance_type: String,
310}
311
312pub fn is_system_index(index_meta: &lance_table::format::IndexMetadata) -> bool {
313 index_meta.name == FRAG_REUSE_INDEX_NAME || index_meta.name == MEM_WAL_INDEX_NAME
314}
315
316pub fn infer_system_index_type(
317 index_meta: &lance_table::format::IndexMetadata,
318) -> Option<IndexType> {
319 if index_meta.name == FRAG_REUSE_INDEX_NAME {
320 Some(IndexType::FragmentReuse)
321 } else if index_meta.name == MEM_WAL_INDEX_NAME {
322 Some(IndexType::MemWal)
323 } else {
324 None
325 }
326}