lance_index/vector/
storage.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Vector Storage, holding (quantized) vectors and providing distance calculation.
5
6use crate::vector::quantizer::QuantizerStorage;
7use arrow::compute::concat_batches;
8use arrow_array::{ArrayRef, RecordBatch};
9use arrow_schema::SchemaRef;
10use deepsize::DeepSizeOf;
11use futures::prelude::stream::TryStreamExt;
12use lance_arrow::RecordBatchExt;
13use lance_core::{Error, ROW_ID, Result};
14use lance_encoding::decoder::FilterExpression;
15use lance_file::reader::FileReader;
16use lance_io::ReadBatchParams;
17use lance_linalg::distance::DistanceType;
18use prost::Message;
19use std::{any::Any, sync::Arc};
20
21use crate::frag_reuse::FragReuseIndex;
22use crate::{
23    pb,
24    vector::{
25        ivf::storage::{IVF_METADATA_KEY, IvfModel},
26        quantizer::Quantization,
27    },
28};
29
30use super::DISTANCE_TYPE_KEY;
31use super::quantizer::{Quantizer, QuantizerMetadata};
32
33/// <section class="warning">
34///  Internal API
35///
36///  API stability is not guaranteed
37/// </section>
38pub trait DistCalculator {
39    fn distance(&self, id: u32) -> f32;
40
41    // return the distances of all rows
42    // k_hint is a hint that can be used for optimization
43    fn distance_all(&self, k_hint: usize) -> Vec<f32>;
44
45    fn prefetch(&self, _id: u32) {}
46}
47
48pub const STORAGE_METADATA_KEY: &str = "storage_metadata";
49
50/// Vector Storage is the abstraction to store the vectors.
51///
52/// It can be in-memory or on-disk, raw vector or quantized vectors.
53///
54/// It abstracts away the logic to compute the distance between vectors.
55///
56/// TODO: should we rename this to "VectorDistance"?;
57///
58/// <section class="warning">
59///  Internal API
60///
61///  API stability is not guaranteed
62/// </section>
63pub trait VectorStore: Send + Sync + Sized + Clone {
64    type DistanceCalculator<'a>: DistCalculator
65    where
66        Self: 'a;
67
68    fn as_any(&self) -> &dyn Any;
69
70    fn schema(&self) -> &SchemaRef;
71
72    fn to_batches(&self) -> Result<impl Iterator<Item = RecordBatch> + Send>;
73
74    fn len(&self) -> usize;
75
76    /// Returns true if this graph is empty.
77    fn is_empty(&self) -> bool {
78        self.len() == 0
79    }
80
81    /// Return [DistanceType].
82    fn distance_type(&self) -> DistanceType;
83
84    /// Get the lance ROW ID from one vector.
85    fn row_id(&self, id: u32) -> u64;
86
87    fn row_ids(&self) -> impl Iterator<Item = &u64>;
88
89    /// Append Raw [RecordBatch] into the Storage.
90    /// The storage implement will perform quantization if necessary.
91    fn append_batch(&self, batch: RecordBatch, vector_column: &str) -> Result<Self>;
92
93    /// Create a [DistCalculator] to compute the distance between the query.
94    ///
95    /// Using dist calculator can be more efficient as it can pre-compute some
96    /// values.
97    fn dist_calculator(&self, query: ArrayRef, dist_q_c: f32) -> Self::DistanceCalculator<'_>;
98
99    fn dist_calculator_from_id(&self, id: u32) -> Self::DistanceCalculator<'_>;
100
101    fn dist_between(&self, u: u32, v: u32) -> f32 {
102        let dist_cal_u = self.dist_calculator_from_id(u);
103        dist_cal_u.distance(v)
104    }
105}
106
107pub struct StorageBuilder<Q: Quantization> {
108    vector_column: String,
109    distance_type: DistanceType,
110    quantizer: Q,
111
112    frag_reuse_index: Option<Arc<FragReuseIndex>>,
113}
114
115impl<Q: Quantization> StorageBuilder<Q> {
116    pub fn new(
117        vector_column: String,
118        distance_type: DistanceType,
119        quantizer: Q,
120        frag_reuse_index: Option<Arc<FragReuseIndex>>,
121    ) -> Result<Self> {
122        Ok(Self {
123            vector_column,
124            distance_type,
125            quantizer,
126            frag_reuse_index,
127        })
128    }
129
130    pub fn build(&self, batches: Vec<RecordBatch>) -> Result<Q::Storage> {
131        let mut batch = concat_batches(batches[0].schema_ref(), batches.iter())?;
132
133        if batch.column_by_name(self.quantizer.column()).is_none() {
134            let vectors = batch
135                .column_by_name(&self.vector_column)
136                .ok_or(Error::index(format!(
137                    "Vector column {} not found in batch",
138                    self.vector_column
139                )))?;
140            let codes = self.quantizer.quantize(vectors)?;
141            batch = batch.drop_column(&self.vector_column)?.try_with_column(
142                arrow_schema::Field::new(self.quantizer.column(), codes.data_type().clone(), true),
143                codes,
144            )?;
145        }
146
147        debug_assert!(batch.column_by_name(ROW_ID).is_some());
148        debug_assert!(batch.column_by_name(self.quantizer.column()).is_some());
149
150        Q::Storage::try_from_batch(
151            batch,
152            &self.quantizer.metadata(None),
153            self.distance_type,
154            self.frag_reuse_index.clone(),
155        )
156    }
157}
158
159/// Loader to load partitioned PQ storage from disk.
160#[derive(Debug)]
161pub struct IvfQuantizationStorage<Q: Quantization> {
162    reader: FileReader,
163
164    distance_type: DistanceType,
165    metadata: Q::Metadata,
166
167    ivf: IvfModel,
168    frag_reuse_index: Option<Arc<FragReuseIndex>>,
169}
170
171impl<Q: Quantization> DeepSizeOf for IvfQuantizationStorage<Q> {
172    fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
173        self.metadata.deep_size_of_children(context) + self.ivf.deep_size_of_children(context)
174    }
175}
176
177impl<Q: Quantization> IvfQuantizationStorage<Q> {
178    /// Open a Loader.
179    ///
180    ///
181    pub async fn try_new(
182        reader: FileReader,
183        frag_reuse_index: Option<Arc<FragReuseIndex>>,
184    ) -> Result<Self> {
185        let schema = reader.schema();
186
187        let distance_type = DistanceType::try_from(
188            schema
189                .metadata
190                .get(DISTANCE_TYPE_KEY)
191                .ok_or(Error::index(format!("{} not found", DISTANCE_TYPE_KEY)))?
192                .as_str(),
193        )?;
194
195        let ivf_pos = schema
196            .metadata
197            .get(IVF_METADATA_KEY)
198            .ok_or(Error::index(format!("{} not found", IVF_METADATA_KEY)))?
199            .parse()
200            .map_err(|e| Error::index(format!("Failed to decode IVF metadata: {}", e)))?;
201        let ivf_bytes = reader.read_global_buffer(ivf_pos).await?;
202        let ivf = IvfModel::try_from(pb::Ivf::decode(ivf_bytes)?)?;
203
204        let mut metadata: Vec<String> = serde_json::from_str(
205            schema
206                .metadata
207                .get(STORAGE_METADATA_KEY)
208                .ok_or(Error::index(format!("{} not found", STORAGE_METADATA_KEY)))?
209                .as_str(),
210        )?;
211        debug_assert_eq!(metadata.len(), 1);
212        // for now the metadata is the same for all partitions, so we just store one
213        let metadata = metadata
214            .pop()
215            .ok_or(Error::index("metadata is empty".to_string()))?;
216        let mut metadata: Q::Metadata = serde_json::from_str(&metadata)?;
217        // we store large metadata (e.g. PQ codebook) in global buffer,
218        // and the schema metadata just contains a pointer to the buffer
219        if let Some(pos) = metadata.buffer_index() {
220            let bytes = reader.read_global_buffer(pos).await?;
221            metadata.parse_buffer(bytes)?;
222        }
223
224        Ok(Self {
225            reader,
226            distance_type,
227            metadata,
228            ivf,
229            frag_reuse_index,
230        })
231    }
232
233    pub fn num_rows(&self) -> u64 {
234        self.reader.num_rows()
235    }
236
237    pub fn partition_size(&self, part_id: usize) -> usize {
238        self.ivf.partition_size(part_id)
239    }
240
241    pub fn quantizer(&self) -> Result<Quantizer> {
242        let metadata = self.metadata();
243        Q::from_metadata(metadata, self.distance_type)
244    }
245
246    pub fn metadata(&self) -> &Q::Metadata {
247        &self.metadata
248    }
249
250    pub fn distance_type(&self) -> DistanceType {
251        self.distance_type
252    }
253
254    pub fn schema(&self) -> SchemaRef {
255        Arc::new(self.reader.schema().as_ref().into())
256    }
257
258    /// Get the number of partitions in the storage.
259    pub fn num_partitions(&self) -> usize {
260        self.ivf.num_partitions()
261    }
262
263    pub async fn load_partition(&self, part_id: usize) -> Result<Q::Storage> {
264        let range = self.ivf.row_range(part_id);
265        let batch = if range.is_empty() {
266            let schema = self.reader.schema();
267            let arrow_schema = arrow_schema::Schema::from(schema.as_ref());
268            RecordBatch::new_empty(Arc::new(arrow_schema))
269        } else {
270            let batches = self
271                .reader
272                .read_stream(
273                    ReadBatchParams::Range(range),
274                    u32::MAX,
275                    1,
276                    FilterExpression::no_filter(),
277                )?
278                .try_collect::<Vec<_>>()
279                .await?;
280            let schema = Arc::new(self.reader.schema().as_ref().into());
281            concat_batches(&schema, batches.iter())?
282        };
283        Q::Storage::try_from_batch(
284            batch,
285            self.metadata(),
286            self.distance_type,
287            self.frag_reuse_index.clone(),
288        )
289    }
290}
lance_index/vector/storage.rs

lance_index/vector/
storage.rs
