lance_index/vector/hnsw/
index.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use lance_core::utils::row_addr_remap::RowAddrRemap;
5use std::{
6    any::Any,
7    fmt::{Debug, Formatter},
8    sync::Arc,
9};
10
11use arrow_array::{Float32Array, RecordBatch, UInt32Array};
12use async_trait::async_trait;
13use datafusion::execution::SendableRecordBatchStream;
14use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
15use lance_arrow::RecordBatchExt;
16use lance_core::ROW_ID;
17use lance_core::deepsize::DeepSizeOf;
18use lance_core::{Error, Result, datatypes::Schema};
19use lance_file::previous::reader::FileReader as PreviousFileReader;
20use lance_io::traits::Reader;
21use lance_linalg::distance::DistanceType;
22use lance_table::format::SelfDescribingFileReader;
23use roaring::RoaringBitmap;
24use serde_json::json;
25use tracing::instrument;
26
27use crate::vector::ivf::storage::IvfModel;
28use crate::vector::quantizer::QuantizationType;
29use crate::vector::v3::subindex::{IvfSubIndex, SubIndexType};
30use crate::{
31    Index, IndexType,
32    vector::{
33        Query, VectorIndex,
34        graph::NEIGHBORS_FIELD,
35        hnsw::{HNSW, HnswMetadata, VECTOR_ID_FIELD},
36        ivf::storage::IVF_PARTITION_KEY,
37        quantizer::{IvfQuantizationStorage, Quantization, Quantizer},
38        storage::VectorStore,
39    },
40};
41use crate::{metrics::MetricsCollector, prefilter::PreFilter};
42
43#[derive(Clone, DeepSizeOf)]
44pub struct HNSWIndexOptions {
45    pub use_residual: bool,
46}
47
48#[derive(Clone, DeepSizeOf)]
49pub struct HNSWIndex<Q: Quantization> {
50    // Some(T) if the index is loaded, None otherwise
51    hnsw: Option<HNSW>,
52    storage: Option<Arc<Q::Storage>>,
53
54    // TODO: move these into IVFIndex after the refactor is complete
55    partition_storage: IvfQuantizationStorage<Q>,
56    partition_metadata: Option<Vec<HnswMetadata>>,
57
58    options: HNSWIndexOptions,
59}
60
61impl<Q: Quantization> Debug for HNSWIndex<Q> {
62    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
63        self.hnsw.fmt(f)
64    }
65}
66
67impl<Q: Quantization> HNSWIndex<Q> {
68    pub async fn try_new(
69        reader: Arc<dyn Reader>,
70        aux_reader: Arc<dyn Reader>,
71        options: HNSWIndexOptions,
72    ) -> Result<Self> {
73        let reader =
74            PreviousFileReader::try_new_self_described_from_reader(reader.clone(), None).await?;
75
76        let partition_metadata = match reader.schema().metadata.get(IVF_PARTITION_KEY) {
77            Some(json) => {
78                let metadata: Vec<HnswMetadata> = serde_json::from_str(json)?;
79                Some(metadata)
80            }
81            None => None,
82        };
83
84        let ivf_store = IvfQuantizationStorage::open(aux_reader).await?;
85        Ok(Self {
86            hnsw: None,
87            storage: None,
88            partition_storage: ivf_store,
89            partition_metadata,
90            options,
91        })
92    }
93
94    pub fn quantizer(&self) -> &Quantizer {
95        self.partition_storage.quantizer()
96    }
97
98    pub fn metadata(&self) -> HnswMetadata {
99        self.partition_metadata.as_ref().unwrap()[0].clone()
100    }
101
102    fn get_partition_metadata(&self, partition_id: usize) -> Result<HnswMetadata> {
103        match self.partition_metadata {
104            Some(ref metadata) => Ok(metadata[partition_id].clone()),
105            None => Err(Error::index("No partition metadata found".to_string())),
106        }
107    }
108}
109
110#[async_trait]
111impl<Q: Quantization + Send + Sync + 'static> Index for HNSWIndex<Q> {
112    /// Cast to [Any].
113    fn as_any(&self) -> &dyn Any {
114        self
115    }
116
117    /// Cast to [Index]
118    fn as_index(self: Arc<Self>) -> Arc<dyn Index> {
119        self
120    }
121
122    /// Retrieve index statistics as a JSON Value
123    fn statistics(&self) -> Result<serde_json::Value> {
124        Ok(json!({
125            "index_type": "HNSW",
126            "distance_type": self.partition_storage.distance_type().to_string(),
127        }))
128    }
129
130    async fn prewarm(&self) -> Result<()> {
131        // TODO: HNSW can (and should) support pre-warming
132        Ok(())
133    }
134
135    /// Get the type of the index
136    fn index_type(&self) -> IndexType {
137        IndexType::Vector
138    }
139
140    /// Read through the index and determine which fragment ids are covered by the index
141    ///
142    /// This is a kind of slow operation.  It's better to use the fragment_bitmap.  This
143    /// only exists for cases where the fragment_bitmap has become corrupted or missing.
144    async fn calculate_included_frags(&self) -> Result<RoaringBitmap> {
145        unimplemented!()
146    }
147}
148
149#[async_trait]
150impl<Q: Quantization + Send + Sync + 'static> VectorIndex for HNSWIndex<Q> {
151    #[instrument(level = "debug", skip_all, name = "HNSWIndex::search")]
152    async fn search(
153        &self,
154        query: &Query,
155        pre_filter: Arc<dyn PreFilter>,
156        metrics: &dyn MetricsCollector,
157    ) -> Result<RecordBatch> {
158        let hnsw = self
159            .hnsw
160            .as_ref()
161            .ok_or(Error::index("HNSW index not loaded".to_string()))?;
162
163        let storage = self
164            .storage
165            .as_ref()
166            .ok_or(Error::index("vector storage not loaded".to_string()))?;
167
168        let refine_factor = query.refine_factor.unwrap_or(1) as usize;
169        let k = query.k * refine_factor;
170
171        hnsw.search(
172            query.key.clone(),
173            k,
174            query.into(),
175            storage.as_ref(),
176            pre_filter,
177            metrics,
178        )
179    }
180
181    fn find_partitions(&self, _: &Query) -> Result<(UInt32Array, Float32Array)> {
182        unimplemented!("only for IVF")
183    }
184
185    fn total_partitions(&self) -> usize {
186        1
187    }
188
189    async fn search_in_partition(
190        &self,
191        _: usize,
192        _: &Query,
193        _: Arc<dyn PreFilter>,
194        _: &dyn MetricsCollector,
195    ) -> Result<RecordBatch> {
196        unimplemented!("only for IVF")
197    }
198
199    fn is_loadable(&self) -> bool {
200        true
201    }
202
203    fn use_residual(&self) -> bool {
204        self.options.use_residual
205    }
206
207    async fn load(
208        &self,
209        reader: Arc<dyn Reader>,
210        _offset: usize,
211        _length: usize,
212    ) -> Result<Box<dyn VectorIndex>> {
213        let schema = Schema::try_from(&arrow_schema::Schema::new(vec![
214            NEIGHBORS_FIELD.clone(),
215            VECTOR_ID_FIELD.clone(),
216        ]))?;
217
218        let reader = PreviousFileReader::try_new_from_reader(
219            reader.path(),
220            reader.clone(),
221            None,
222            schema,
223            0,
224            0,
225            2,
226            None,
227        )
228        .await?;
229
230        let storage = Arc::new(self.partition_storage.load_partition(0).await?);
231        let batch = reader.read_range(0..reader.len(), reader.schema()).await?;
232        let hnsw = HNSW::load(batch)?;
233
234        Ok(Box::new(Self {
235            hnsw: Some(hnsw),
236            storage: Some(storage),
237            partition_storage: self.partition_storage.clone(),
238            partition_metadata: self.partition_metadata.clone(),
239            options: self.options.clone(),
240        }))
241    }
242
243    async fn load_partition(
244        &self,
245        reader: Arc<dyn Reader>,
246        offset: usize,
247        length: usize,
248        partition_id: usize,
249    ) -> Result<Box<dyn VectorIndex>> {
250        let reader = PreviousFileReader::try_new_self_described_from_reader(reader, None).await?;
251
252        let metadata = self.get_partition_metadata(partition_id)?;
253        let storage = Arc::new(self.partition_storage.load_partition(partition_id).await?);
254        let batch = reader
255            .read_range(offset..offset + length, reader.schema())
256            .await?;
257        let mut schema = batch.schema_ref().as_ref().clone();
258        schema.metadata.insert(
259            HNSW::metadata_key().to_owned(),
260            serde_json::to_string(&metadata)?,
261        );
262        let batch = batch.with_schema(schema.into())?;
263        let hnsw = HNSW::load(batch)?;
264
265        Ok(Box::new(Self {
266            hnsw: Some(hnsw),
267            storage: Some(storage),
268            partition_storage: self.partition_storage.clone(),
269            partition_metadata: self.partition_metadata.clone(),
270            options: self.options.clone(),
271        }))
272    }
273
274    async fn to_batch_stream(&self, with_vector: bool) -> Result<SendableRecordBatchStream> {
275        let store = self
276            .storage
277            .as_ref()
278            .ok_or(Error::index("vector storage not loaded".to_string()))?;
279
280        let schema = if with_vector {
281            store.schema().clone()
282        } else {
283            let schema = store.schema();
284            let row_id_idx = schema.index_of(ROW_ID)?;
285            Arc::new(schema.project(&[row_id_idx])?)
286        };
287
288        let batches = store
289            .to_batches()?
290            .map(|b| {
291                let batch = b.project_by_schema(&schema)?;
292                Ok(batch)
293            })
294            .collect::<Vec<_>>();
295        let stream = futures::stream::iter(batches);
296        let stream = RecordBatchStreamAdapter::new(schema, stream);
297        Ok(Box::pin(stream))
298    }
299
300    fn num_rows(&self) -> u64 {
301        self.hnsw
302            .as_ref()
303            .map_or(0, |hnsw| hnsw.num_nodes(0) as u64)
304    }
305
306    fn row_ids(&self) -> Box<dyn Iterator<Item = &'_ u64> + '_> {
307        Box::new(self.storage.as_ref().unwrap().row_ids())
308    }
309
310    async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> {
311        Err(Error::index(
312            "Remapping HNSW in this way not supported".to_string(),
313        ))
314    }
315
316    fn ivf_model(&self) -> &IvfModel {
317        unimplemented!("only for IVF")
318    }
319
320    fn quantizer(&self) -> Quantizer {
321        self.partition_storage.quantizer().clone()
322    }
323
324    fn partition_size(&self, _: usize) -> usize {
325        unimplemented!("only for IVF")
326    }
327
328    fn sub_index_type(&self) -> (SubIndexType, QuantizationType) {
329        (
330            SubIndexType::Hnsw,
331            self.partition_storage.quantizer().quantization_type(),
332        )
333    }
334
335    fn metric_type(&self) -> DistanceType {
336        self.partition_storage.distance_type()
337    }
338}
lance_index/vector/hnsw/index.rs

lance_index/vector/hnsw/
index.rs
