feat: expose hnsw indices (#1595)

PR closes #1522

---------

Co-authored-by: Will Jones <willjones127@gmail.com>

python/src/index.rs

@@ -16,7 +16,11 @@ use std::sync::Mutex;
 
 use lancedb::index::scalar::FtsIndexBuilder;
 use lancedb::{
-    index::{scalar::BTreeIndexBuilder, vector::IvfPqIndexBuilder, Index as LanceDbIndex},
+    index::{
+        scalar::BTreeIndexBuilder,
+        vector::{IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder},
+        Index as LanceDbIndex,
+    },
     DistanceType,
 };
 use pyo3::{
@@ -24,6 +28,8 @@ use pyo3::{
     pyclass, pymethods, PyResult,
 };
 
+use crate::util::parse_distance_type;
+
 #[pyclass]
 pub struct Index {
     inner: Mutex<Option<LanceDbIndex>>,
@@ -110,6 +116,78 @@ impl Index {
             inner: Mutex::new(Some(LanceDbIndex::FTS(opts))),
         }
     }
+
+    #[staticmethod]
+    pub fn hnsw_pq(
+        distance_type: Option<String>,
+        num_partitions: Option<u32>,
+        num_sub_vectors: Option<u32>,
+        max_iterations: Option<u32>,
+        sample_rate: Option<u32>,
+        m: Option<u32>,
+        ef_construction: Option<u32>,
+    ) -> PyResult<Self> {
+        let mut hnsw_pq_builder = IvfHnswPqIndexBuilder::default();
+        if let Some(distance_type) = distance_type {
+            let distance_type = parse_distance_type(distance_type)?;
+            hnsw_pq_builder = hnsw_pq_builder.distance_type(distance_type);
+        }
+        if let Some(num_partitions) = num_partitions {
+            hnsw_pq_builder = hnsw_pq_builder.num_partitions(num_partitions);
+        }
+        if let Some(num_sub_vectors) = num_sub_vectors {
+            hnsw_pq_builder = hnsw_pq_builder.num_sub_vectors(num_sub_vectors);
+        }
+        if let Some(max_iterations) = max_iterations {
+            hnsw_pq_builder = hnsw_pq_builder.max_iterations(max_iterations);
+        }
+        if let Some(sample_rate) = sample_rate {
+            hnsw_pq_builder = hnsw_pq_builder.sample_rate(sample_rate);
+        }
+        if let Some(m) = m {
+            hnsw_pq_builder = hnsw_pq_builder.num_edges(m);
+        }
+        if let Some(ef_construction) = ef_construction {
+            hnsw_pq_builder = hnsw_pq_builder.ef_construction(ef_construction);
+        }
+        Ok(Self {
+            inner: Mutex::new(Some(LanceDbIndex::IvfHnswPq(hnsw_pq_builder))),
+        })
+    }
+
+    #[staticmethod]
+    pub fn hnsw_sq(
+        distance_type: Option<String>,
+        num_partitions: Option<u32>,
+        max_iterations: Option<u32>,
+        sample_rate: Option<u32>,
+        m: Option<u32>,
+        ef_construction: Option<u32>,
+    ) -> PyResult<Self> {
+        let mut hnsw_sq_builder = IvfHnswSqIndexBuilder::default();
+        if let Some(distance_type) = distance_type {
+            let distance_type = parse_distance_type(distance_type)?;
+            hnsw_sq_builder = hnsw_sq_builder.distance_type(distance_type);
+        }
+        if let Some(num_partitions) = num_partitions {
+            hnsw_sq_builder = hnsw_sq_builder.num_partitions(num_partitions);
+        }
+        if let Some(max_iterations) = max_iterations {
+            hnsw_sq_builder = hnsw_sq_builder.max_iterations(max_iterations);
+        }
+        if let Some(sample_rate) = sample_rate {
+            hnsw_sq_builder = hnsw_sq_builder.sample_rate(sample_rate);
+        }
+        if let Some(m) = m {
+            hnsw_sq_builder = hnsw_sq_builder.num_edges(m);
+        }
+        if let Some(ef_construction) = ef_construction {
+            hnsw_sq_builder = hnsw_sq_builder.ef_construction(ef_construction);
+        }
+        Ok(Self {
+            inner: Mutex::new(Some(LanceDbIndex::IvfHnswSq(hnsw_sq_builder))),
+        })
+    }
 }
 
 #[pyclass(get_all)]