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`.

docs/src/js/classes/Table.md

@@ -385,6 +385,13 @@ await table.createIndex("vector", {
 });
 ```
 
+```ts
+// Or size the partitions by row count instead of choosing how many.
+await table.createIndex("vector", {
+  config: lancedb.Index.ivfPq({ targetPartitionSize: 4096 }),
+});
+```
+
 ```ts
 // Or create a Scalar index
 await table.createIndex("my_float_col");

docs/src/js/interfaces/HnswPqOptions.md

@@ -100,6 +100,8 @@ optional numPartitions: number;
 
 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
@@ -147,3 +149,23 @@ 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.
+
+***
+
+### targetPartitionSize?
+
+```ts
+optional targetPartitionSize: 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.

docs/src/js/interfaces/HnswSqOptions.md

@@ -100,6 +100,8 @@ optional numPartitions: number;
 
 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
@@ -126,3 +128,23 @@ 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.
+
+***
+
+### targetPartitionSize?
+
+```ts
+optional targetPartitionSize: 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.

docs/src/js/interfaces/IvfFlatOptions.md

@@ -82,8 +82,9 @@ optional numPartitions: number;
 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
@@ -110,3 +111,23 @@ Increasing this value might improve the quality of the index but in most cases t
 default should be sufficient.
 
 The default value is 256.
+
+***
+
+### targetPartitionSize?
+
+```ts
+optional targetPartitionSize: 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.

docs/src/js/interfaces/IvfPqOptions.md

@@ -92,8 +92,9 @@ optional numPartitions: number;
 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
@@ -141,3 +142,23 @@ Increasing this value might improve the quality of the index but in most cases t
 default should be sufficient.
 
 The default value is 256.
+
+***
+
+### targetPartitionSize?
+
+```ts
+optional targetPartitionSize: 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.

docs/src/js/interfaces/IvfRqOptions.md

@@ -72,8 +72,9 @@ optional numPartitions: number;
 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
@@ -100,3 +101,23 @@ Increasing this value might improve the quality of the index but in most cases t
 default should be sufficient.
 
 The default value is 256.
+
+***
+
+### targetPartitionSize?
+
+```ts
+optional targetPartitionSize: 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.

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);
         }