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 lance_core::utils::row_addr_remap::RowAddrRemap;
8use std::any::Any;
9use std::fmt::Debug;
10use std::sync::Arc;
11
12use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array};
13use arrow_schema::Field;
14use async_trait::async_trait;
15use datafusion::execution::SendableRecordBatchStream;
16use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
17use futures::stream;
18use ivf::storage::IvfModel;
19use lance_core::deepsize::DeepSizeOf;
20use lance_core::{Error, ROW_ID_FIELD, Result};
21use lance_io::traits::Reader;
22use lance_linalg::distance::DistanceType;
23use quantizer::{QuantizationType, Quantizer};
24use std::sync::LazyLock;
25use v3::subindex::SubIndexType;
26
27pub mod bq;
28pub mod distributed;
29pub mod flat;
30pub mod graph;
31pub mod hnsw;
32pub mod ivf;
33pub mod kmeans;
34pub mod pq;
35pub mod quantizer;
36pub mod residual;
37pub mod shared;
38pub mod sq;
39pub mod storage;
40pub mod transform;
41pub mod utils;
42pub mod v3;
43
44use super::pb;
45use crate::metrics::MetricsCollector;
46use crate::{Index, prefilter::PreFilter};
47
48// TODO: Make these crate private once the migration from lance to lance-index is done.
49pub const DIST_COL: &str = "_distance";
50pub const DISTANCE_TYPE_KEY: &str = "distance_type";
51pub const INDEX_UUID_COLUMN: &str = "__index_uuid";
52pub const PART_ID_COLUMN: &str = "__ivf_part_id";
53pub const DIST_Q_C_COLUMN: &str = "__dist_q_c";
54// dist from vector to centroid
55pub const CENTROID_DIST_COLUMN: &str = "__centroid_dist";
56pub const PQ_CODE_COLUMN: &str = "__pq_code";
57pub const SQ_CODE_COLUMN: &str = "__sq_code";
58pub const LOSS_METADATA_KEY: &str = "_loss";
59
60pub type PreparedPartitionSearchHandle = Box<dyn Any + Send>;
61
62/// Controls when a multi-partition search should stop producing more partition results.
63pub trait PartitionSearchControl: Send + Sync {
64    fn should_stop(&self) -> bool;
65
66    fn record_batch(&self, _batch: &RecordBatch) {}
67}
68
69pub static VECTOR_RESULT_SCHEMA: LazyLock<arrow_schema::SchemaRef> = LazyLock::new(|| {
70    arrow_schema::SchemaRef::new(arrow_schema::Schema::new(vec![
71        Field::new(DIST_COL, arrow_schema::DataType::Float32, true),
72        ROW_ID_FIELD.clone(),
73    ]))
74});
75
76pub static PART_ID_FIELD: LazyLock<arrow_schema::Field> = LazyLock::new(|| {
77    arrow_schema::Field::new(PART_ID_COLUMN, arrow_schema::DataType::UInt32, true)
78});
79
80pub static CENTROID_DIST_FIELD: LazyLock<arrow_schema::Field> = LazyLock::new(|| {
81    arrow_schema::Field::new(CENTROID_DIST_COLUMN, arrow_schema::DataType::Float32, true)
82});
83
84pub const DEFAULT_QUERY_PARALLELISM: i32 = 0;
85
86/// Controls the speed / accuracy tradeoff for approximate vector search.
87///
88/// This currently only affects RQ-quantized vector indexes, such as IVF_RQ.
89/// Other index types ignore this setting.
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
91pub enum ApproxMode {
92    /// Prefer lower query latency, which can reduce recall.
93    Fast,
94
95    /// Use the default balance between query latency and recall.
96    #[default]
97    Normal,
98
99    /// Prefer higher recall, which can increase query latency.
100    Accurate,
101}
102
103/// Query parameters for the vector indices
104
105#[derive(Debug, Clone)]
106pub struct Query {
107    /// The column to be searched.
108    pub column: String,
109
110    /// The vector to be searched.
111    pub key: ArrayRef,
112
113    /// Top k results to return.
114    pub k: usize,
115
116    /// The lower bound (inclusive) of the distance to be searched.
117    pub lower_bound: Option<f32>,
118
119    /// The upper bound (exclusive) of the distance to be searched.
120    pub upper_bound: Option<f32>,
121
122    /// The minimum number of probes to load and search.  More partitions
123    /// will only be loaded if we have not found k results, or the algorithm
124    /// determines more partitions are needed to satisfy recall requirements.
125    ///
126    /// The planner will always search at least this many partitions. Defaults to 1.
127    pub minimum_nprobes: usize,
128
129    /// The maximum number of probes to load and search.  If not set then
130    /// ALL partitions will be searched, if needed, to satisfy k results.
131    pub maximum_nprobes: Option<usize>,
132
133    /// The number of candidates to reserve while searching.
134    /// this is an optional parameter for HNSW related index types.
135    pub ef: Option<usize>,
136
137    /// If presented, apply a refine step.
138    /// TODO: should we support fraction / float number here?
139    pub refine_factor: Option<u32>,
140
141    /// Distance metric type. If None, uses the index's metric (if available)
142    /// or the default for the data type.
143    pub metric_type: Option<DistanceType>,
144
145    /// Whether to use an ANN index if available
146    pub use_index: bool,
147
148    /// Maximum partition-search concurrency for a single vector query.
149    ///
150    /// The default is 0.
151    /// Value 0 selects the automatic policy; today this resolves to 1 for the
152    /// sequential fast path unless an index implementation overrides it.
153    /// Value -1 uses the CPU pool size.
154    /// Value 1 uses the single-worker sequential partition search path.
155    /// Values >= 2 use the partition-parallel path and are clamped to the CPU
156    /// pool size by the execution layer.
157    pub query_parallelism: i32,
158
159    /// the distance between the query and the centroid
160    /// this is only used for IVF index with Rabit quantization
161    pub dist_q_c: f32,
162
163    /// Controls the speed / accuracy tradeoff for approximate vector search.
164    ///
165    /// This currently only affects RQ-quantized vector indexes, such as IVF_RQ.
166    /// Other index types ignore this setting.
167    pub approx_mode: ApproxMode,
168}
169
170impl From<pb::VectorMetricType> for DistanceType {
171    fn from(proto: pb::VectorMetricType) -> Self {
172        match proto {
173            pb::VectorMetricType::L2 => Self::L2,
174            pb::VectorMetricType::Cosine => Self::Cosine,
175            pb::VectorMetricType::Dot => Self::Dot,
176            pb::VectorMetricType::Hamming => Self::Hamming,
177        }
178    }
179}
180
181impl From<DistanceType> for pb::VectorMetricType {
182    fn from(mt: DistanceType) -> Self {
183        match mt {
184            DistanceType::L2 => Self::L2,
185            DistanceType::Cosine => Self::Cosine,
186            DistanceType::Dot => Self::Dot,
187            DistanceType::Hamming => Self::Hamming,
188        }
189    }
190}
191
192/// Vector Index for (Approximate) Nearest Neighbor (ANN) Search.
193///
194/// Vector indices are often built as a chain of indices.  For example, IVF -> PQ
195/// or IVF -> HNSW -> SQ.
196///
197/// We use one trait for both the top-level and the sub-indices.  Typically the top-level
198/// search is a partition-aware search and all sub-indices are whole-index searches.
199#[async_trait]
200#[allow(clippy::redundant_pub_crate)]
201pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index {
202    /// Search entire index for k nearest neighbors.
203    ///
204    /// It returns a [RecordBatch] with Schema of:
205    ///
206    /// ```
207    /// use arrow_schema::{Schema, Field, DataType};
208    ///
209    /// Schema::new(vec![
210    ///   Field::new("_rowid", DataType::UInt64, true),
211    ///   Field::new("_distance", DataType::Float32, true),
212    /// ]);
213    /// ```
214    ///
215    /// The `pre_filter` argument is used to filter out row ids that we know are
216    /// not relevant to the query. For example, it removes deleted rows or rows that
217    /// do not match a user-provided filter.
218    async fn search(
219        &self,
220        query: &Query,
221        pre_filter: Arc<dyn PreFilter>,
222        metrics: &dyn MetricsCollector,
223    ) -> Result<RecordBatch>;
224
225    /// Find partitions that may contain nearest neighbors.
226    ///
227    /// If maximum_nprobes is set then this method will return the partitions
228    /// that are most likely to contain the nearest neighbors (e.g. the closest
229    /// partitions to the query vector).
230    ///
231    /// Return the partition ids and the distances between the query and the centroids,
232    /// the results should be in sorted order from closest to farthest.
233    fn find_partitions(&self, query: &Query) -> Result<(UInt32Array, Float32Array)>;
234
235    /// Get the total number of partitions in the index.
236    fn total_partitions(&self) -> usize;
237
238    /// Search a single partition for nearest neighbors.
239    ///
240    /// This method should return the same results as [`VectorIndex::search`] method except
241    /// that it will only search a single partition.
242    async fn search_in_partition(
243        &self,
244        partition_id: usize,
245        query: &Query,
246        pre_filter: Arc<dyn PreFilter>,
247        metrics: &dyn MetricsCollector,
248    ) -> Result<RecordBatch>;
249
250    /// Asynchronously prepare a single-partition search so the CPU-heavy portion
251    /// can be executed separately.
252    async fn prepare_partition_search(
253        &self,
254        _partition_id: usize,
255        _query: &Query,
256        _pre_filter: Arc<dyn PreFilter>,
257        _metrics: &dyn MetricsCollector,
258    ) -> Result<PreparedPartitionSearchHandle> {
259        unimplemented!("prepared partition search is not supported for this index")
260    }
261
262    /// Execute the synchronous portion of a previously prepared partition search.
263    fn search_prepared_partition(
264        &self,
265        _prepared: PreparedPartitionSearchHandle,
266        _metrics: &dyn MetricsCollector,
267    ) -> Result<RecordBatch> {
268        unimplemented!("prepared partition search is not supported for this index")
269    }
270
271    /// Return true if the index supports splitting partition search into async
272    /// prepare and sync execute phases.
273    fn supports_prepared_partition_search(&self) -> bool {
274        false
275    }
276
277    /// Choose partition search concurrency for `query_parallelism = 0`.
278    ///
279    /// The default keeps the single-worker sequential path. Index
280    /// implementations can override this when their sub-index search work does
281    /// not benefit from the sequential fast path.
282    fn auto_query_parallelism(&self, _cpu_pool_size: usize) -> usize {
283        1
284    }
285
286    /// Search a range of partitions and return a stream of per-partition result batches.
287    ///
288    /// The default implementation searches each partition sequentially with
289    /// [`VectorIndex::search_in_partition`]. Implementations can override this
290    /// to use a more efficient execution strategy.
291    #[allow(clippy::too_many_arguments)]
292    async fn search_partitions(
293        self: Arc<Self>,
294        query: Query,
295        partitions: Arc<UInt32Array>,
296        q_c_dists: Arc<Float32Array>,
297        start_idx: usize,
298        end_idx: usize,
299        pre_filter: Arc<dyn PreFilter>,
300        control: Option<Arc<dyn PartitionSearchControl>>,
301        metrics: Arc<dyn MetricsCollector>,
302    ) -> Result<SendableRecordBatchStream>
303    where
304        Self: 'static,
305    {
306        if partitions.len() != q_c_dists.len() {
307            return Err(Error::invalid_input(format!(
308                "partition count {} does not match centroid distance count {}",
309                partitions.len(),
310                q_c_dists.len()
311            )));
312        }
313        if start_idx > end_idx || end_idx > partitions.len() {
314            return Err(Error::invalid_input(format!(
315                "invalid partition search range [{start_idx}, {end_idx}) for {} partitions",
316                partitions.len()
317            )));
318        }
319
320        let stream = stream::try_unfold(start_idx, move |idx| {
321            let index = self.clone();
322            let partitions = partitions.clone();
323            let q_c_dists = q_c_dists.clone();
324            let query = query.clone();
325            let pre_filter = pre_filter.clone();
326            let control = control.clone();
327            let metrics = metrics.clone();
328            async move {
329                if idx >= end_idx
330                    || control
331                        .as_ref()
332                        .is_some_and(|control| control.should_stop())
333                {
334                    return Ok(None);
335                }
336                let part_id = partitions.value(idx);
337                let mut query = query;
338                query.dist_q_c = q_c_dists.value(idx);
339                index
340                    .search_in_partition(part_id as usize, &query, pre_filter, metrics.as_ref())
341                    .await
342                    .map(|batch| {
343                        if let Some(control) = control.as_ref() {
344                            control.record_batch(&batch);
345                        }
346                        Some((batch, idx + 1))
347                    })
348                    .map_err(Into::into)
349            }
350        });
351        Ok(Box::pin(RecordBatchStreamAdapter::new(
352            VECTOR_RESULT_SCHEMA.clone(),
353            stream,
354        )))
355    }
356
357    /// If the index is loadable by IVF, so it can be a sub-index that
358    /// is loaded on demand by IVF.
359    fn is_loadable(&self) -> bool;
360
361    /// Use residual vector to search.
362    fn use_residual(&self) -> bool;
363
364    // async fn append(&self, batches: Vec<RecordBatch>) -> Result<()>;
365    // async fn merge(&self, indices: Vec<Arc<dyn VectorIndex>>) -> Result<()>;
366
367    /// Load the index from the reader on-demand.
368    async fn load(
369        &self,
370        reader: Arc<dyn Reader>,
371        offset: usize,
372        length: usize,
373    ) -> Result<Box<dyn VectorIndex>>;
374
375    /// Load the partition from the reader on-demand.
376    async fn load_partition(
377        &self,
378        reader: Arc<dyn Reader>,
379        offset: usize,
380        length: usize,
381        _partition_id: usize,
382    ) -> Result<Box<dyn VectorIndex>> {
383        self.load(reader, offset, length).await
384    }
385
386    // for IVF only
387    async fn partition_reader(
388        &self,
389        _partition_id: usize,
390        _with_vector: bool,
391        _metrics: &dyn MetricsCollector,
392    ) -> Result<SendableRecordBatchStream> {
393        unimplemented!("only for IVF")
394    }
395
396    // for SubIndex only
397    async fn to_batch_stream(&self, with_vector: bool) -> Result<SendableRecordBatchStream>;
398
399    fn num_rows(&self) -> u64;
400
401    /// Return the IDs of rows in the index.
402    fn row_ids(&self) -> Box<dyn Iterator<Item = &'_ u64> + '_>;
403
404    /// Remap the index according to mapping
405    ///
406    /// Each item in mapping describes an old row id -> new row id
407    /// pair.  If old row id -> None then that row id has been
408    /// deleted and can be removed from the index.
409    ///
410    /// If an old row id is not in the mapping then it should be
411    /// left alone.
412    async fn remap(&mut self, mapping: &RowAddrRemap) -> Result<()>;
413
414    /// The metric type of this vector index.
415    fn metric_type(&self) -> DistanceType;
416
417    fn ivf_model(&self) -> &IvfModel;
418    fn quantizer(&self) -> Quantizer;
419    fn partition_size(&self, part_id: usize) -> usize;
420
421    /// the index type of this vector index.
422    fn sub_index_type(&self) -> (SubIndexType, QuantizationType);
423
424    /// The cumulative I/O performed while opening this index (file footers, IVF
425    /// centroids, quantization metadata).  This is a one-time cost; it is
426    /// reported once, on the query that actually opens the index, and is `None`
427    /// for index implementations that do not track it.
428    fn open_io_stats(&self) -> Option<lance_io::scheduler::ScanStats> {
429        None
430    }
431}
432
433// it can be an IVF index or a partition of IVF index
434pub trait VectorIndexCacheEntry: Debug + Send + Sync + DeepSizeOf {
435    fn as_any(&self) -> &dyn Any;
436}
lance_index/vector.rs

lance_index/
vector.rs
