feat(index): add IVF_RQ index type (#2687)

this expose IVF_RQ (RabitQ quantization) index type to lancedb

---------

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

nodejs/__test__/table.test.ts

@@ -861,6 +861,15 @@ describe("When creating an index", () => {
     });
   });
 
+  it("should be able to create IVF_RQ", async () => {
+    await tbl.createIndex("vec", {
+      config: Index.ivfRq({
+        numPartitions: 10,
+        numBits: 1,
+      }),
+    });
+  });
+
   it("should allow me to replace (or not) an existing index", async () => {
     await tbl.createIndex("id");
     // Default is replace=true

nodejs/lancedb/index.ts

@@ -85,6 +85,7 @@ export {
   Index,
   IndexOptions,
   IvfPqOptions,
+  IvfRqOptions,
   IvfFlatOptions,
   HnswPqOptions,
   HnswSqOptions,

nodejs/lancedb/indices.ts

@@ -112,6 +112,77 @@ export interface IvfPqOptions {
   sampleRate?: number;
 }
 
+export interface IvfRqOptions {
+  /**
+   * The number of IVF partitions to create.
+   *
+   * This value should generally scale with the number of rows in the dataset.
+   * By default the number of partitions is the square root of the number of
+   * rows.
+   *
+   * If this value is too large then the first part of the search (picking the
+   * right partition) will be slow. If this value is too small then the second
+   * part of the search (searching within a partition) will be slow.
+   */
+  numPartitions?: number;
+
+  /**
+   * Number of bits per dimension for residual quantization.
+   *
+   * This value controls how much each residual component is compressed. The more
+   * bits, the more accurate the index will be but the slower search. Typical values
+   * are small integers; the default is 1 bit per dimension.
+   */
+  numBits?: number;
+
+  /**
+   * Distance type to use to build the index.
+   *
+   * Default value is "l2".
+   *
+   * This is used when training the index to calculate the IVF partitions
+   * (vectors are grouped in partitions with similar vectors according to this
+   * distance type) and during quantization.
+   *
+   * The distance type used to train an index MUST match the distance type used
+   * to search the index. Failure to do so will yield inaccurate results.
+   *
+   * The following distance types are available:
+   *
+   * "l2" - Euclidean distance.
+   * "cosine" - Cosine distance.
+   * "dot" - Dot product.
+   */
+  distanceType?: "l2" | "cosine" | "dot";
+
+  /**
+   * Max iterations to train IVF kmeans.
+   *
+   * When training an IVF index we use kmeans to calculate the partitions. This parameter
+   * controls how many iterations of kmeans to run.
+   *
+   * The default value is 50.
+   */
+  maxIterations?: number;
+
+  /**
+   * The number of vectors, per partition, to sample when training IVF kmeans.
+   *
+   * When an IVF index is trained, we need to calculate partitions. These are groups
+   * of vectors that are similar to each other. To do this we use an algorithm called kmeans.
+   *
+   * Running kmeans on a large dataset can be slow. To speed this up we run kmeans on a
+   * random sample of the data. This parameter controls the size of the sample. The total
+   * number of vectors used to train the index is `sample_rate * num_partitions`.
+   *
+   * Increasing this value might improve the quality of the index but in most cases the
+   * default should be sufficient.
+   *
+   * The default value is 256.
+   */
+  sampleRate?: number;
+}
+
 /**
  * Options to create an `HNSW_PQ` index
  */
@@ -523,6 +594,35 @@ export class Index {
         options?.distanceType,
         options?.numPartitions,
         options?.numSubVectors,
+        options?.numBits,
+        options?.maxIterations,
+        options?.sampleRate,
+      ),
+    );
+  }
+
+  /**
+   * Create an IvfRq index
+   *
+   * IVF-RQ (RabitQ Quantization) compresses vectors using RabitQ quantization
+   * and organizes them into IVF partitions.
+   *
+   * The compression scheme is called RabitQ quantization. Each dimension is quantized into a small number of bits.
+   * The parameters `num_bits` and `num_partitions` control this process, providing a tradeoff
+   * between index size (and thus search speed) and index accuracy.
+   *
+   * The partitioning process is called IVF and the `num_partitions` parameter controls how
+   * many groups to create.
+   *
+   * Note that training an IVF RQ index on a large dataset is a slow operation and
+   * currently is also a memory intensive operation.
+   */
+  static ivfRq(options?: Partial<IvfRqOptions>) {
+    return new Index(
+      LanceDbIndex.ivfRq(
+        options?.distanceType,
+        options?.numPartitions,
+        options?.numBits,
         options?.maxIterations,
         options?.sampleRate,
       ),

nodejs/src/index.rs

@@ -6,6 +6,7 @@ use std::sync::Mutex;
 use lancedb::index::scalar::{BTreeIndexBuilder, FtsIndexBuilder};
 use lancedb::index::vector::{
     IvfFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder,
+    IvfRqIndexBuilder,
 };
 use lancedb::index::Index as LanceDbIndex;
 use napi_derive::napi;
@@ -65,6 +66,36 @@ impl Index {
         })
     }
 
+    #[napi(factory)]
+    pub fn ivf_rq(
+        distance_type: Option<String>,
+        num_partitions: Option<u32>,
+        num_bits: Option<u32>,
+        max_iterations: Option<u32>,
+        sample_rate: Option<u32>,
+    ) -> napi::Result<Self> {
+        let mut ivf_rq_builder = IvfRqIndexBuilder::default();
+        if let Some(distance_type) = distance_type {
+            let distance_type = parse_distance_type(distance_type)?;
+            ivf_rq_builder = ivf_rq_builder.distance_type(distance_type);
+        }
+        if let Some(num_partitions) = num_partitions {
+            ivf_rq_builder = ivf_rq_builder.num_partitions(num_partitions);
+        }
+        if let Some(num_bits) = num_bits {
+            ivf_rq_builder = ivf_rq_builder.num_bits(num_bits);
+        }
+        if let Some(max_iterations) = max_iterations {
+            ivf_rq_builder = ivf_rq_builder.max_iterations(max_iterations);
+        }
+        if let Some(sample_rate) = sample_rate {
+            ivf_rq_builder = ivf_rq_builder.sample_rate(sample_rate);
+        }
+        Ok(Self {
+            inner: Mutex::new(Some(LanceDbIndex::IvfRq(ivf_rq_builder))),
+        })
+    }
+
     #[napi(factory)]
     pub fn ivf_flat(
         distance_type: Option<String>,