feat(napi): Provide a new createIndex API in the napi SDK. (#857)

nodejs/src/index.rs

@@ -0,0 +1,101 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use lance_linalg::distance::MetricType as LanceMetricType;
+use napi_derive::napi;
+
+#[napi]
+pub enum IndexType {
+    Scalar,
+    IvfPq,
+}
+
+#[napi]
+pub enum MetricType {
+    L2,
+    Cosine,
+    Dot,
+}
+
+impl From<MetricType> for LanceMetricType {
+    fn from(metric: MetricType) -> Self {
+        match metric {
+            MetricType::L2 => Self::L2,
+            MetricType::Cosine => Self::Cosine,
+            MetricType::Dot => Self::Dot,
+        }
+    }
+}
+
+#[napi]
+pub struct IndexBuilder {
+    inner: vectordb::index::IndexBuilder,
+}
+
+#[napi]
+impl IndexBuilder {
+    pub fn new(tbl: &dyn vectordb::Table) -> Self {
+        let inner = tbl.create_index(&[]);
+        Self { inner }
+    }
+
+    #[napi]
+    pub unsafe fn replace(&mut self, v: bool) {
+        self.inner.replace(v);
+    }
+
+    #[napi]
+    pub unsafe fn column(&mut self, c: String) {
+        self.inner.columns(&[c.as_str()]);
+    }
+
+    #[napi]
+    pub unsafe fn name(&mut self, name: String) {
+        self.inner.name(name.as_str());
+    }
+
+    #[napi]
+    pub unsafe fn ivf_pq(
+        &mut self,
+        metric_type: Option<MetricType>,
+        num_partitions: Option<u32>,
+        num_sub_vectors: Option<u32>,
+        num_bits: Option<u32>,
+        max_iterations: Option<u32>,
+        sample_rate: Option<u32>,
+    ) {
+        self.inner.ivf_pq();
+        metric_type.map(|m| self.inner.metric_type(m.into()));
+        num_partitions.map(|p| self.inner.num_partitions(p));
+        num_sub_vectors.map(|s| self.inner.num_sub_vectors(s));
+        num_bits.map(|b| self.inner.num_bits(b));
+        max_iterations.map(|i| self.inner.max_iterations(i));
+        sample_rate.map(|s| self.inner.sample_rate(s));
+    }
+
+    #[napi]
+    pub unsafe fn scalar(&mut self) {
+        self.inner.scalar();
+    }
+
+    #[napi]
+    pub async fn build(&self) -> napi::Result<()> {
+        println!("nodejs::index.rs : build");
+        self.inner
+            .build()
+            .await
+            .map_err(|e| napi::Error::from_reason(format!("Failed to build index: {}", e)))?;
+        Ok(())
+    }
+}

nodejs/src/lib.rs

@@ -16,6 +16,7 @@ use connection::Connection;
 use napi_derive::*;
 
 mod connection;
+mod index;
 mod query;
 mod table;
 

nodejs/src/table.rs

@@ -12,12 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::query::Query;
 use arrow_ipc::writer::FileWriter;
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
 use vectordb::{ipc::ipc_file_to_batches, table::TableRef};
 
+use crate::index::IndexBuilder;
+use crate::query::Query;
+
 #[napi]
 pub struct Table {
     pub(crate) table: TableRef,
@@ -43,7 +45,7 @@ impl Table {
     }
 
     #[napi]
-    pub async unsafe fn add(&mut self, buf: Buffer) -> napi::Result<()> {
+    pub async fn add(&self, buf: Buffer) -> napi::Result<()> {
         let batches = ipc_file_to_batches(buf.to_vec())
             .map_err(|e| napi::Error::from_reason(format!("Failed to read IPC file: {}", e)))?;
         self.table.add(Box::new(batches), None).await.map_err(|e| {
@@ -65,7 +67,7 @@ impl Table {
     }
 
     #[napi]
-    pub async unsafe fn delete(&mut self, predicate: String) -> napi::Result<()> {
+    pub async fn delete(&self, predicate: String) -> napi::Result<()> {
         self.table.delete(&predicate).await.map_err(|e| {
             napi::Error::from_reason(format!(
                 "Failed to delete rows in table {}: predicate={}",
@@ -74,6 +76,11 @@ impl Table {
         })
     }
 
+    #[napi]
+    pub fn create_index(&self) -> IndexBuilder {
+        IndexBuilder::new(self.table.as_ref())
+    }
+
     #[napi]
     pub fn query(&self) -> Query {
         Query::new(self)