feat(node): expose targetPartitionSize on IVF index options (#4594)

The Node SDK's vector index options only take `numPartitions`. The Rust index builders (`target_partition_size` in `rust/lancedb/src/index/vector.rs`) and the Python SDK already support a target partition size, so this adds `targetPartitionSize` to `IvfPqOptions`, `IvfRqOptions`, `IvfFlatOptions`, `HnswPqOptions` and `HnswSqOptions` and passes it to their builders. When both are set, `numPartitions` is used, as it is for Rust and Python callers.

The reason to want it: an index built from a partition count records no target partition size, so `table.optimize()` splits and joins it against the index type's default and can merge small partitions down to one or two. Building with a target records it in the index details, and optimize keeps to it. https://github.com/lance-format/lance/pull/10154 fixes this on the Lance side for indexes built from a partition count; this lets Node users set the size directly.

The new tests build IvfPq, IvfFlat, IvfRq, IvfHnswPq and IvfHnswSq indexes with `targetPartitionSize` and check the size recorded in the index details, and check that `numPartitions` wins when both are set. Each option's doc comment states that index type's default target. Docs regenerated with `pnpm run docs`.

nodejs/__test__/table.test.ts

@@ -1707,6 +1707,38 @@ describe("When creating an index", () => {
     });
   });
 
+  it.each([
+    ["IvfPq", Index.ivfPq({ targetPartitionSize: 100 })],
+    ["IvfFlat", Index.ivfFlat({ targetPartitionSize: 100 })],
+    ["IvfRq", Index.ivfRq({ targetPartitionSize: 100 })],
+    ["IvfHnswPq", Index.hnswPq({ targetPartitionSize: 100 })],
+    ["IvfHnswSq", Index.hnswSq({ targetPartitionSize: 100 })],
+  ])("should create %s with targetPartitionSize", async (indexType, config) => {
+    await tbl.createIndex("vec", { config });
+
+    const indices = await tbl.listIndices();
+    expect(indices).toHaveLength(1);
+    expect(indices[0].indexType).toBe(indexType);
+    expect(
+      (indices[0].indexDetails as Record<string, unknown>)[
+        "target_partition_size"
+      ],
+    ).toBe(100);
+  });
+
+  it("should prefer numPartitions over targetPartitionSize", async () => {
+    await tbl.createIndex("vec", {
+      config: Index.ivfFlat({ numPartitions: 2, targetPartitionSize: 100 }),
+    });
+
+    const indices = await tbl.listIndices();
+    expect(
+      (indices[0].indexDetails as Record<string, unknown>)[
+        "target_partition_size"
+      ],
+    ).toBeUndefined();
+  });
+
   it("should allow me to replace (or not) an existing index", async () => {
     await tbl.createIndex("id");
     // Default is replace=true

nodejs/lancedb/indices.ts

@@ -11,8 +11,9 @@ export interface IvfPqOptions {
    * 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.
+   * When set, this takes precedence over `targetPartitionSize`. When omitted,
+   * the partition count is derived from `targetPartitionSize`, which defaults
+   * to 8192 rows for this index type.
    *
    * 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
@@ -20,6 +21,21 @@ export interface IvfPqOptions {
    */
   numPartitions?: number;
 
+  /**
+   * The target number of rows in each IVF partition.
+   *
+   * This is an alternative to `numPartitions`. When it is set and
+   * `numPartitions` is not, the number of partitions is derived by dividing
+   * the number of rows by this value. When both are set, `numPartitions` is
+   * used and this value is ignored. When neither is set, the target is
+   * 8192 rows.
+   *
+   * Larger partitions mean fewer of them, which makes the first part of the
+   * search (picking the right partitions) faster and the second part
+   * (searching within a partition) slower.
+   */
+  targetPartitionSize?: number;
+
   /**
    * Number of sub-vectors of PQ.
    *
@@ -117,8 +133,9 @@ 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.
+   * When set, this takes precedence over `targetPartitionSize`. When omitted,
+   * the partition count is derived from `targetPartitionSize`, which defaults
+   * to 4096 rows for this index type.
    *
    * 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
@@ -126,6 +143,21 @@ export interface IvfRqOptions {
    */
   numPartitions?: number;
 
+  /**
+   * The target number of rows in each IVF partition.
+   *
+   * This is an alternative to `numPartitions`. When it is set and
+   * `numPartitions` is not, the number of partitions is derived by dividing
+   * the number of rows by this value. When both are set, `numPartitions` is
+   * used and this value is ignored. When neither is set, the target is
+   * 4096 rows.
+   *
+   * Larger partitions mean fewer of them, which makes the first part of the
+   * search (picking the right partitions) faster and the second part
+   * (searching within a partition) slower.
+   */
+  targetPartitionSize?: number;
+
   /**
    * Number of bits per dimension for residual quantization.
    *
@@ -215,6 +247,8 @@ export interface HnswPqOptions {
   /**
    * The number of IVF partitions to create.
    *
+   * When set, this takes precedence over `targetPartitionSize`.
+   *
    * For HNSW, we recommend a small number of partitions. Setting this to 1 works
    * well for most tables. For very large tables, training just one HNSW graph
    * will require too much memory. Each partition becomes its own HNSW graph, so
@@ -223,6 +257,21 @@ export interface HnswPqOptions {
    */
   numPartitions?: number;
 
+  /**
+   * The target number of rows in each IVF partition.
+   *
+   * This is an alternative to `numPartitions`. When it is set and
+   * `numPartitions` is not, the number of partitions is derived by dividing
+   * the number of rows by this value. When both are set, `numPartitions` is
+   * used and this value is ignored. When neither is set, the target is
+   * 1,048,576 rows.
+   *
+   * Larger partitions mean fewer of them, which makes the first part of the
+   * search (picking the right partitions) faster and the second part
+   * (searching within a partition) slower.
+   */
+  targetPartitionSize?: number;
+
   /**
    * Number of sub-vectors of PQ.
    *
@@ -331,6 +380,8 @@ export interface HnswSqOptions {
   /**
    * The number of IVF partitions to create.
    *
+   * When set, this takes precedence over `targetPartitionSize`.
+   *
    * For HNSW, we recommend a small number of partitions. Setting this to 1 works
    * well for most tables. For very large tables, training just one HNSW graph
    * will require too much memory. Each partition becomes its own HNSW graph, so
@@ -339,6 +390,21 @@ export interface HnswSqOptions {
    */
   numPartitions?: number;
 
+  /**
+   * The target number of rows in each IVF partition.
+   *
+   * This is an alternative to `numPartitions`. When it is set and
+   * `numPartitions` is not, the number of partitions is derived by dividing
+   * the number of rows by this value. When both are set, `numPartitions` is
+   * used and this value is ignored. When neither is set, the target is
+   * 1,048,576 rows.
+   *
+   * Larger partitions mean fewer of them, which makes the first part of the
+   * search (picking the right partitions) faster and the second part
+   * (searching within a partition) slower.
+   */
+  targetPartitionSize?: number;
+
   /**
    * Max iterations to train kmeans.
    *
@@ -407,8 +473,9 @@ export interface IvfFlatOptions {
    * 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.
+   * When set, this takes precedence over `targetPartitionSize`. When omitted,
+   * the partition count is derived from `targetPartitionSize`, which defaults
+   * to 4096 rows for this index type.
    *
    * 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
@@ -416,6 +483,21 @@ export interface IvfFlatOptions {
    */
   numPartitions?: number;
 
+  /**
+   * The target number of rows in each IVF partition.
+   *
+   * This is an alternative to `numPartitions`. When it is set and
+   * `numPartitions` is not, the number of partitions is derived by dividing
+   * the number of rows by this value. When both are set, `numPartitions` is
+   * used and this value is ignored. When neither is set, the target is
+   * 4096 rows.
+   *
+   * Larger partitions mean fewer of them, which makes the first part of the
+   * search (picking the right partitions) faster and the second part
+   * (searching within a partition) slower.
+   */
+  targetPartitionSize?: number;
+
   /**
    * Distance type to use to build the index.
    *
@@ -633,6 +715,7 @@ export class Index {
         options?.numBits,
         options?.maxIterations,
         options?.sampleRate,
+        options?.targetPartitionSize,
       ),
     );
   }
@@ -661,6 +744,7 @@ export class Index {
         options?.numBits,
         options?.maxIterations,
         options?.sampleRate,
+        options?.targetPartitionSize,
       ),
     );
   }
@@ -688,6 +772,7 @@ export class Index {
         options?.numPartitions,
         options?.maxIterations,
         options?.sampleRate,
+        options?.targetPartitionSize,
       ),
     );
   }
@@ -798,6 +883,7 @@ export class Index {
         options?.sampleRate,
         options?.m,
         options?.efConstruction,
+        options?.targetPartitionSize,
       ),
     );
   }
@@ -820,6 +906,7 @@ export class Index {
         options?.sampleRate,
         options?.m,
         options?.efConstruction,
+        options?.targetPartitionSize,
       ),
     );
   }

nodejs/lancedb/table.ts

@@ -381,6 +381,11 @@ export abstract class Table {
    *   }),
    * });
    * @example
+   * // Or size the partitions by row count instead of choosing how many.
+   * await table.createIndex("vector", {
+   *   config: lancedb.Index.ivfPq({ targetPartitionSize: 4096 }),
+   * });
+   * @example
    * // Or create a Scalar index
    * await table.createIndex("my_float_col");
    */

nodejs/src/index.rs

@@ -104,6 +104,7 @@ impl Index {
         num_bits: Option<u32>,
         max_iterations: Option<u32>,
         sample_rate: Option<u32>,
+        target_partition_size: Option<u32>,
     ) -> napi::Result<Self> {
         let mut ivf_pq_builder = IvfPqIndexBuilder::default();
         if let Some(distance_type) = distance_type {
@@ -113,6 +114,9 @@ impl Index {
         if let Some(num_partitions) = num_partitions {
             ivf_pq_builder = ivf_pq_builder.num_partitions(num_partitions);
         }
+        if let Some(target_partition_size) = target_partition_size {
+            ivf_pq_builder = ivf_pq_builder.target_partition_size(target_partition_size);
+        }
         if let Some(num_sub_vectors) = num_sub_vectors {
             ivf_pq_builder = ivf_pq_builder.num_sub_vectors(num_sub_vectors);
         }
@@ -137,6 +141,7 @@ impl Index {
         num_bits: Option<u32>,
         max_iterations: Option<u32>,
         sample_rate: Option<u32>,
+        target_partition_size: Option<u32>,
     ) -> napi::Result<Self> {
         let mut ivf_rq_builder = IvfRqIndexBuilder::default();
         if let Some(distance_type) = distance_type {
@@ -146,6 +151,9 @@ impl Index {
         if let Some(num_partitions) = num_partitions {
             ivf_rq_builder = ivf_rq_builder.num_partitions(num_partitions);
         }
+        if let Some(target_partition_size) = target_partition_size {
+            ivf_rq_builder = ivf_rq_builder.target_partition_size(target_partition_size);
+        }
         if let Some(num_bits) = num_bits {
             ivf_rq_builder = ivf_rq_builder.num_bits(num_bits);
         }
@@ -166,6 +174,7 @@ impl Index {
         num_partitions: Option<u32>,
         max_iterations: Option<u32>,
         sample_rate: Option<u32>,
+        target_partition_size: Option<u32>,
     ) -> napi::Result<Self> {
         let mut ivf_flat_builder = IvfFlatIndexBuilder::default();
         if let Some(distance_type) = distance_type {
@@ -175,6 +184,9 @@ impl Index {
         if let Some(num_partitions) = num_partitions {
             ivf_flat_builder = ivf_flat_builder.num_partitions(num_partitions);
         }
+        if let Some(target_partition_size) = target_partition_size {
+            ivf_flat_builder = ivf_flat_builder.target_partition_size(target_partition_size);
+        }
         if let Some(max_iterations) = max_iterations {
             ivf_flat_builder = ivf_flat_builder.max_iterations(max_iterations);
         }
@@ -283,6 +295,7 @@ impl Index {
     }
 
     #[napi(factory)]
+    #[allow(clippy::too_many_arguments)]
     pub fn hnsw_pq(
         distance_type: Option<String>,
         num_partitions: Option<u32>,
@@ -291,6 +304,7 @@ impl Index {
         sample_rate: Option<u32>,
         m: Option<u32>,
         ef_construction: Option<u32>,
+        target_partition_size: Option<u32>,
     ) -> napi::Result<Self> {
         let mut hnsw_pq_builder = IvfHnswPqIndexBuilder::default();
         if let Some(distance_type) = distance_type {
@@ -300,6 +314,9 @@ impl Index {
         if let Some(num_partitions) = num_partitions {
             hnsw_pq_builder = hnsw_pq_builder.num_partitions(num_partitions);
         }
+        if let Some(target_partition_size) = target_partition_size {
+            hnsw_pq_builder = hnsw_pq_builder.target_partition_size(target_partition_size);
+        }
         if let Some(num_sub_vectors) = num_sub_vectors {
             hnsw_pq_builder = hnsw_pq_builder.num_sub_vectors(num_sub_vectors);
         }
@@ -328,6 +345,7 @@ impl Index {
         sample_rate: Option<u32>,
         m: Option<u32>,
         ef_construction: Option<u32>,
+        target_partition_size: Option<u32>,
     ) -> napi::Result<Self> {
         let mut hnsw_sq_builder = IvfHnswSqIndexBuilder::default();
         if let Some(distance_type) = distance_type {
@@ -337,6 +355,9 @@ impl Index {
         if let Some(num_partitions) = num_partitions {
             hnsw_sq_builder = hnsw_sq_builder.num_partitions(num_partitions);
         }
+        if let Some(target_partition_size) = target_partition_size {
+            hnsw_sq_builder = hnsw_sq_builder.target_partition_size(target_partition_size);
+        }
         if let Some(max_iterations) = max_iterations {
             hnsw_sq_builder = hnsw_sq_builder.max_iterations(max_iterations);
         }