diff --git a/Cargo.lock b/Cargo.lock index 88aabd4f5..8d7636c80 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3585,7 +3585,7 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-array", "rand 0.9.5", @@ -5060,7 +5060,7 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arc-swap", "arrow", @@ -5134,7 +5134,7 @@ dependencies = [ [[package]] name = "lance-arrow" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" 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?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" 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?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-array", "arrow-schema", @@ -5180,7 +5180,7 @@ dependencies = [ [[package]] name = "lance-bitpacking" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrayref", "crunchy", @@ -5191,7 +5191,7 @@ dependencies = [ [[package]] name = "lance-core" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-array", "arrow-buffer", @@ -5229,7 +5229,7 @@ dependencies = [ [[package]] name = "lance-datafusion" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow", "arrow-array", @@ -5261,7 +5261,7 @@ dependencies = [ [[package]] name = "lance-datagen" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow", "arrow-array", @@ -5279,7 +5279,7 @@ dependencies = [ [[package]] name = "lance-derive" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "proc-macro2", "quote", @@ -5289,7 +5289,7 @@ dependencies = [ [[package]] name = "lance-encoding" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-arith", "arrow-array", @@ -5323,7 +5323,7 @@ dependencies = [ [[package]] name = "lance-file" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-arith", "arrow-array", @@ -5356,7 +5356,7 @@ dependencies = [ [[package]] name = "lance-geo" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "datafusion", "geo-traits", @@ -5371,7 +5371,7 @@ dependencies = [ [[package]] name = "lance-index" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arc-swap", "arrow", @@ -5438,12 +5438,13 @@ dependencies = [ "tokio", "tracing", "twox-hash", + "uuid", ] [[package]] name = "lance-index-core" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-array", "arrow-schema", @@ -5467,7 +5468,7 @@ dependencies = [ [[package]] name = "lance-io" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow", "arrow-array", @@ -5511,7 +5512,7 @@ dependencies = [ [[package]] name = "lance-jsonb" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "byteorder", "ethnum", @@ -5529,7 +5530,7 @@ dependencies = [ [[package]] name = "lance-linalg" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-array", "arrow-schema", @@ -5544,7 +5545,7 @@ dependencies = [ [[package]] name = "lance-namespace" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow", "async-trait", @@ -5559,7 +5560,7 @@ dependencies = [ [[package]] name = "lance-namespace-impls" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arc-swap", "arrow", @@ -5614,7 +5615,7 @@ dependencies = [ [[package]] name = "lance-select" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-array", "arrow-buffer", @@ -5629,7 +5630,7 @@ dependencies = [ [[package]] name = "lance-table" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow", "arrow-array", @@ -5670,7 +5671,7 @@ dependencies = [ [[package]] name = "lance-testing" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "arrow-array", "arrow-schema", @@ -5684,7 +5685,7 @@ dependencies = [ [[package]] name = "lance-tokenizer" version = "14.0.0-beta.18" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.18#f80bf1d1e59f40d249bce2539b78e41d93dee9ec" +source = "git+https://github.com/lance-format/lance.git?rev=823aaac48d42d095205b0db0162e1a5f8aeac843#823aaac48d42d095205b0db0162e1a5f8aeac843" dependencies = [ "frostem", "icu_segmenter", diff --git a/Cargo.toml b/Cargo.toml index c6b9e84fe..d9bc9c796 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.18", default-features = false, "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=14.0.0-beta.18", default-features = false, "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=14.0.0-beta.18", default-features = false, "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=14.0.0-beta.18", "tag" = "v14.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=14.0.0-beta.18", default-features = false, "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=14.0.0-beta.18", default-features = false, "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=14.0.0-beta.18", default-features = false, "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=14.0.0-beta.18", "rev" = "823aaac48d42d095205b0db0162e1a5f8aeac843", "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 6c770494b..0effb7327 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 d39053775..3e16ef938 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3612,6 +3612,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 717490dcd..c1b0ba322 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1477,7 +1477,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): @@ -1550,7 +1551,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 ) @@ -1574,6 +1577,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 7d7a6301a..293262658 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -7333,6 +7333,7 @@ mod tests { "metric_type": "cosine", "num_partitions": 64, "num_bits": 8, + "layered": false, "sample_rate": 256, "max_iterations": 50, }), @@ -7340,9 +7341,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;