1use std::{any::Any, sync::Arc};
5
6use async_trait::async_trait;
7use lance_core::deepsize::DeepSizeOf;
8use lance_core::{Error, Result};
9use roaring::RoaringBitmap;
10use serde::{Deserialize, Serialize};
11use std::convert::TryFrom;
12
13pub mod metrics;
14pub mod scalar;
15
16#[async_trait]
19pub trait Index: Send + Sync + DeepSizeOf {
20 fn as_any(&self) -> &dyn Any;
22
23 fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
25
26 fn statistics(&self) -> Result<serde_json::Value>;
28
29 async fn prewarm(&self) -> Result<()>;
33
34 fn index_type(&self) -> IndexType;
36
37 async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
42}
43
44#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf, Serialize, Deserialize)]
46pub enum IndexType {
47 Scalar = 0, BTree = 1, Bitmap = 2, LabelList = 3, Inverted = 4, NGram = 5, FragmentReuse = 6,
61
62 MemWal = 7,
63
64 ZoneMap = 8, BloomFilter = 9, RTree = 10, Fm = 11, Vector = 100, IvfFlat = 101,
76 IvfSq = 102,
77 IvfPq = 103,
78 IvfHnswSq = 104,
79 IvfHnswPq = 105,
80 IvfHnswFlat = 106,
81 IvfRq = 107,
82}
83
84impl std::fmt::Display for IndexType {
85 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
86 match self {
87 Self::Scalar | Self::BTree => write!(f, "BTree"),
88 Self::Bitmap => write!(f, "Bitmap"),
89 Self::LabelList => write!(f, "LabelList"),
90 Self::Inverted => write!(f, "Inverted"),
91 Self::NGram => write!(f, "NGram"),
92 Self::FragmentReuse => write!(f, "FragmentReuse"),
93 Self::MemWal => write!(f, "MemWal"),
94 Self::ZoneMap => write!(f, "ZoneMap"),
95 Self::BloomFilter => write!(f, "BloomFilter"),
96 Self::RTree => write!(f, "RTree"),
97 Self::Fm => write!(f, "Fm"),
98 Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"),
99 Self::IvfFlat => write!(f, "IVF_FLAT"),
100 Self::IvfSq => write!(f, "IVF_SQ"),
101 Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"),
102 Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"),
103 Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"),
104 Self::IvfRq => write!(f, "IVF_RQ"),
105 }
106 }
107}
108
109impl TryFrom<i32> for IndexType {
110 type Error = Error;
111
112 fn try_from(value: i32) -> Result<Self> {
113 match value {
114 v if v == Self::Scalar as i32 => Ok(Self::Scalar),
115 v if v == Self::BTree as i32 => Ok(Self::BTree),
116 v if v == Self::Bitmap as i32 => Ok(Self::Bitmap),
117 v if v == Self::LabelList as i32 => Ok(Self::LabelList),
118 v if v == Self::NGram as i32 => Ok(Self::NGram),
119 v if v == Self::Inverted as i32 => Ok(Self::Inverted),
120 v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse),
121 v if v == Self::MemWal as i32 => Ok(Self::MemWal),
122 v if v == Self::ZoneMap as i32 => Ok(Self::ZoneMap),
123 v if v == Self::BloomFilter as i32 => Ok(Self::BloomFilter),
124 v if v == Self::RTree as i32 => Ok(Self::RTree),
125 v if v == Self::Fm as i32 => Ok(Self::Fm),
126 v if v == Self::Vector as i32 => Ok(Self::Vector),
127 v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat),
128 v if v == Self::IvfSq as i32 => Ok(Self::IvfSq),
129 v if v == Self::IvfPq as i32 => Ok(Self::IvfPq),
130 v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq),
131 v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq),
132 v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat),
133 v if v == Self::IvfRq as i32 => Ok(Self::IvfRq),
134 _ => Err(Error::invalid_input_source(
135 format!("the input value {} is not a valid IndexType", value).into(),
136 )),
137 }
138 }
139}
140
141impl TryFrom<&str> for IndexType {
142 type Error = Error;
143
144 fn try_from(value: &str) -> Result<Self> {
145 match value {
146 "BTree" | "BTREE" => Ok(Self::BTree),
147 "Bitmap" | "BITMAP" => Ok(Self::Bitmap),
148 "LabelList" | "LABELLIST" => Ok(Self::LabelList),
149 "Inverted" | "INVERTED" => Ok(Self::Inverted),
150 "NGram" | "NGRAM" => Ok(Self::NGram),
151 "ZoneMap" | "ZONEMAP" => Ok(Self::ZoneMap),
152 "BloomFilter" | "BLOOMFILTER" | "BLOOM_FILTER" => Ok(Self::BloomFilter),
153 "RTree" | "RTREE" | "R_TREE" => Ok(Self::RTree),
154 "Fm" | "FM" => Ok(Self::Fm),
155 "Vector" | "VECTOR" => Ok(Self::Vector),
156 "IVF_FLAT" => Ok(Self::IvfFlat),
157 "IVF_SQ" => Ok(Self::IvfSq),
158 "IVF_PQ" => Ok(Self::IvfPq),
159 "IVF_RQ" => Ok(Self::IvfRq),
160 "IVF_HNSW_FLAT" => Ok(Self::IvfHnswFlat),
161 "IVF_HNSW_SQ" => Ok(Self::IvfHnswSq),
162 "IVF_HNSW_PQ" => Ok(Self::IvfHnswPq),
163 "FragmentReuse" => Ok(Self::FragmentReuse),
164 "MemWal" => Ok(Self::MemWal),
165 _ => Err(Error::invalid_input(format!(
166 "invalid index type: {}",
167 value
168 ))),
169 }
170 }
171}
172
173impl IndexType {
174 pub fn is_scalar(&self) -> bool {
175 matches!(
176 self,
177 Self::Scalar
178 | Self::BTree
179 | Self::Bitmap
180 | Self::LabelList
181 | Self::Inverted
182 | Self::NGram
183 | Self::ZoneMap
184 | Self::BloomFilter
185 | Self::RTree
186 | Self::Fm,
187 )
188 }
189
190 pub fn is_vector(&self) -> bool {
191 matches!(
192 self,
193 Self::Vector
194 | Self::IvfPq
195 | Self::IvfHnswSq
196 | Self::IvfHnswPq
197 | Self::IvfHnswFlat
198 | Self::IvfFlat
199 | Self::IvfSq
200 | Self::IvfRq
201 )
202 }
203
204 pub fn is_system(&self) -> bool {
205 matches!(self, Self::FragmentReuse | Self::MemWal)
206 }
207
208 pub fn version(&self) -> i32 {
214 match self {
215 Self::Scalar => 0,
216 Self::BTree => 0,
217 Self::Bitmap => 0,
218 Self::LabelList => 0,
219 Self::Inverted => 0,
220 Self::NGram => 0,
221 Self::FragmentReuse => 0,
222 Self::MemWal => 0,
223 Self::ZoneMap => 0,
224 Self::BloomFilter => 0,
225 Self::RTree => 0,
226 Self::Fm => 0,
227
228 Self::Vector
235 | Self::IvfFlat
236 | Self::IvfSq
237 | Self::IvfPq
238 | Self::IvfHnswSq
239 | Self::IvfHnswPq
240 | Self::IvfHnswFlat => 1,
241 Self::IvfRq => 2,
242 }
243 }
244
245 pub fn target_partition_size(&self) -> usize {
252 match self {
253 Self::Vector => 8192,
254 Self::IvfFlat => 4096,
255 Self::IvfSq => 8192,
256 Self::IvfPq => 8192,
257 Self::IvfRq => 4096,
258 Self::IvfHnswFlat => 1 << 20,
259 Self::IvfHnswSq => 1 << 20,
260 Self::IvfHnswPq => 1 << 20,
261 _ => 8192,
262 }
263 }
264
265 pub fn max_vector_version() -> u32 {
267 [
268 Self::Vector,
269 Self::IvfFlat,
270 Self::IvfSq,
271 Self::IvfPq,
272 Self::IvfHnswSq,
273 Self::IvfHnswPq,
274 Self::IvfHnswFlat,
275 Self::IvfRq,
276 ]
277 .into_iter()
278 .map(|index_type| index_type.version() as u32)
279 .max()
280 .unwrap_or(1)
281 }
282}
283
284pub trait IndexParams: Send + Sync {
285 fn as_any(&self) -> &dyn Any;
286
287 fn index_name(&self) -> &str;
288}