Skip to main content

lance_index_core/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use 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/// Generic methods common across all types of secondary indices
17///
18#[async_trait]
19pub trait Index: Send + Sync + DeepSizeOf {
20    /// Cast to [Any].
21    fn as_any(&self) -> &dyn Any;
22
23    /// Cast to [Index]
24    fn as_index(self: Arc<Self>) -> Arc<dyn Index>;
25
26    /// Retrieve index statistics as a JSON Value
27    fn statistics(&self) -> Result<serde_json::Value>;
28
29    /// Prewarm the index.
30    ///
31    /// This will load the index into memory and cache it.
32    async fn prewarm(&self) -> Result<()>;
33
34    /// Get the type of the index
35    fn index_type(&self) -> IndexType;
36
37    /// Read through the index and determine which fragment ids are covered by the index
38    ///
39    /// This is a kind of slow operation.  It's better to use the fragment_bitmap.  This
40    /// only exists for cases where the fragment_bitmap has become corrupted or missing.
41    async fn calculate_included_frags(&self) -> Result<RoaringBitmap>;
42}
43
44/// Index Type
45#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf, Serialize, Deserialize)]
46pub enum IndexType {
47    // Preserve 0-100 for simple indices.
48    Scalar = 0, // Legacy scalar index, alias to BTree
49
50    BTree = 1, // BTree
51
52    Bitmap = 2, // Bitmap
53
54    LabelList = 3, // LabelList
55
56    Inverted = 4, // Inverted
57
58    NGram = 5, // NGram
59
60    FragmentReuse = 6,
61
62    MemWal = 7,
63
64    ZoneMap = 8, // ZoneMap
65
66    BloomFilter = 9, // Bloom filter
67
68    RTree = 10, // RTree
69
70    Fm = 11, // FM-Index
71
72    // 100+ and up for vector index.
73    /// Flat vector index.
74    Vector = 100, // Legacy vector index, alias to IvfPq
75    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    /// Returns the current format version of the index type,
209    /// bump this when the index format changes.
210    /// Indices which higher version than these will be ignored for compatibility,
211    /// This would happen when creating index in a newer version of Lance,
212    /// but then opening the index in older version of Lance
213    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            // IMPORTANT: if any vector index subtype needs a format bump that is
229            // not backward compatible, its new version must be set to
230            // (current max vector index version + 1), even if only one subtype
231            // changed. Compatibility filtering currently cannot distinguish vector
232            // subtypes from details-only metadata, so vector versions effectively
233            // share one global monotonic compatibility level.
234            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    /// Returns the target partition size for the index type.
246    ///
247    /// This is used to compute the number of partitions for the index.
248    /// The partition size is optimized for the best performance of the index.
249    ///
250    /// This is for vector indices only.
251    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    /// Returns the highest supported vector index version in this Lance build.
266    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}