lance_index/vector/flat/
index.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Flat Vector Index.
5//!
6
7use std::collections::{BinaryHeap, HashMap};
8use std::sync::Arc;
9
10use arrow::array::AsArray;
11use arrow_array::{Array, ArrayRef, Float32Array, RecordBatch, UInt64Array};
12use arrow_schema::{DataType, Field, Schema, SchemaRef};
13use deepsize::DeepSizeOf;
14use lance_core::{Error, ROW_ID_FIELD, Result};
15use lance_file::previous::reader::FileReader as PreviousFileReader;
16use lance_linalg::distance::DistanceType;
17use serde::{Deserialize, Serialize};
18
19use crate::{
20    metrics::MetricsCollector,
21    prefilter::PreFilter,
22    vector::{
23        DIST_COL, Query,
24        graph::OrderedNode,
25        quantizer::{Quantization, QuantizationType, Quantizer, QuantizerMetadata},
26        storage::{DistCalculator, VectorStore},
27        v3::subindex::IvfSubIndex,
28    },
29};
30
31use super::storage::{FLAT_COLUMN, FlatBinStorage, FlatFloatStorage};
32
33/// A Flat index is any index that stores no metadata, and
34/// during query, it simply scans over the storage and returns the top k results
35#[derive(Debug, Clone, Default, DeepSizeOf)]
36pub struct FlatIndex {}
37
38use std::sync::LazyLock;
39
40static ANN_SEARCH_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
41    Schema::new(vec![
42        Field::new(DIST_COL, DataType::Float32, true),
43        ROW_ID_FIELD.clone(),
44    ])
45    .into()
46});
47
48#[derive(Default)]
49pub struct FlatQueryParams {
50    lower_bound: Option<f32>,
51    upper_bound: Option<f32>,
52    dist_q_c: f32,
53}
54
55impl From<&Query> for FlatQueryParams {
56    fn from(q: &Query) -> Self {
57        Self {
58            lower_bound: q.lower_bound,
59            upper_bound: q.upper_bound,
60            dist_q_c: q.dist_q_c,
61        }
62    }
63}
64
65impl IvfSubIndex for FlatIndex {
66    type QueryParams = FlatQueryParams;
67    type BuildParams = ();
68
69    fn name() -> &'static str {
70        "FLAT"
71    }
72
73    fn metadata_key() -> &'static str {
74        "lance:flat"
75    }
76
77    fn schema() -> arrow_schema::SchemaRef {
78        Schema::new(vec![Field::new("__flat_marker", DataType::UInt64, false)]).into()
79    }
80
81    fn search(
82        &self,
83        query: ArrayRef,
84        k: usize,
85        params: Self::QueryParams,
86        storage: &impl VectorStore,
87        prefilter: Arc<dyn PreFilter>,
88        metrics: &dyn MetricsCollector,
89    ) -> Result<RecordBatch> {
90        let is_range_query = params.lower_bound.is_some() || params.upper_bound.is_some();
91        let row_ids = storage.row_ids();
92        let dist_calc = storage.dist_calculator(query, params.dist_q_c);
93        let mut res = BinaryHeap::with_capacity(k);
94        metrics.record_comparisons(storage.len());
95
96        match prefilter.is_empty() {
97            true => {
98                let dists = dist_calc.distance_all(k);
99
100                if is_range_query {
101                    let lower_bound = params.lower_bound.unwrap_or(f32::MIN).into();
102                    let upper_bound = params.upper_bound.unwrap_or(f32::MAX).into();
103
104                    for (&row_id, dist) in row_ids.zip(dists) {
105                        let dist = dist.into();
106                        if dist < lower_bound || dist >= upper_bound {
107                            continue;
108                        }
109                        if res.len() < k {
110                            res.push(OrderedNode::new(row_id, dist));
111                        } else if res.peek().unwrap().dist > dist {
112                            res.pop();
113                            res.push(OrderedNode::new(row_id, dist));
114                        }
115                    }
116                } else {
117                    for (&row_id, dist) in row_ids.zip(dists) {
118                        let dist = dist.into();
119                        if res.len() < k {
120                            res.push(OrderedNode::new(row_id, dist));
121                        } else if res.peek().unwrap().dist > dist {
122                            res.pop();
123                            res.push(OrderedNode::new(row_id, dist));
124                        }
125                    }
126                }
127            }
128            false => {
129                let row_addr_mask = prefilter.mask();
130                if is_range_query {
131                    let lower_bound = params.lower_bound.unwrap_or(f32::MIN).into();
132                    let upper_bound = params.upper_bound.unwrap_or(f32::MAX).into();
133                    for (id, &row_addr) in row_ids.enumerate() {
134                        if !row_addr_mask.selected(row_addr) {
135                            continue;
136                        }
137                        let dist = dist_calc.distance(id as u32).into();
138                        if dist < lower_bound || dist >= upper_bound {
139                            continue;
140                        }
141
142                        if res.len() < k {
143                            res.push(OrderedNode::new(row_addr, dist));
144                        } else if res.peek().unwrap().dist > dist {
145                            res.pop();
146                            res.push(OrderedNode::new(row_addr, dist));
147                        }
148                    }
149                } else {
150                    for (id, &row_addr) in row_ids.enumerate() {
151                        if !row_addr_mask.selected(row_addr) {
152                            continue;
153                        }
154
155                        let dist = dist_calc.distance(id as u32).into();
156                        if res.len() < k {
157                            res.push(OrderedNode::new(row_addr, dist));
158                        } else if res.peek().unwrap().dist > dist {
159                            res.pop();
160                            res.push(OrderedNode::new(row_addr, dist));
161                        }
162                    }
163                }
164            }
165        };
166
167        // we don't need to sort the results by distances here
168        // because there's a SortExec node in the query plan which sorts the results from all partitions
169        let (row_ids, dists): (Vec<_>, Vec<_>) = res.into_iter().map(|r| (r.id, r.dist.0)).unzip();
170        let (row_ids, dists) = (UInt64Array::from(row_ids), Float32Array::from(dists));
171
172        Ok(RecordBatch::try_new(
173            ANN_SEARCH_SCHEMA.clone(),
174            vec![Arc::new(dists), Arc::new(row_ids)],
175        )?)
176    }
177
178    fn load(_: RecordBatch) -> Result<Self> {
179        Ok(Self {})
180    }
181
182    fn index_vectors(_: &impl VectorStore, _: Self::BuildParams) -> Result<Self>
183    where
184        Self: Sized,
185    {
186        Ok(Self {})
187    }
188
189    fn remap(&self, _: &HashMap<u64, Option<u64>>, _: &impl VectorStore) -> Result<Self> {
190        Ok(self.clone())
191    }
192
193    fn to_batch(&self) -> Result<RecordBatch> {
194        Ok(RecordBatch::new_empty(Schema::empty().into()))
195    }
196}
197
198#[derive(Debug, Clone, Serialize, Deserialize, DeepSizeOf)]
199pub struct FlatMetadata {
200    pub dim: usize,
201}
202
203#[async_trait::async_trait]
204impl QuantizerMetadata for FlatMetadata {
205    async fn load(_: &PreviousFileReader) -> Result<Self> {
206        unimplemented!("Flat will be used in new index builder which doesn't require this")
207    }
208}
209
210#[derive(Debug, Clone, DeepSizeOf)]
211pub struct FlatQuantizer {
212    dim: usize,
213    distance_type: DistanceType,
214}
215
216impl FlatQuantizer {
217    pub fn new(dim: usize, distance_type: DistanceType) -> Self {
218        Self { dim, distance_type }
219    }
220}
221
222impl Quantization for FlatQuantizer {
223    type BuildParams = ();
224    type Metadata = FlatMetadata;
225    type Storage = FlatFloatStorage;
226
227    fn build(data: &dyn Array, distance_type: DistanceType, _: &Self::BuildParams) -> Result<Self> {
228        let dim = data.as_fixed_size_list().value_length();
229        Ok(Self::new(dim as usize, distance_type))
230    }
231
232    fn retrain(&mut self, _: &dyn Array) -> Result<()> {
233        Ok(())
234    }
235
236    fn code_dim(&self) -> usize {
237        self.dim
238    }
239
240    fn column(&self) -> &'static str {
241        FLAT_COLUMN
242    }
243
244    fn from_metadata(metadata: &Self::Metadata, distance_type: DistanceType) -> Result<Quantizer> {
245        Ok(Quantizer::Flat(Self {
246            dim: metadata.dim,
247            distance_type,
248        }))
249    }
250
251    fn metadata(&self, _: Option<crate::vector::quantizer::QuantizationMetadata>) -> FlatMetadata {
252        FlatMetadata { dim: self.dim }
253    }
254
255    fn metadata_key() -> &'static str {
256        "flat"
257    }
258
259    fn quantization_type() -> QuantizationType {
260        QuantizationType::Flat
261    }
262
263    fn quantize(&self, vectors: &dyn Array) -> Result<ArrayRef> {
264        Ok(vectors.slice(0, vectors.len()))
265    }
266
267    fn field(&self) -> Field {
268        Field::new(
269            FLAT_COLUMN,
270            DataType::FixedSizeList(
271                Arc::new(Field::new("item", DataType::Float32, true)),
272                self.dim as i32,
273            ),
274            true,
275        )
276    }
277}
278
279impl From<FlatQuantizer> for Quantizer {
280    fn from(value: FlatQuantizer) -> Self {
281        Self::Flat(value)
282    }
283}
284
285impl TryFrom<Quantizer> for FlatQuantizer {
286    type Error = Error;
287
288    fn try_from(value: Quantizer) -> Result<Self> {
289        match value {
290            Quantizer::Flat(quantizer) => Ok(quantizer),
291            _ => Err(Error::invalid_input("quantizer is not FlatQuantizer")),
292        }
293    }
294}
295
296#[derive(Debug, Clone, DeepSizeOf)]
297pub struct FlatBinQuantizer {
298    dim: usize,
299    distance_type: DistanceType,
300}
301
302impl FlatBinQuantizer {
303    pub fn new(dim: usize, distance_type: DistanceType) -> Self {
304        Self { dim, distance_type }
305    }
306}
307
308impl Quantization for FlatBinQuantizer {
309    type BuildParams = ();
310    type Metadata = FlatMetadata;
311    type Storage = FlatBinStorage;
312
313    fn build(data: &dyn Array, distance_type: DistanceType, _: &Self::BuildParams) -> Result<Self> {
314        let dim = data.as_fixed_size_list().value_length();
315        Ok(Self::new(dim as usize, distance_type))
316    }
317
318    fn retrain(&mut self, _: &dyn Array) -> Result<()> {
319        Ok(())
320    }
321
322    fn code_dim(&self) -> usize {
323        self.dim
324    }
325
326    fn column(&self) -> &'static str {
327        FLAT_COLUMN
328    }
329
330    fn from_metadata(metadata: &Self::Metadata, distance_type: DistanceType) -> Result<Quantizer> {
331        Ok(Quantizer::FlatBin(Self {
332            dim: metadata.dim,
333            distance_type,
334        }))
335    }
336
337    fn metadata(&self, _: Option<crate::vector::quantizer::QuantizationMetadata>) -> FlatMetadata {
338        FlatMetadata { dim: self.dim }
339    }
340
341    fn metadata_key() -> &'static str {
342        "flat"
343    }
344
345    fn quantization_type() -> QuantizationType {
346        QuantizationType::Flat
347    }
348
349    fn quantize(&self, vectors: &dyn Array) -> Result<ArrayRef> {
350        Ok(vectors.slice(0, vectors.len()))
351    }
352
353    fn field(&self) -> Field {
354        Field::new(
355            FLAT_COLUMN,
356            DataType::FixedSizeList(
357                Arc::new(Field::new("item", DataType::UInt8, true)),
358                self.dim as i32,
359            ),
360            true,
361        )
362    }
363}
364
365impl From<FlatBinQuantizer> for Quantizer {
366    fn from(value: FlatBinQuantizer) -> Self {
367        Self::FlatBin(value)
368    }
369}
370
371impl TryFrom<Quantizer> for FlatBinQuantizer {
372    type Error = Error;
373
374    fn try_from(value: Quantizer) -> Result<Self> {
375        match value {
376            Quantizer::FlatBin(quantizer) => Ok(quantizer),
377            _ => Err(Error::invalid_input("quantizer is not FlatBinQuantizer")),
378        }
379    }
380}
lance_index/vector/flat/index.rs

lance_index/vector/flat/
index.rs
