diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index b1929019c..4398ee3ad 100644 --- a/docs/src/js/classes/Table.md +++ b/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"); diff --git a/docs/src/js/interfaces/HnswPqOptions.md b/docs/src/js/interfaces/HnswPqOptions.md index 6d42ed28f..f044e0e11 100644 --- a/docs/src/js/interfaces/HnswPqOptions.md +++ b/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. diff --git a/docs/src/js/interfaces/HnswSqOptions.md b/docs/src/js/interfaces/HnswSqOptions.md index 585737a5a..332854490 100644 --- a/docs/src/js/interfaces/HnswSqOptions.md +++ b/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. diff --git a/docs/src/js/interfaces/IvfFlatOptions.md b/docs/src/js/interfaces/IvfFlatOptions.md index 36da98907..5b8024bc3 100644 --- a/docs/src/js/interfaces/IvfFlatOptions.md +++ b/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. diff --git a/docs/src/js/interfaces/IvfPqOptions.md b/docs/src/js/interfaces/IvfPqOptions.md index 7109e448c..f1dd47621 100644 --- a/docs/src/js/interfaces/IvfPqOptions.md +++ b/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. diff --git a/docs/src/js/interfaces/IvfRqOptions.md b/docs/src/js/interfaces/IvfRqOptions.md index 298565a62..eb2cfe0ac 100644 --- a/docs/src/js/interfaces/IvfRqOptions.md +++ b/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. diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 6db5b2f4d..c98292fd4 100644 --- a/nodejs/__test__/table.test.ts +++ b/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)[ + "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)[ + "target_partition_size" + ], + ).toBeUndefined(); + }); + it("should allow me to replace (or not) an existing index", async () => { await tbl.createIndex("id"); // Default is replace=true diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index 0e383e984..4b1b422bf 100644 --- a/nodejs/lancedb/indices.ts +++ b/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, ), ); } diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 8e7cd02a1..0806f7c6d 100644 --- a/nodejs/lancedb/table.ts +++ b/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"); */ diff --git a/nodejs/src/index.rs b/nodejs/src/index.rs index e5884aa4d..39aa6f8b8 100644 --- a/nodejs/src/index.rs +++ b/nodejs/src/index.rs @@ -104,6 +104,7 @@ impl Index { num_bits: Option, max_iterations: Option, sample_rate: Option, + target_partition_size: Option, ) -> napi::Result { 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, max_iterations: Option, sample_rate: Option, + target_partition_size: Option, ) -> napi::Result { 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, max_iterations: Option, sample_rate: Option, + target_partition_size: Option, ) -> napi::Result { 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, num_partitions: Option, @@ -291,6 +304,7 @@ impl Index { sample_rate: Option, m: Option, ef_construction: Option, + target_partition_size: Option, ) -> napi::Result { 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, m: Option, ef_construction: Option, + target_partition_size: Option, ) -> napi::Result { 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); }