lance_index/vector/
ivf.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! IVF - Inverted File Index
5
6use std::ops::Range;
7use std::sync::Arc;
8
9use arrow_array::{Array, FixedSizeListArray, Float32Array, RecordBatch, UInt32Array};
10
11pub use builder::IvfBuildParams;
12use lance_core::Result;
13use lance_linalg::distance::{DistanceType, MetricType};
14use tracing::instrument;
15
16use crate::vector::bq::builder::RabitQuantizer;
17use crate::vector::bq::transform::RQTransformer;
18use crate::vector::ivf::transform::PartitionTransformer;
19use crate::vector::kmeans::{compute_partitions_arrow_array, kmeans_find_partitions_arrow_array};
20use crate::vector::{pq::ProductQuantizer, transform::Transformer};
21
22use super::flat::transform::FlatTransformer;
23use super::pq::transform::PQTransformer;
24use super::quantizer::Quantization;
25use super::residual::ResidualTransform;
26use super::sq::ScalarQuantizer;
27use super::sq::transform::SQTransformer;
28use super::transform::KeepFiniteVectors;
29use super::{PART_ID_COLUMN, PQ_CODE_COLUMN, SQ_CODE_COLUMN};
30use super::{quantizer::Quantizer, residual::compute_residual};
31
32pub mod builder;
33pub mod shuffler;
34pub mod storage;
35mod transform;
36
37/// Create an IVF from the flatten centroids.
38///
39/// Parameters
40/// ----------
41/// - *centroids*: a flatten floating number array of centroids.
42/// - *dimension*: dimension of the vector.
43/// - *metric_type*: metric type to compute pair-wise vector distance.
44/// - *transforms*: a list of transforms to apply to the vector column.
45/// - *range*: only covers a range of partitions. Default is None
46pub fn new_ivf_transformer(
47    centroids: FixedSizeListArray,
48    metric_type: DistanceType,
49    transforms: Vec<Arc<dyn Transformer>>,
50) -> IvfTransformer {
51    IvfTransformer::new(centroids, metric_type, transforms)
52}
53
54pub fn new_ivf_transformer_with_quantizer(
55    centroids: FixedSizeListArray,
56    metric_type: MetricType,
57    vector_column: &str,
58    quantizer: Quantizer,
59    range: Option<Range<u32>>,
60) -> Result<IvfTransformer> {
61    match quantizer {
62        Quantizer::Flat(_) | Quantizer::FlatBin(_) => Ok(IvfTransformer::new_flat(
63            centroids,
64            metric_type,
65            vector_column,
66            range,
67        )),
68        Quantizer::Product(pq) => Ok(IvfTransformer::with_pq(
69            centroids,
70            metric_type,
71            vector_column,
72            pq,
73            range,
74        )),
75        Quantizer::Scalar(sq) => Ok(IvfTransformer::with_sq(
76            centroids,
77            metric_type,
78            vector_column,
79            sq,
80            range,
81        )),
82        Quantizer::Rabit(rq) => {
83            IvfTransformer::with_rq(centroids, metric_type, vector_column, rq, range)
84        }
85    }
86}
87
88/// IVF - IVF file partition
89///
90#[derive(Debug)]
91pub struct IvfTransformer {
92    /// Centroids of a cluster algorithm, to run IVF.
93    ///
94    /// It is a 2-D `(num_partitions * dimension)` of floating array.
95    centroids: FixedSizeListArray,
96
97    /// Transform applied to each partition.
98    transforms: Vec<Arc<dyn Transformer>>,
99
100    /// Metric type to compute pair-wise vector distance.
101    distance_type: DistanceType,
102}
103
104impl IvfTransformer {
105    /// Create a new Ivf model.
106    pub fn new(
107        centroids: FixedSizeListArray,
108        metric_type: MetricType,
109        transforms: Vec<Arc<dyn Transformer>>,
110    ) -> Self {
111        Self {
112            centroids,
113            distance_type: metric_type,
114            transforms,
115        }
116    }
117
118    pub fn new_partition_transformer(
119        centroids: FixedSizeListArray,
120        distance_type: DistanceType,
121        vector_column: &str,
122    ) -> Self {
123        let mut transforms: Vec<Arc<dyn Transformer>> =
124            vec![Arc::new(super::transform::Flatten::new(vector_column))];
125
126        let distance_type = if distance_type == MetricType::Cosine {
127            transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
128                vector_column,
129            )));
130            MetricType::L2
131        } else {
132            distance_type
133        };
134        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
135
136        let partition_transform = Arc::new(PartitionTransformer::new(
137            centroids.clone(),
138            distance_type,
139            vector_column,
140        ));
141        transforms.push(partition_transform);
142        Self::new(centroids, distance_type, transforms)
143    }
144
145    pub fn new_flat(
146        centroids: FixedSizeListArray,
147        distance_type: DistanceType,
148        vector_column: &str,
149        range: Option<Range<u32>>,
150    ) -> Self {
151        let mut transforms: Vec<Arc<dyn Transformer>> =
152            vec![Arc::new(super::transform::Flatten::new(vector_column))];
153
154        let dt = if distance_type == DistanceType::Cosine {
155            transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
156                vector_column,
157            )));
158            MetricType::L2
159        } else {
160            distance_type
161        };
162        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
163
164        let ivf_transform = Arc::new(PartitionTransformer::new(
165            centroids.clone(),
166            dt,
167            vector_column,
168        ));
169        transforms.push(ivf_transform);
170
171        if let Some(range) = range {
172            transforms.push(Arc::new(transform::PartitionFilter::new(
173                PART_ID_COLUMN,
174                range,
175            )));
176        }
177
178        transforms.push(Arc::new(FlatTransformer::new(vector_column)));
179
180        Self::new(centroids, distance_type, transforms)
181    }
182
183    /// Create a IVF_PQ struct.
184    pub fn with_pq(
185        centroids: FixedSizeListArray,
186        distance_type: DistanceType,
187        vector_column: &str,
188        pq: ProductQuantizer,
189        range: Option<Range<u32>>,
190    ) -> Self {
191        let mut transforms: Vec<Arc<dyn Transformer>> =
192            vec![Arc::new(super::transform::Flatten::new(vector_column))];
193
194        let distance_type = if distance_type == MetricType::Cosine {
195            transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
196                vector_column,
197            )));
198            MetricType::L2
199        } else {
200            distance_type
201        };
202        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
203
204        let partition_transform = Arc::new(PartitionTransformer::new(
205            centroids.clone(),
206            distance_type,
207            vector_column,
208        ));
209        transforms.push(partition_transform);
210
211        if let Some(range) = range {
212            transforms.push(Arc::new(transform::PartitionFilter::new(
213                PART_ID_COLUMN,
214                range,
215            )));
216        }
217
218        if ProductQuantizer::use_residual(distance_type) {
219            transforms.push(Arc::new(ResidualTransform::new(
220                centroids.clone(),
221                PART_ID_COLUMN,
222                vector_column,
223            )));
224        }
225        transforms.push(Arc::new(PQTransformer::new(
226            pq,
227            vector_column,
228            PQ_CODE_COLUMN,
229        )));
230
231        Self::new(centroids, distance_type, transforms)
232    }
233
234    fn with_sq(
235        centroids: FixedSizeListArray,
236        metric_type: MetricType,
237        vector_column: &str,
238        sq: ScalarQuantizer,
239        range: Option<Range<u32>>,
240    ) -> Self {
241        let mut transforms: Vec<Arc<dyn Transformer>> =
242            vec![Arc::new(super::transform::Flatten::new(vector_column))];
243
244        let distance_type = if metric_type == MetricType::Cosine {
245            transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
246                vector_column,
247            )));
248            MetricType::L2
249        } else {
250            metric_type
251        };
252        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
253
254        let partition_transformer = Arc::new(PartitionTransformer::new(
255            centroids.clone(),
256            distance_type,
257            vector_column,
258        ));
259        transforms.push(partition_transformer);
260
261        if let Some(range) = range {
262            transforms.push(Arc::new(transform::PartitionFilter::new(
263                PART_ID_COLUMN,
264                range,
265            )));
266        }
267
268        transforms.push(Arc::new(SQTransformer::new(
269            sq,
270            vector_column.to_owned(),
271            SQ_CODE_COLUMN.to_owned(),
272        )));
273
274        Self::new(centroids, distance_type, transforms)
275    }
276
277    fn with_rq(
278        centroids: FixedSizeListArray,
279        distance_type: DistanceType,
280        vector_column: &str,
281        rq: RabitQuantizer,
282        range: Option<Range<u32>>,
283    ) -> Result<Self> {
284        let mut transforms: Vec<Arc<dyn Transformer>> =
285            vec![Arc::new(super::transform::Flatten::new(vector_column))];
286
287        let distance_type = if distance_type == MetricType::Cosine {
288            transforms.push(Arc::new(super::transform::NormalizeTransformer::new(
289                vector_column,
290            )));
291            MetricType::L2
292        } else {
293            distance_type
294        };
295        transforms.push(Arc::new(KeepFiniteVectors::new(vector_column)));
296
297        let partition_transform = Arc::new(
298            PartitionTransformer::new(centroids.clone(), distance_type, vector_column)
299                .with_distance(true),
300        );
301        transforms.push(partition_transform);
302
303        if let Some(range) = range {
304            transforms.push(Arc::new(transform::PartitionFilter::new(
305                PART_ID_COLUMN,
306                range,
307            )));
308        }
309
310        transforms.push(Arc::new(ResidualTransform::new(
311            centroids.clone(),
312            PART_ID_COLUMN,
313            vector_column,
314        )));
315
316        transforms.push(Arc::new(RQTransformer::new(
317            rq,
318            distance_type,
319            centroids.clone(),
320            vector_column,
321        )?));
322
323        Ok(Self::new(centroids, distance_type, transforms))
324    }
325
326    #[inline]
327    pub fn compute_residual(&self, data: &FixedSizeListArray) -> Result<FixedSizeListArray> {
328        compute_residual(&self.centroids, data, Some(self.distance_type), None)
329    }
330
331    #[inline]
332    pub fn compute_partitions(&self, data: &FixedSizeListArray) -> Result<UInt32Array> {
333        Ok(
334            compute_partitions_arrow_array(&self.centroids, data, self.distance_type)
335                .map(|(part_ids, _)| part_ids.into())?,
336        )
337    }
338
339    pub fn find_partitions(
340        &self,
341        query: &dyn Array,
342        nprobes: usize,
343    ) -> Result<(UInt32Array, Float32Array)> {
344        Ok(kmeans_find_partitions_arrow_array(
345            &self.centroids,
346            query,
347            nprobes,
348            self.distance_type,
349        )?)
350    }
351}
352
353impl Transformer for IvfTransformer {
354    #[instrument(name = "IvfTransformer::transform", level = "debug", skip_all)]
355    fn transform(&self, batch: &RecordBatch) -> Result<RecordBatch> {
356        let mut batch = batch.clone();
357        for transform in self.transforms.as_slice() {
358            batch = transform.transform(&batch)?;
359        }
360        Ok(batch)
361    }
362}
lance_index/vector/ivf.rs

lance_index/vector/
ivf.rs
