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

lance_index/vector/
storage.rs
