From bb48446f20fb7abb4426d0281edf6cb5c524b907 Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:30:09 +0800 Subject: [PATCH] feat(index): expose layered RaBitQ defaults across SDKs --- Cargo.lock | 48 +++++++++++----------- Cargo.toml | 28 ++++++------- docs/src/js/interfaces/IvfRqOptions.md | 13 +++++- nodejs/__test__/table.test.ts | 40 ++++++++++++++---- nodejs/lancedb/indices.ts | 9 ++++- nodejs/src/index.rs | 4 ++ python/python/lancedb/index.py | 9 ++++- python/python/lancedb/table.py | 1 + python/python/tests/test_index.py | 20 ++++++++- python/python/tests/test_remote_db.py | 12 +++++- python/src/index.rs | 4 ++ rust/lancedb/src/index/vector.rs | 16 ++++++++ rust/lancedb/src/remote/table.rs | 14 ++++++- rust/lancedb/src/table/create_index.rs | 56 ++++++++++++++++++++++++-- 14 files changed, 216 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5c713d99a..091ed859e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3585,7 +3585,7 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "rand 0.9.5", @@ -5060,7 +5060,7 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arc-swap", "arrow", @@ -5134,7 +5134,7 @@ dependencies = [ [[package]] name = "lance-arrow" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-buffer", @@ -5156,7 +5156,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-buffer", @@ -5170,7 +5170,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-schema", @@ -5180,7 +5180,7 @@ dependencies = [ [[package]] name = "lance-bitpacking" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrayref", "crunchy", @@ -5191,7 +5191,7 @@ dependencies = [ [[package]] name = "lance-core" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-buffer", @@ -5229,7 +5229,7 @@ dependencies = [ [[package]] name = "lance-datafusion" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow", "arrow-array", @@ -5261,7 +5261,7 @@ dependencies = [ [[package]] name = "lance-datagen" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow", "arrow-array", @@ -5279,7 +5279,7 @@ dependencies = [ [[package]] name = "lance-derive" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "proc-macro2", "quote", @@ -5289,7 +5289,7 @@ dependencies = [ [[package]] name = "lance-encoding" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-arith", "arrow-array", @@ -5323,7 +5323,7 @@ dependencies = [ [[package]] name = "lance-file" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-arith", "arrow-array", @@ -5356,7 +5356,7 @@ dependencies = [ [[package]] name = "lance-geo" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "datafusion", "geo-traits", @@ -5371,7 +5371,7 @@ dependencies = [ [[package]] name = "lance-index" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arc-swap", "arrow", @@ -5444,7 +5444,7 @@ dependencies = [ [[package]] name = "lance-index-core" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-schema", @@ -5468,7 +5468,7 @@ dependencies = [ [[package]] name = "lance-io" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow", "arrow-array", @@ -5511,7 +5511,7 @@ dependencies = [ [[package]] name = "lance-jsonb" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "byteorder", "ethnum", @@ -5529,7 +5529,7 @@ dependencies = [ [[package]] name = "lance-linalg" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-schema", @@ -5544,7 +5544,7 @@ dependencies = [ [[package]] name = "lance-namespace" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow", "async-trait", @@ -5559,7 +5559,7 @@ dependencies = [ [[package]] name = "lance-namespace-impls" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arc-swap", "arrow", @@ -5614,7 +5614,7 @@ dependencies = [ [[package]] name = "lance-select" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-buffer", @@ -5629,7 +5629,7 @@ dependencies = [ [[package]] name = "lance-table" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow", "arrow-array", @@ -5670,7 +5670,7 @@ dependencies = [ [[package]] name = "lance-testing" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "arrow-array", "arrow-schema", @@ -5684,7 +5684,7 @@ dependencies = [ [[package]] name = "lance-tokenizer" version = "14.0.0-beta.12" -source = "git+https://github.com/lance-format/lance.git?rev=c140ffc9eedee6015145ebc63efa9083068e1887#c140ffc9eedee6015145ebc63efa9083068e1887" +source = "git+https://github.com/lance-format/lance.git?rev=4571ad6059cfe20c8e45998afa01d8749965a23c#4571ad6059cfe20c8e45998afa01d8749965a23c" dependencies = [ "frostem", "icu_segmenter", diff --git a/Cargo.toml b/Cargo.toml index 956624b3b..3c078189c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.94.0" [workspace.dependencies] -lance = { "version" = "=14.0.0-beta.12", default-features = false, "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=14.0.0-beta.12", default-features = false, "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=14.0.0-beta.12", default-features = false, "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=14.0.0-beta.12", "rev" = "c140ffc9eedee6015145ebc63efa9083068e1887", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=14.0.0-beta.12", default-features = false, "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=14.0.0-beta.12", default-features = false, "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=14.0.0-beta.12", default-features = false, "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=14.0.0-beta.12", "rev" = "4571ad6059cfe20c8e45998afa01d8749965a23c", "git" = "https://github.com/lance-format/lance.git" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow diff --git a/docs/src/js/interfaces/IvfRqOptions.md b/docs/src/js/interfaces/IvfRqOptions.md index 6b48cf653..3a809762b 100644 --- a/docs/src/js/interfaces/IvfRqOptions.md +++ b/docs/src/js/interfaces/IvfRqOptions.md @@ -33,6 +33,17 @@ The following distance types are available: *** +### layered? + +```ts +optional layered: boolean; +``` + +Use prefix planes. Unset uses the product default: false in OSS and true +in Enterprise. Ignored for 1-bit and 2-bit codes. + +*** + ### maxIterations? ```ts @@ -58,7 +69,7 @@ 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. +are small integers; the default is 5 bits per dimension. *** diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index d2244b48c..051114d75 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -1695,14 +1695,38 @@ 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, - }), - }); - }); + test.each([ + [1, true], + [2, true], + [5, undefined], + [5, false], + [5, true], + ] as const)( + "should create IVF_RQ with bits=%s layered=%s", + async (numBits, layered) => { + await tbl.createIndex("vec", { + config: Index.ivfRq({ numPartitions: 10, numBits, layered }), + }); + const indices = await tbl.listIndices(); + expect(indices).toHaveLength(1); + expect(indices[0].indexType).toBe("IvfRq"); + const expected = await tbl + .vectorSearch(queryVec) + .bypassVectorIndex() + .limit(10) + .toArray(); + const actual = await tbl + .vectorSearch(queryVec) + .nprobes(10) + .refineFactor(30) + .limit(10) + .toArray(); + const ids = new Set(expected.map((row) => row.id)); + expect( + actual.filter((row) => ids.has(row.id)).length, + ).toBeGreaterThanOrEqual(5); + }, + ); it("should allow me to replace (or not) an existing index", async () => { await tbl.createIndex("id"); diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index 43d1cf1b5..8daeec89b 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -131,10 +131,16 @@ export interface IvfRqOptions { * * 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. + * are small integers; the default is 5 bits per dimension. */ numBits?: number; + /** + * Use prefix planes. Unset uses the product default: false in OSS and true + * in Enterprise. Ignored for 1-bit and 2-bit codes. + */ + layered?: boolean; + /** * Distance type to use to build the index. * @@ -660,6 +666,7 @@ export class Index { options?.numBits, options?.maxIterations, options?.sampleRate, + options?.layered, ), ); } diff --git a/nodejs/src/index.rs b/nodejs/src/index.rs index e5884aa4d..6e5cbf462 100644 --- a/nodejs/src/index.rs +++ b/nodejs/src/index.rs @@ -137,8 +137,12 @@ impl Index { num_bits: Option, max_iterations: Option, sample_rate: Option, + layered: Option, ) -> napi::Result { let mut ivf_rq_builder = IvfRqIndexBuilder::default(); + if let Some(layered) = layered { + ivf_rq_builder = ivf_rq_builder.layered(layered); + } 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); diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index cf3bf8e4f..db7b4919d 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -882,9 +882,13 @@ class IvfRq: num_partitions: int, default sqrt(num_rows) Number of IVF partitions to create. - num_bits: int, default 1 + num_bits: int, default 5 Number of bits to encode each dimension in the RabitQ codebook. + layered: Optional[bool], default None + Use prefix planes. None uses the product default: False in OSS and True + in Enterprise. Ignored for 1-bit and 2-bit codes. + max_iterations: int, default 50 Max iterations to train kmeans when computing IVF partitions. @@ -897,13 +901,14 @@ class IvfRq: distance_type: Literal["l2", "cosine", "dot"] = "l2" num_partitions: Optional[int] = None - num_bits: int = 1 + num_bits: int = 5 max_iterations: int = 50 sample_rate: int = 256 target_partition_size: Optional[int] = None # Name of the accelerator (e.g. "cuda") to use for IVF training. When set, # create_index() dispatches to pylance to build the index on the accelerator. accelerator: Optional[str] = None + layered: Optional[bool] = None # The API reference renders this module with a single mkdocstrings directive, diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 21838ecef..2e5bdd156 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3491,6 +3491,7 @@ class LanceTable(Table): replace=replace, accelerator=acc, num_bits=getattr(config, "num_bits", 8), + layered=getattr(config, "layered", None), m=getattr(config, "m", 20), ef_construction=getattr(config, "ef_construction", 300), target_partition_size=getattr( diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index fe6ebe87a..19bddacd2 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -421,13 +421,29 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable): @pytest.mark.asyncio -async def test_create_ivfrq_index(some_table: AsyncTable): - await some_table.create_index("vector", config=IvfRq(num_bits=1)) +@pytest.mark.parametrize("num_bits", [1, 2, 5]) +@pytest.mark.parametrize("layered", [None, False, True]) +async def test_create_ivfrq_index(some_table: AsyncTable, num_bits, layered): + await some_table.create_index( + "vector", config=IvfRq(num_bits=num_bits, layered=layered) + ) indices = await some_table.list_indices() assert len(indices) == 1 assert indices[0].index_type == "IvfRq" assert indices[0].columns == ["vector"] assert indices[0].name == "vector_idx" + query = list(map(float, range(DIM))) + expected = ( + await some_table.vector_search(query).bypass_vector_index().limit(10).to_arrow() + ) + actual = ( + await some_table.vector_search(query) + .nprobes(NROWS) + .refine_factor(26) + .limit(10) + .to_arrow() + ) + assert len(set(actual["id"].to_pylist()) & set(expected["id"].to_pylist())) >= 5 @pytest.mark.asyncio diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 3e66ac0e0..760b8bfa3 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1476,7 +1476,8 @@ def test_remote_json_fts_config(server_version): assert request["block_size"] == 256 -def test_remote_create_index_new_api(caplog): +@pytest.mark.parametrize("layered", [None, False, True]) +def test_remote_create_index_new_api(caplog, layered): received_requests = [] def handler(request): @@ -1549,7 +1550,9 @@ def test_remote_create_index_new_api(caplog): table.create_index("category", config=BTree()) table.create_index("text", config=FTS(block_size=256)) # IvfRq via new API - table.create_index("vector", config=IvfRq(distance_type="l2")) + table.create_index( + "vector", config=IvfRq(distance_type="l2", layered=layered) + ) table.create_index( "vector", config=IvfPq(distance_type="l2"), replace=False ) @@ -1573,6 +1576,11 @@ def test_remote_create_index_new_api(caplog): "vector", "vector", ] + assert received_requests[3]["num_bits"] == 5 + if layered is None: + assert "layered" not in received_requests[3] + else: + assert received_requests[3]["layered"] is layered assert received_requests[2]["block_size"] == 256 assert received_requests[4]["replace"] is False assert received_requests[5]["replace"] is False diff --git a/python/src/index.rs b/python/src/index.rs index d4e16a67e..414dde515 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -138,6 +138,9 @@ pub fn extract_index_params(source: &Option>) -> PyResult, num_bits: u32, + layered: Option, max_iterations: u32, sample_rate: u32, target_partition_size: Option, diff --git a/rust/lancedb/src/index/vector.rs b/rust/lancedb/src/index/vector.rs index bce77a610..39d51bb9b 100644 --- a/rust/lancedb/src/index/vector.rs +++ b/rust/lancedb/src/index/vector.rs @@ -339,6 +339,8 @@ pub struct IvfRqIndexBuilder { pub(crate) num_partitions: Option, #[serde(skip_serializing_if = "Option::is_none")] pub(crate) num_bits: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) layered: Option, pub(crate) sample_rate: u32, pub(crate) max_iterations: u32, #[serde(skip_serializing_if = "Option::is_none")] @@ -351,6 +353,7 @@ impl Default for IvfRqIndexBuilder { distance_type: DistanceType::L2, num_partitions: None, num_bits: None, + layered: None, sample_rate: 256, max_iterations: 50, target_partition_size: None, @@ -362,10 +365,23 @@ impl IvfRqIndexBuilder { impl_distance_type_setter!(); impl_ivf_params_setter!(); + /// Bits per dimension, defaulting to 5 in both OSS and Enterprise. pub fn num_bits(mut self, num_bits: u32) -> Self { self.num_bits = Some(num_bits); self } + + /// Override the prefix-plane layout default (false in OSS, true in Enterprise). + /// Ignored for 1-bit and 2-bit codes. Leaving this unset preserves the product default. + /// + /// ``` + /// use lancedb::index::vector::IvfRqIndexBuilder; + /// let index = IvfRqIndexBuilder::default().num_bits(5).layered(false); + /// ``` + pub fn layered(mut self, layered: bool) -> Self { + self.layered = Some(layered); + self + } } /// Builder for an IVF HNSW PQ index. diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 119577dc5..5f61317c7 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -7024,6 +7024,7 @@ mod tests { "metric_type": "cosine", "num_partitions": 64, "num_bits": 8, + "layered": false, "sample_rate": 256, "max_iterations": 50, }), @@ -7031,9 +7032,20 @@ mod tests { IvfRqIndexBuilder::default() .distance_type(DistanceType::Cosine) .num_partitions(64) - .num_bits(8), + .num_bits(8) + .layered(false), ), ), + ( + "IVF_RQ", + json!({ + "metric_type": "l2", + "layered": true, + "sample_rate": 256, + "max_iterations": 50, + }), + Index::IvfRq(IvfRqIndexBuilder::default().layered(true)), + ), // HNSW_PQ isn't yet supported on SaaS ("BTREE", json!({}), Index::BTree(Default::default())), ("BITMAP", json!({}), Index::Bitmap(Default::default())), diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 3ac759bdd..c1dbe7fe8 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -15,7 +15,7 @@ use lance::index::vector::utils::infer_vector_dim; use lance_arrow::json::is_json_field; use lance_index::IndexType; use lance_index::scalar::{BuiltinIndexType, ScalarIndexParams}; -use lance_index::vector::bq::RQBuildParams; +use lance_index::vector::bq::{RABIT_DEFAULT_NUM_BITS, RQBuildParams}; use lance_index::vector::hnsw::builder::HnswBuildParams; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; @@ -355,7 +355,10 @@ impl NativeTable { index.sample_rate, index.max_iterations, ); - let rq_params = RQBuildParams::new(index.num_bits.unwrap_or(1) as u8); + let mut rq_params = RQBuildParams::new( + index.num_bits.unwrap_or(u32::from(RABIT_DEFAULT_NUM_BITS)) as u8, + ); + rq_params.layered = index.layered; let lance_idx_params = VectorIndexParams::with_ivf_rq_params( index.distance_type.into(), ivf_params, @@ -469,6 +472,7 @@ impl NativeTable { #[cfg(test)] mod tests { use lance::index::DatasetIndexExt; + use lance::index::vector::{StageParams, VectorIndexParams}; use std::sync::Arc; use std::time::Duration; @@ -492,7 +496,7 @@ mod tests { FtsIndexBuilder, ZoneMapIndexBuilder, }; use crate::index::vector::{ - IvfHnswFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, + IvfHnswFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfRqIndexBuilder, }; use crate::query::{ExecutableQuery, QueryBase}; use crate::table::optimize::{CompactionOptions, OptimizeAction}; @@ -516,6 +520,52 @@ mod tests { Ok(FixedSizeListArray::from(data)) } + #[rstest::rstest] + #[tokio::test] + async fn test_rq_params_preserve_product_default( + #[values(None, Some(1), Some(2), Some(5))] num_bits: Option, + #[values(None, Some(false), Some(true))] layered: Option, + ) { + let dir = tempdir().unwrap(); + let connection = connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let field = Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 8), + true, + ); + let table = connection + .create_empty_table("test", Arc::new(Schema::new(vec![field.clone()]))) + .execute() + .await + .unwrap(); + let mut builder = IvfRqIndexBuilder::default(); + if let Some(num_bits) = num_bits { + builder = builder.num_bits(num_bits); + } + if let Some(layered) = layered { + builder = builder.layered(layered); + } + let params = table + .as_native() + .unwrap() + .make_index_params(&field, Index::IvfRq(builder)) + .await + .unwrap(); + let params = params.as_any().downcast_ref::().unwrap(); + let Some(StageParams::RQ(rq)) = params.stages.last() else { + panic!("expected RQ stage"); + }; + assert_eq!(u32::from(rq.num_bits), num_bits.unwrap_or(5)); + assert_eq!(rq.layered, layered); + assert_eq!( + rq.is_layered(), + num_bits.unwrap_or(5) > 2 && layered.unwrap_or(false) + ); + } + #[tokio::test] async fn test_create_index() { use std::iter::repeat_with;