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::versions::v1::reader::FileReader as V1FileReader;
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 = V1FileReader::try_new_self_described_from_reader(reader.clone(), None).await?;
74
75        let partition_metadata = match reader.schema().metadata.get(IVF_PARTITION_KEY) {
76            Some(json) => {
77                let metadata: Vec<HnswMetadata> = serde_json::from_str(json)?;
78                Some(metadata)
79            }
80            None => None,
81        };
82
83        let ivf_store = IvfQuantizationStorage::open(aux_reader).await?;
84        Ok(Self {
85            hnsw: None,
86            storage: None,
87            partition_storage: ivf_store,
88            partition_metadata,
89            options,
90        })
91    }
92
93    pub fn quantizer(&self) -> &Quantizer {
94        self.partition_storage.quantizer()
95    }
96
97    pub fn metadata(&self) -> HnswMetadata {
98        self.partition_metadata.as_ref().unwrap()[0].clone()
99    }
100
101    fn get_partition_metadata(&self, partition_id: usize) -> Result<HnswMetadata> {
102        match self.partition_metadata {
103            Some(ref metadata) => Ok(metadata[partition_id].clone()),
104            None => Err(Error::index("No partition metadata found".to_string())),
105        }
106    }
107}
108
109#[async_trait]
110impl<Q: Quantization + Send + Sync + 'static> Index for HNSWIndex<Q> {
111    /// Cast to [Any].
112    fn as_any(&self) -> &dyn Any {
113        self
114    }
115
116    /// Cast to [Index]
117    fn as_index(self: Arc<Self>) -> Arc<dyn Index> {
118        self
119    }
120
121    /// Retrieve index statistics as a JSON Value
122    fn statistics(&self) -> Result<serde_json::Value> {
123        Ok(json!({
124            "index_type": "HNSW",
125            "distance_type": self.partition_storage.distance_type().to_string(),
126        }))
127    }
128
129    async fn prewarm(&self) -> Result<()> {
130        // TODO: HNSW can (and should) support pre-warming
131        Ok(())
132    }
133
134    /// Get the type of the index
135    fn index_type(&self) -> IndexType {
136        IndexType::Vector
137    }
138
139    /// Read through the index and determine which fragment ids are covered by the index
140    ///
141    /// This is a kind of slow operation.  It's better to use the fragment_bitmap.  This
142    /// only exists for cases where the fragment_bitmap has become corrupted or missing.
143    async fn calculate_included_frags(&self) -> Result<RoaringBitmap> {
144        unimplemented!()
145    }
146}
147
148#[async_trait]
149impl<Q: Quantization + Send + Sync + 'static> VectorIndex for HNSWIndex<Q> {
150    #[instrument(level = "debug", skip_all, name = "HNSWIndex::search")]
151    async fn search(
152        &self,
153        query: &Query,
154        pre_filter: Arc<dyn PreFilter>,
155        metrics: &dyn MetricsCollector,
156    ) -> Result<RecordBatch> {
157        let hnsw = self
158            .hnsw
159            .as_ref()
160            .ok_or(Error::index("HNSW index not loaded".to_string()))?;
161
162        let storage = self
163            .storage
164            .as_ref()
165            .ok_or(Error::index("vector storage not loaded".to_string()))?;
166
167        let refine_factor = query.refine_factor.unwrap_or(1) as usize;
168        let k = query.k * refine_factor;
169
170        hnsw.search(
171            query.key.clone(),
172            k,
173            query.into(),
174            storage.as_ref(),
175            pre_filter,
176            metrics,
177        )
178    }
179
180    fn find_partitions(&self, _: &Query) -> Result<(UInt32Array, Float32Array)> {
181        unimplemented!("only for IVF")
182    }
183
184    fn total_partitions(&self) -> usize {
185        1
186    }
187
188    async fn search_in_partition(
189        &self,
190        _: usize,
191        _: &Query,
192        _: Arc<dyn PreFilter>,
193        _: &dyn MetricsCollector,
194    ) -> Result<RecordBatch> {
195        unimplemented!("only for IVF")
196    }
197
198    fn is_loadable(&self) -> bool {
199        true
200    }
201
202    fn use_residual(&self) -> bool {
203        self.options.use_residual
204    }
205
206    async fn load(
207        &self,
208        reader: Arc<dyn Reader>,
209        _offset: usize,
210        _length: usize,
211    ) -> Result<Box<dyn VectorIndex>> {
212        let schema = Schema::try_from(&arrow_schema::Schema::new(vec![
213            NEIGHBORS_FIELD.clone(),
214            VECTOR_ID_FIELD.clone(),
215        ]))?;
216
217        let reader = V1FileReader::try_new_from_reader(
218            reader.path(),
219            reader.clone(),
220            None,
221            schema,
222            0,
223            0,
224            2,
225            None,
226        )
227        .await?;
228
229        let storage = Arc::new(self.partition_storage.load_partition(0).await?);
230        let batch = reader.read_range(0..reader.len(), reader.schema()).await?;
231        let hnsw = HNSW::load(batch)?;
232
233        Ok(Box::new(Self {
234            hnsw: Some(hnsw),
235            storage: Some(storage),
236            partition_storage: self.partition_storage.clone(),
237            partition_metadata: self.partition_metadata.clone(),
238            options: self.options.clone(),
239        }))
240    }
241
242    async fn load_partition(
243        &self,
244        reader: Arc<dyn Reader>,
245        offset: usize,
246        length: usize,
247        partition_id: usize,
248    ) -> Result<Box<dyn VectorIndex>> {
249        let reader = V1FileReader::try_new_self_described_from_reader(reader, None).await?;
250
251        let metadata = self.get_partition_metadata(partition_id)?;
252        let storage = Arc::new(self.partition_storage.load_partition(partition_id).await?);
253        let batch = reader
254            .read_range(offset..offset + length, reader.schema())
255            .await?;
256        let mut schema = batch.schema_ref().as_ref().clone();
257        schema.metadata.insert(
258            HNSW::metadata_key().to_owned(),
259            serde_json::to_string(&metadata)?,
260        );
261        let batch = batch.with_schema(schema.into())?;
262        let hnsw = HNSW::load(batch)?;
263
264        Ok(Box::new(Self {
265            hnsw: Some(hnsw),
266            storage: Some(storage),
267            partition_storage: self.partition_storage.clone(),
268            partition_metadata: self.partition_metadata.clone(),
269            options: self.options.clone(),
270        }))
271    }
272
273    async fn to_batch_stream(&self, with_vector: bool) -> Result<SendableRecordBatchStream> {
274        let store = self
275            .storage
276            .as_ref()
277            .ok_or(Error::index("vector storage not loaded".to_string()))?;
278
279        let schema = if with_vector {
280            store.schema().clone()
281        } else {
282            let schema = store.schema();
283            let row_id_idx = schema.index_of(ROW_ID)?;
284            Arc::new(schema.project(&[row_id_idx])?)
285        };
286
287        let batches = store
288            .to_batches()?
289            .map(|b| {
290                let batch = b.project_by_schema(&schema)?;
291                Ok(batch)
292            })
293            .collect::<Vec<_>>();
294        let stream = futures::stream::iter(batches);
295        let stream = RecordBatchStreamAdapter::new(schema, stream);
296        Ok(Box::pin(stream))
297    }
298
299    fn num_rows(&self) -> u64 {
300        self.hnsw
301            .as_ref()
302            .map_or(0, |hnsw| hnsw.num_nodes(0) as u64)
303    }
304
305    fn row_ids(&self) -> Box<dyn Iterator<Item = &'_ u64> + '_> {
306        Box::new(self.storage.as_ref().unwrap().row_ids())
307    }
308
309    async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> {
310        Err(Error::index(
311            "Remapping HNSW in this way not supported".to_string(),
312        ))
313    }
314
315    fn ivf_model(&self) -> &IvfModel {
316        unimplemented!("only for IVF")
317    }
318
319    fn quantizer(&self) -> Quantizer {
320        self.partition_storage.quantizer().clone()
321    }
322
323    fn partition_size(&self, _: usize) -> usize {
324        unimplemented!("only for IVF")
325    }
326
327    fn sub_index_type(&self) -> (SubIndexType, QuantizationType) {
328        (
329            SubIndexType::Hnsw,
330            self.partition_storage.quantizer().quantization_type(),
331        )
332    }
333
334    fn metric_type(&self) -> DistanceType {
335        self.partition_storage.distance_type()
336    }
337}
lance_index/vector/hnsw/index.rs

lance_index/vector/hnsw/
index.rs
