lance_index/
vector.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Vector Index
5//!
6
7use std::any::Any;
8use std::fmt::Debug;
9use std::{collections::HashMap, sync::Arc};
10
11use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array};
12use arrow_schema::Field;
13use async_trait::async_trait;
14use datafusion::execution::SendableRecordBatchStream;
15use deepsize::DeepSizeOf;
16use ivf::storage::IvfModel;
17use lance_core::{Result, ROW_ID_FIELD};
18use lance_io::traits::Reader;
19use lance_linalg::distance::DistanceType;
20use quantizer::{QuantizationType, Quantizer};
21use std::sync::LazyLock;
22use v3::subindex::SubIndexType;
23
24pub mod bq;
25pub mod flat;
26pub mod graph;
27pub mod hnsw;
28pub mod ivf;
29pub mod kmeans;
30pub mod pq;
31pub mod quantizer;
32pub mod residual;
33pub mod sq;
34pub mod storage;
35pub mod transform;
36pub mod utils;
37pub mod v3;
38
39use super::pb;
40use crate::metrics::MetricsCollector;
41use crate::{prefilter::PreFilter, Index};
42
43// TODO: Make these crate private once the migration from lance to lance-index is done.
44pub const DIST_COL: &str = "_distance";
45pub const DISTANCE_TYPE_KEY: &str = "distance_type";
46pub const INDEX_UUID_COLUMN: &str = "__index_uuid";
47pub const PART_ID_COLUMN: &str = "__ivf_part_id";
48pub const DIST_Q_C_COLUMN: &str = "__dist_q_c";
49// dist from vector to centroid
50pub const CENTROID_DIST_COLUMN: &str = "__centroid_dist";
51pub const PQ_CODE_COLUMN: &str = "__pq_code";
52pub const SQ_CODE_COLUMN: &str = "__sq_code";
53pub const LOSS_METADATA_KEY: &str = "_loss";
54
55pub static VECTOR_RESULT_SCHEMA: LazyLock<arrow_schema::SchemaRef> = LazyLock::new(|| {
56    arrow_schema::SchemaRef::new(arrow_schema::Schema::new(vec![
57        Field::new(DIST_COL, arrow_schema::DataType::Float32, false),
58        ROW_ID_FIELD.clone(),
59    ]))
60});
61
62pub static PART_ID_FIELD: LazyLock<arrow_schema::Field> = LazyLock::new(|| {
63    arrow_schema::Field::new(PART_ID_COLUMN, arrow_schema::DataType::UInt32, true)
64});
65
66pub static CENTROID_DIST_FIELD: LazyLock<arrow_schema::Field> = LazyLock::new(|| {
67    arrow_schema::Field::new(CENTROID_DIST_COLUMN, arrow_schema::DataType::Float32, true)
68});
69
70/// Query parameters for the vector indices
71#[derive(Debug, Clone)]
72pub struct Query {
73    /// The column to be searched.
74    pub column: String,
75
76    /// The vector to be searched.
77    pub key: ArrayRef,
78
79    /// Top k results to return.
80    pub k: usize,
81
82    /// The lower bound (inclusive) of the distance to be searched.
83    pub lower_bound: Option<f32>,
84
85    /// The upper bound (exclusive) of the distance to be searched.
86    pub upper_bound: Option<f32>,
87
88    /// The minimum number of probes to load and search.  More partitions
89    /// will only be loaded if we have not found k results, or the the algorithm
90    /// determines more partitions are needed to satisfy recall requirements.
91    ///
92    /// The planner will always search at least this many partitions. Defaults to 1.
93    pub minimum_nprobes: usize,
94
95    /// The maximum number of probes to load and search.  If not set then
96    /// ALL partitions will be searched, if needed, to satisfy k results.
97    pub maximum_nprobes: Option<usize>,
98
99    /// The number of candidates to reserve while searching.
100    /// this is an optional parameter for HNSW related index types.
101    pub ef: Option<usize>,
102
103    /// If presented, apply a refine step.
104    /// TODO: should we support fraction / float number here?
105    pub refine_factor: Option<u32>,
106
107    /// Distance metric type
108    pub metric_type: DistanceType,
109
110    /// Whether to use an ANN index if available
111    pub use_index: bool,
112
113    /// the distance between the query and the centroid
114    /// this is only used for IVF index with Rabit quantization
115    pub dist_q_c: f32,
116}
117
118impl From<pb::VectorMetricType> for DistanceType {
119    fn from(proto: pb::VectorMetricType) -> Self {
120        match proto {
121            pb::VectorMetricType::L2 => Self::L2,
122            pb::VectorMetricType::Cosine => Self::Cosine,
123            pb::VectorMetricType::Dot => Self::Dot,
124            pb::VectorMetricType::Hamming => Self::Hamming,
125        }
126    }
127}
128
129impl From<DistanceType> for pb::VectorMetricType {
130    fn from(mt: DistanceType) -> Self {
131        match mt {
132            DistanceType::L2 => Self::L2,
133            DistanceType::Cosine => Self::Cosine,
134            DistanceType::Dot => Self::Dot,
135            DistanceType::Hamming => Self::Hamming,
136        }
137    }
138}
139
140/// Vector Index for (Approximate) Nearest Neighbor (ANN) Search.
141///
142/// Vector indices are often built as a chain of indices.  For example, IVF -> PQ
143/// or IVF -> HNSW -> SQ.
144///
145/// We use one trait for both the top-level and the sub-indices.  Typically the top-level
146/// search is a partition-aware search and all sub-indices are whole-index searches.
147#[async_trait]
148#[allow(clippy::redundant_pub_crate)]
149pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index {
150    /// Search entire index for k nearest neighbors.
151    ///
152    /// It returns a [RecordBatch] with Schema of:
153    ///
154    /// ```
155    /// use arrow_schema::{Schema, Field, DataType};
156    ///
157    /// Schema::new(vec![
158    ///   Field::new("_rowid", DataType::UInt64, true),
159    ///   Field::new("_distance", DataType::Float32, false),
160    /// ]);
161    /// ```
162    ///
163    /// The `pre_filter` argument is used to filter out row ids that we know are
164    /// not relevant to the query. For example, it removes deleted rows or rows that
165    /// do not match a user-provided filter.
166    async fn search(
167        &self,
168        query: &Query,
169        pre_filter: Arc<dyn PreFilter>,
170        metrics: &dyn MetricsCollector,
171    ) -> Result<RecordBatch>;
172
173    /// Find partitions that may contain nearest neighbors.
174    ///
175    /// If maximum_nprobes is set then this method will return the partitions
176    /// that are most likely to contain the nearest neighbors (e.g. the closest
177    /// partitions to the query vector).
178    ///
179    /// Return the partition ids and the distances between the query and the centroids,
180    /// the results should be in sorted order from closest to farthest.
181    fn find_partitions(&self, query: &Query) -> Result<(UInt32Array, Float32Array)>;
182
183    /// Get the total number of partitions in the index.
184    fn total_partitions(&self) -> usize;
185
186    /// Search a single partition for nearest neighbors.
187    ///
188    /// This method should return the same results as [`VectorIndex::search`] method except
189    /// that it will only search a single partition.
190    async fn search_in_partition(
191        &self,
192        partition_id: usize,
193        query: &Query,
194        pre_filter: Arc<dyn PreFilter>,
195        metrics: &dyn MetricsCollector,
196    ) -> Result<RecordBatch>;
197
198    /// If the index is loadable by IVF, so it can be a sub-index that
199    /// is loaded on demand by IVF.
200    fn is_loadable(&self) -> bool;
201
202    /// Use residual vector to search.
203    fn use_residual(&self) -> bool;
204
205    // async fn append(&self, batches: Vec<RecordBatch>) -> Result<()>;
206    // async fn merge(&self, indices: Vec<Arc<dyn VectorIndex>>) -> Result<()>;
207
208    /// Load the index from the reader on-demand.
209    async fn load(
210        &self,
211        reader: Arc<dyn Reader>,
212        offset: usize,
213        length: usize,
214    ) -> Result<Box<dyn VectorIndex>>;
215
216    /// Load the partition from the reader on-demand.
217    async fn load_partition(
218        &self,
219        reader: Arc<dyn Reader>,
220        offset: usize,
221        length: usize,
222        _partition_id: usize,
223    ) -> Result<Box<dyn VectorIndex>> {
224        self.load(reader, offset, length).await
225    }
226
227    // for IVF only
228    async fn partition_reader(
229        &self,
230        _partition_id: usize,
231        _with_vector: bool,
232        _metrics: &dyn MetricsCollector,
233    ) -> Result<SendableRecordBatchStream> {
234        unimplemented!("only for IVF")
235    }
236
237    // for SubIndex only
238    async fn to_batch_stream(&self, with_vector: bool) -> Result<SendableRecordBatchStream>;
239
240    fn num_rows(&self) -> u64;
241
242    /// Return the IDs of rows in the index.
243    fn row_ids(&self) -> Box<dyn Iterator<Item = &'_ u64> + '_>;
244
245    /// Remap the index according to mapping
246    ///
247    /// Each item in mapping describes an old row id -> new row id
248    /// pair.  If old row id -> None then that row id has been
249    /// deleted and can be removed from the index.
250    ///
251    /// If an old row id is not in the mapping then it should be
252    /// left alone.
253    async fn remap(&mut self, mapping: &HashMap<u64, Option<u64>>) -> Result<()>;
254
255    /// The metric type of this vector index.
256    fn metric_type(&self) -> DistanceType;
257
258    fn ivf_model(&self) -> &IvfModel;
259    fn quantizer(&self) -> Quantizer;
260    fn partition_size(&self, part_id: usize) -> usize;
261
262    /// the index type of this vector index.
263    fn sub_index_type(&self) -> (SubIndexType, QuantizationType);
264}
265
266// it can be an IVF index or a partition of IVF index
267pub trait VectorIndexCacheEntry: Debug + Send + Sync + DeepSizeOf {
268    fn as_any(&self) -> &dyn Any;
269}
lance_index/vector.rs

lance_index/
vector.rs
