diff --git a/Cargo.lock b/Cargo.lock index 87bef3a25..5400ebddd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "rand 0.9.5", @@ -5143,8 +5143,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arc-swap", "arrow", @@ -5217,8 +5217,8 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-buffer", @@ -5240,7 +5240,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-buffer", @@ -5254,7 +5254,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-schema", @@ -5263,8 +5263,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrayref", "crunchy", @@ -5274,8 +5274,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-buffer", @@ -5312,8 +5312,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow", "arrow-array", @@ -5344,8 +5344,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow", "arrow-array", @@ -5362,8 +5362,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "proc-macro2", "quote", @@ -5372,8 +5372,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-arith", "arrow-array", @@ -5406,8 +5406,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-arith", "arrow-array", @@ -5439,8 +5439,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "datafusion", "geo-traits", @@ -5454,8 +5454,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arc-swap", "arrow", @@ -5525,8 +5525,8 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-schema", @@ -5549,8 +5549,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow", "arrow-array", @@ -5592,8 +5592,8 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-schema", @@ -5607,8 +5607,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow", "async-trait", @@ -5622,9 +5622,10 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ + "arc-swap", "arrow", "arrow-ipc", "arrow-schema", @@ -5676,8 +5677,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-buffer", @@ -5691,8 +5692,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow", "arrow-array", @@ -5732,8 +5733,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "arrow-array", "arrow-schema", @@ -5746,8 +5747,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "13.0.0-beta.15" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e" +version = "13.0.0-beta.18" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.18#b646eda6a891f481b6f6d6baef9cd17656aef65b" dependencies = [ "frostem", "icu_segmenter", diff --git a/Cargo.toml b/Cargo.toml index 1bf7ce86a..0f8bbfc93 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.91.0" [workspace.dependencies] -lance = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=13.0.0-beta.18", default-features = false, "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=13.0.0-beta.18", default-features = false, "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=13.0.0-beta.18", default-features = false, "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=13.0.0-beta.18", "tag" = "v13.0.0-beta.18", "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/classes/BlobFile.md b/docs/src/js/classes/BlobFile.md index 84a596d1a..ee9f9ad28 100644 --- a/docs/src/js/classes/BlobFile.md +++ b/docs/src/js/classes/BlobFile.md @@ -10,16 +10,50 @@ A lazy handle to blob bytes. Create one with [Table.fetchBlobFiles](Table.md#fet ## Methods +### close() + +```ts +close(): Promise +``` + +Releases the handle. Reads after `close()` fail. Calling it again does +nothing. + +#### Returns + +`Promise`<`void`> + +*** + +### isClosed() + +```ts +isClosed(): boolean +``` + +Returns true after [BlobFile.close](BlobFile.md#close). + +#### Returns + +`boolean` + +*** + ### read() ```ts -read(): Promise +read(maxBytes?): Promise ``` -Reads from the cursor to the end and advances the cursor. +Reads from the cursor and advances the cursor. -A second call returns an empty buffer. [BlobFile.readRange](BlobFile.md#readrange) does -not move the cursor. +Reads to the end when `maxBytes` is omitted, or at most `maxBytes` bytes +otherwise. Returns an empty buffer at the end of the blob. +[BlobFile.readRange](BlobFile.md#readrange) does not move the cursor. + +#### Parameters + +* **maxBytes?**: `bigint` #### Returns @@ -49,6 +83,45 @@ Fails when `end` is past the blob size. Does not move the cursor. *** +### readRanges() + +```ts +readRanges(ranges): Promise +``` + +Reads several half-open byte ranges. Returns one buffer per range, in +the order given. + +Fails when any `end` is past the blob size. Does not move the cursor. + +#### Parameters + +* **ranges**: [`BlobRange`](../type-aliases/BlobRange.md)[] + +#### Returns + +`Promise`<`Buffer`[]> + +*** + +### seek() + +```ts +seek(position): Promise +``` + +Moves the cursor to `position`, in bytes from the start of the blob. + +#### Parameters + +* **position**: `bigint` + +#### Returns + +`Promise`<`void`> + +*** + ### size() ```ts @@ -60,3 +133,17 @@ Returns the blob size in bytes. #### Returns `bigint` + +*** + +### tell() + +```ts +tell(): Promise +``` + +Returns the cursor position in bytes. + +#### Returns + +`Promise`<`bigint`> diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 2f49578b1..9b9f76334 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -158,6 +158,7 @@ - [AnalyzePlanDistributedMetrics](type-aliases/AnalyzePlanDistributedMetrics.md) - [BaseTokenizer](type-aliases/BaseTokenizer.md) - [BlobOptions](type-aliases/BlobOptions.md) +- [BlobRange](type-aliases/BlobRange.md) - [Data](type-aliases/Data.md) - [DataLike](type-aliases/DataLike.md) - [FieldLike](type-aliases/FieldLike.md) diff --git a/docs/src/js/type-aliases/BlobRange.md b/docs/src/js/type-aliases/BlobRange.md new file mode 100644 index 000000000..85cd48242 --- /dev/null +++ b/docs/src/js/type-aliases/BlobRange.md @@ -0,0 +1,27 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BlobRange + +# Type Alias: BlobRange + +```ts +type BlobRange: object; +``` + +A half-open byte range `[start, end)` for [BlobFile.readRanges](../classes/BlobFile.md#readranges). + +## Type declaration + +### end + +```ts +end: bigint; +``` + +### start + +```ts +start: bigint; +``` diff --git a/java/pom.xml b/java/pom.xml index 752dacefd..7b9009ccd 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 13.0.0-beta.14 + 13.0.0-beta.18 false 2.30.0 1.7 diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 5f15297ff..f718b2c89 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -92,6 +92,38 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( await expect(table.countRows()).resolves.toBe(3); }); + it("creates and adds Arrow tables with large string and binary columns", async () => { + const data = new arrow.Table({ + text: arrow.vectorFromArray(["alpha", "beta"], new arrow.LargeUtf8()), + bytes: arrow.vectorFromArray( + [Buffer.from("one"), Buffer.from("two")], + new arrow.LargeBinary(), + ), + }); + const conn = await connect(tmpDir.name); + const largeTable = await conn.createTable("large_columns", data); + await largeTable.add(data); + + const rows = await largeTable.query().toArray(); + expect(rows.map((row) => row.text)).toEqual([ + "alpha", + "beta", + "alpha", + "beta", + ]); + expect(rows.map((row) => Buffer.from(row.bytes).toString())).toEqual([ + "one", + "two", + "one", + "two", + ]); + const fields = (await largeTable.schema()).fields; + expect(fields.map((field) => field.type.typeId)).toEqual([ + arrow.Type.LargeUtf8, + arrow.Type.LargeBinary, + ]); + }); + it("should support a foreign Float64 vector schema end to end", async () => { const conn = await connect(tmpDir.name); const schema = new arrow.Schema([ @@ -2613,6 +2645,27 @@ describe("when dealing with blob columns", () => { ); }); + it("adds a LargeBinary Arrow column to a blob table", async () => { + const db = await connect(tmpDir.name); + const schema = new Schema([new Field("id", new Int64()), blob("image")]); + const table = await db.createTable( + "blobs", + [{ id: 1n, image: Buffer.from("alpha") }], + { schema }, + ); + const payload = Buffer.from("beta"); + const data = new arrow18.Table({ + id: arrow18.vectorFromArray([2n], new arrow18.Int64()), + image: arrow18.vectorFromArray([payload], new arrow18.LargeBinary()), + }); + await table.add(data); + + const rows = await table.query().withRowId().toArray(); + const rowId = rows.find((row) => row.id === 2n)!._rowid as bigint; + const [actual] = await table.fetchBlobs("image", [rowId]); + expect(actual).toEqual(payload); + }); + it("reads a half-open range", async () => { const { table, rowIds } = await openBlobTable(); const files = await table.fetchBlobFiles("image", rowIds); @@ -2637,6 +2690,74 @@ describe("when dealing with blob columns", () => { ).rejects.toThrow(/exceeds blob size/); }); + it("reads at most maxBytes and advances the cursor", async () => { + const { table, rowIds } = await openBlobTable(); + const [handle] = await table.fetchBlobFiles("image", rowIds); + expect((await handle!.read(2n)).toString()).toBe("al"); + expect(await handle!.tell()).toBe(2n); + expect((await handle!.read(2n)).toString()).toBe("ph"); + expect((await handle!.read(10n)).toString()).toBe("a"); + expect(await handle!.tell()).toBe(5n); + expect(await handle!.read(1n)).toEqual(Buffer.alloc(0)); + }); + + it("seeks the cursor", async () => { + const { table, rowIds, alpha } = await openBlobTable(); + const [handle] = await table.fetchBlobFiles("image", rowIds); + await handle!.seek(3n); + expect(await handle!.tell()).toBe(3n); + expect((await handle!.read()).toString()).toBe("ha"); + await handle!.seek(0n); + expect(await handle!.read()).toEqual(alpha); + }); + + it("rejects a negative seek position", async () => { + const { table, rowIds } = await openBlobTable(); + const [handle] = await table.fetchBlobFiles("image", rowIds); + await expect(handle!.seek(-1n)).rejects.toThrow(/cannot be negative/); + }); + + it("readRanges returns one buffer per range in order", async () => { + const { table, rowIds } = await openBlobTable(); + const [handle] = await table.fetchBlobFiles("image", rowIds); + const buffers = await handle!.readRanges([ + { start: 3n, end: 5n }, + { start: 0n, end: 2n }, + { start: 3n, end: 5n }, + ]); + expect(buffers.map((b) => b.toString())).toEqual(["ha", "al", "ha"]); + expect(await handle!.tell()).toBe(0n); + }); + + it("fails when a readRanges end is past the blob size", async () => { + const { table, rowIds, alpha } = await openBlobTable(); + const [handle] = await table.fetchBlobFiles("image", rowIds); + await expect( + handle!.readRanges([ + { start: 0n, end: 1n }, + { start: 0n, end: BigInt(alpha.length + 1) }, + ]), + ).rejects.toThrow(/exceeds blob size/); + }); + + it("fails reads after close and closes once", async () => { + const { table, rowIds } = await openBlobTable(); + const [handle] = await table.fetchBlobFiles("image", rowIds); + expect(handle!.isClosed()).toBe(false); + await handle!.close(); + expect(handle!.isClosed()).toBe(true); + await expect(handle!.close()).resolves.toBeUndefined(); + await expect(handle!.read()).rejects.toThrow(/already closed/); + await expect(handle!.read(1n)).rejects.toThrow(/already closed/); + await expect(handle!.readRange(0n, 1n)).rejects.toThrow(/already closed/); + await expect(handle!.readRanges([{ start: 0n, end: 1n }])).rejects.toThrow( + /already closed/, + ); + await expect(handle!.seek(0n)).rejects.toThrow(/already closed/); + await expect(handle!.tell()).rejects.toThrow(/already closed/); + expect(handle!.size()).toBe(5n); + }); + it("rejects fetchBlobs on a non-blob column", async () => { const { table, rowIds } = await openBlobTable(); await expect(table.fetchBlobs("id", rowIds)).rejects.toThrow(/blob/i); diff --git a/nodejs/lancedb/blob.ts b/nodejs/lancedb/blob.ts index 244faac09..fc8c43196 100644 --- a/nodejs/lancedb/blob.ts +++ b/nodejs/lancedb/blob.ts @@ -141,13 +141,14 @@ export class BlobFile { } /** - * Reads from the cursor to the end and advances the cursor. + * Reads from the cursor and advances the cursor. * - * A second call returns an empty buffer. {@link BlobFile.readRange} does - * not move the cursor. + * Reads to the end when `maxBytes` is omitted, or at most `maxBytes` bytes + * otherwise. Returns an empty buffer at the end of the blob. + * {@link BlobFile.readRange} does not move the cursor. */ - read(): Promise { - return this.inner.read(); + read(maxBytes?: bigint): Promise { + return this.inner.read(maxBytes); } /** @@ -158,8 +159,47 @@ export class BlobFile { readRange(start: bigint, end: bigint): Promise { return this.inner.readRange(start, end); } + + /** + * Reads several half-open byte ranges. Returns one buffer per range, in + * the order given. + * + * Fails when any `end` is past the blob size. Does not move the cursor. + */ + readRanges(ranges: BlobRange[]): Promise { + return this.inner.readRanges(ranges); + } + + /** Moves the cursor to `position`, in bytes from the start of the blob. */ + seek(position: bigint): Promise { + return this.inner.seek(position); + } + + /** Returns the cursor position in bytes. */ + tell(): Promise { + return this.inner.tell(); + } + + /** + * Releases the handle. Reads after `close()` fail. Calling it again does + * nothing. + */ + close(): Promise { + return this.inner.close(); + } + + /** Returns true after {@link BlobFile.close}. */ + isClosed(): boolean { + return this.inner.isClosed(); + } } +/** A half-open byte range `[start, end)` for {@link BlobFile.readRanges}. */ +export type BlobRange = { + start: bigint; + end: bigint; +}; + export function coerceBlobValue(value: unknown): BlobInput | null { if (value == null) { return null; diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index ed6889232..2e444784b 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -84,7 +84,7 @@ export { } from "./arrow"; export { blob, isBlobField, BlobFile } from "./blob"; -export type { BlobOptions } from "./blob"; +export type { BlobOptions, BlobRange } from "./blob"; export { Connection, diff --git a/nodejs/lancedb/sanitize.ts b/nodejs/lancedb/sanitize.ts index 3e5d583c4..98a144451 100644 --- a/nodejs/lancedb/sanitize.ts +++ b/nodejs/lancedb/sanitize.ts @@ -43,6 +43,8 @@ import { Interval, IntervalDayTime, IntervalYearMonth, + LargeBinary, + LargeUtf8, List, Map_, Null, @@ -465,8 +467,12 @@ function sanitizeTypeById( return sanitizeFloat(typeLike); case Type.Binary: return new Binary(); + case Type.LargeBinary: + return new LargeBinary(); case Type.Utf8: return new Utf8(); + case Type.LargeUtf8: + return new LargeUtf8(); case Type.Bool: return new Bool(); case Type.Decimal: diff --git a/nodejs/src/blob.rs b/nodejs/src/blob.rs index 230a11bae..413bf6d97 100644 --- a/nodejs/src/blob.rs +++ b/nodejs/src/blob.rs @@ -3,6 +3,7 @@ use std::ops::Range; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use arrow_array::{Array, LargeBinaryArray}; use lancedb::blob::BlobFile as LanceBlobFile; @@ -11,15 +12,31 @@ use napi_derive::napi; use crate::error::convert_error; +#[napi(object)] +pub struct BlobRange { + pub start: BigInt, + pub end: BigInt, +} + #[napi] pub struct BlobFile { inner: Arc, + closed: AtomicBool, } impl BlobFile { pub(crate) fn new(inner: LanceBlobFile) -> Self { Self { inner: Arc::new(inner), + closed: AtomicBool::new(false), + } + } + + fn ensure_open(&self) -> napi::Result<()> { + if self.closed.load(Ordering::Acquire) { + Err(napi::Error::from_reason("blob file is already closed")) + } else { + Ok(()) } } } @@ -32,13 +49,23 @@ impl BlobFile { } #[napi] - pub async fn read(&self) -> napi::Result { - let bytes = self.inner.read().await.map_err(|err| convert_error(&err))?; + pub async fn read(&self, max_bytes: Option) -> napi::Result { + self.ensure_open()?; + let bytes = match max_bytes { + None => self.inner.read().await, + Some(max_bytes) => { + let max_bytes = usize::try_from(parse_u64(max_bytes, "maxBytes")?) + .map_err(|_| napi::Error::from_reason("maxBytes is too large"))?; + self.inner.read_up_to(max_bytes).await + } + } + .map_err(|err| convert_error(&err))?; Ok(Buffer::from(bytes.as_ref())) } #[napi] pub async fn read_range(&self, start: BigInt, end: BigInt) -> napi::Result { + self.ensure_open()?; let range = bigint_range(start, end)?; let bytes = self .inner @@ -47,6 +74,54 @@ impl BlobFile { .map_err(|err| convert_error(&err))?; Ok(Buffer::from(bytes.as_ref())) } + + #[napi] + pub async fn read_ranges(&self, ranges: Vec) -> napi::Result> { + self.ensure_open()?; + let ranges = ranges + .into_iter() + .map(|range| bigint_range(range.start, range.end)) + .collect::>>()?; + let buffers = self + .inner + .read_ranges(&ranges) + .await + .map_err(|err| convert_error(&err))?; + Ok(buffers + .iter() + .map(|bytes| Buffer::from(bytes.as_ref())) + .collect()) + } + + #[napi] + pub async fn seek(&self, position: BigInt) -> napi::Result<()> { + self.ensure_open()?; + let position = parse_u64(position, "position")?; + self.inner + .seek(position) + .await + .map_err(|err| convert_error(&err)) + } + + #[napi] + pub async fn tell(&self) -> napi::Result { + self.ensure_open()?; + let position = self.inner.tell().await.map_err(|err| convert_error(&err))?; + Ok(BigInt::from(position)) + } + + #[napi] + pub async fn close(&self) -> napi::Result<()> { + if self.closed.swap(true, Ordering::AcqRel) { + return Ok(()); + } + self.inner.close().await.map_err(|err| convert_error(&err)) + } + + #[napi] + pub fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } } fn bigint_range(start: BigInt, end: BigInt) -> napi::Result> { diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index f9fe91d35..f542fe856 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -658,6 +658,7 @@ impl JsFullTextQuery { FtsQuery::Phrase(_) => "phrase".to_string(), FtsQuery::Boost(_) => "boost".to_string(), FtsQuery::MultiMatch(_) => "multi_match".to_string(), + FtsQuery::CombinedFields(_) => "combined_fields".to_string(), FtsQuery::Boolean(_) => "boolean".to_string(), } } diff --git a/python/python/lancedb/embeddings/base.py b/python/python/lancedb/embeddings/base.py index 149b9c089..fbe43860d 100644 --- a/python/python/lancedb/embeddings/base.py +++ b/python/python/lancedb/embeddings/base.py @@ -182,7 +182,11 @@ class EmbeddingFunction(BaseModel, ABC): def VectorField(self, **kwargs): """ Creates a pydantic Field that can automatically annotate - the target vector column for this embedding function + the target vector column for this embedding function. + + The field can be annotated as ``list[float]`` for compatibility with + static type checkers. LanceDB will infer the fixed vector dimension from + this embedding function. """ return Field(json_schema_extra={"vector_column_for": self}, **kwargs) diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index 528d865d7..313c514c3 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -88,6 +88,17 @@ def Vector( ... pa.field("url", pa.utf8(), False), ... pa.field("embeddings", pa.list_(pa.float32(), 768)) ... ]) + + Notes + ----- + ``Vector`` creates a type dynamically, so calls such as ``Vector(768)`` are + not valid static type annotations. For an embedding field, use the standard + ``list[float]`` annotation when running mypy; ``VectorField`` supplies the + fixed dimension to LanceDB:: + + class MyModel(LanceModel): + text: str = embeddings.SourceField() + vector: list[float] = embeddings.VectorField() """ # TODO: make a public parameterized type. @@ -313,6 +324,10 @@ def _unwrap_optional_annotation(annotation: Any) -> Any | None: def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType: """Convert a Pydantic FieldInfo to Arrow DataType""" + embedding_vector_type = _embedding_vector_to_arrow_type(field) + if embedding_vector_type is not None: + return embedding_vector_type + unwrapped = _unwrap_optional_annotation(field.annotation) if unwrapped is not None: return _pydantic_type_to_arrow_type(unwrapped, field) @@ -326,8 +341,32 @@ def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType: return _pydantic_type_to_arrow_type(field.annotation, field) +def _embedding_vector_to_arrow_type(field: FieldInfo) -> pa.DataType | None: + """Infer a fixed-size vector type from ``VectorField`` metadata.""" + if not _is_embedding_vector_annotation(field): + return None + + function = get_extras(field, "vector_column_for") + return pa.list_(pa.float32(), function.ndims()) + + +def _is_embedding_vector_annotation(field: FieldInfo) -> bool: + if get_extras(field, "vector_column_for") is None: + return False + + annotation = _unwrap_optional_annotation(field.annotation) + if annotation is None: + annotation = field.annotation + + origin = getattr(annotation, "__origin__", None) + args = getattr(annotation, "__args__", ()) + return origin is list and args == (float,) + + def is_nullable(field: FieldInfo) -> bool: """Check if a Pydantic FieldInfo is nullable.""" + if _is_embedding_vector_annotation(field): + return True if _unwrap_optional_annotation(field.annotation) is not None: return True if isinstance(field.annotation, (_GenericAlias, GenericAlias)): diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index a54b2c479..3b3875fb7 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -261,6 +261,101 @@ IndexConfigType = Union[ KNOWN_METRICS = {"l2", "cosine", "dot", "hamming"} +def _blob_value_to_storage(value: Any) -> Optional[dict]: + """Keep a Python blob's inline data or external URI before Arrow infers it.""" + if value is None: + return None + if isinstance(value, (bytes, bytearray, memoryview)): + return {"data": bytes(value)} + if isinstance(value, str): + if not value: + raise ValueError("Blob uri cannot be empty") + return {"uri": value} + if isinstance(value, dict): + unknown = value.keys() - {"data", "uri", "position", "size"} + if unknown: + raise ValueError(f"Unknown blob fields: {sorted(unknown)}") + return value + + try: + from lance.blob import Blob + except ModuleNotFoundError as err: + if err.name not in ("lance", "lance.blob"): + raise + else: + if isinstance(value, Blob): + return { + "data": value.data, + "uri": value.uri, + "position": value.position, + "size": value.size, + } + raise TypeError(f"Unsupported blob value: {type(value).__name__}") + + +def _blob_input_to_arrow(data: Any, schema: Optional[pa.Schema]) -> Optional[pa.Table]: + """Convert Python rows with blob fields before Arrow loses their value types.""" + if schema is None: + return None + + if isinstance(data, list) and data and isinstance(data[0], dict): + names = { + field.name + for field in schema + if is_blob_v2_field(field) + and any(isinstance(row, dict) and field.name in row for row in data) + } + if not names: + return None + values = {name: [row.get(name) for row in data] for name in names} + rows = [{**row, **{name: None for name in names}} for row in data] + table = pa.Table.from_pylist(rows) + elif _check_for_pandas(data) and isinstance(data, pd.DataFrame): + names = { + field.name + for field in schema + if is_blob_v2_field(field) and field.name in data.columns + } + if not names: + return None + values = { + name: [ + None + if value is pd.NA or isinstance(value, float) and np.isnan(value) + else value + for value in data[name] + ] + for name in names + } + table = pa.Table.from_pandas( + data.assign(**{name: None for name in names}), preserve_index=False + ).replace_schema_metadata(None) + else: + return None + + for field in schema: + if field.name not in names: + continue + storage_type = ( + field.type.storage_type + if isinstance(field.type, pa.ExtensionType) + else field.type + ) + storage = pa.array( + [_blob_value_to_storage(value) for value in values[field.name]], + type=storage_type, + ) + column = ( + pa.ExtensionArray.from_storage(field.type, storage) + if isinstance(field.type, pa.ExtensionType) + else storage + ) + table = table.set_column( + table.schema.get_field_index(field.name), field, column + ) + return table + + def _into_pyarrow_reader( data, schema: Optional[pa.Schema] = None ) -> pa.RecordBatchReader: @@ -303,9 +398,14 @@ def _into_pyarrow_reader( return pa.Table.from_batches(data).to_reader() else: data = _serialize_json_values(data, schema) - return pa.Table.from_pylist(data).to_reader() + table = _blob_input_to_arrow(data, schema) + return ( + table if table is not None else pa.Table.from_pylist(data) + ).to_reader() elif _check_for_pandas(data) and isinstance(data, pd.DataFrame): - table = pa.Table.from_pandas(data, preserve_index=False) + table = _blob_input_to_arrow(data, schema) + if table is None: + table = pa.Table.from_pandas(data, preserve_index=False) # Do not serialize Pandas metadata meta = table.schema.metadata if table.schema.metadata is not None else {} meta = {k: v for k, v in meta.items() if k != b"pandas"} @@ -5907,6 +6007,9 @@ class AsyncTable: """ schema = await self.schema() data = _serialize_json_values(data, schema) + blob_table = _blob_input_to_arrow(data, schema) + if blob_table is not None: + data = blob_table if on_bad_vectors is None: on_bad_vectors = "error" if fill_value is None: diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 6f0caf80b..b3508d683 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -7,10 +7,11 @@ import sys import textwrap import lance +import pandas as pd import pyarrow as pa import pyarrow.compute as pc import pytest -from lance.blob import BlobType as LanceBlobType +from lance.blob import Blob, BlobType as LanceBlobType import lancedb from lancedb._blob import ( @@ -1514,3 +1515,129 @@ def test_add_external_uri_string_round_trips_with_flag(tmp_path): hits = table.search().to_arrow() blobs = table.fetch_blobs("image", hits) assert blobs[0].as_py() == payload + + +@pytest.mark.parametrize("as_pandas", [False, True]) +def test_add_bytes_and_uri_in_one_batch_preserves_external_uri(tmp_path, as_pandas): + payload = b"external-payload" + blob_path = tmp_path / "payload.bin" + blob_path.write_bytes(payload) + rows = [ + {"id": 1, "image": b"inline"}, + {"id": 2, "image": blob_path.as_uri()}, + ] + db = lancedb.connect(tmp_path / "db") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("mixed", schema=schema) + + table.add( + pd.DataFrame(rows) if as_pandas else rows, + allow_external_blob_outside_bases=True, + ) + + by_id = _row_ids_by_id(table) + assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [ + b"inline", + payload, + ] + + +@pytest.mark.parametrize("as_pandas", [False, True]) +def test_add_mixed_python_blob_values_preserves_external_uris(tmp_path, as_pandas): + payload = b"external-payload" + blob_path = tmp_path / "payload.bin" + blob_path.write_bytes(payload) + uri = blob_path.as_uri() + rows = [ + {"id": 1, "image": b"inline"}, + {"id": 2, "image": uri}, + {"id": 3, "image": {"data": b"from-dict"}}, + {"id": 4, "image": Blob.from_bytes(b"from-blob")}, + {"id": 5, "image": Blob.from_uri(uri)}, + {"id": 6, "image": {"uri": uri}}, + {"id": 7, "image": None}, + ] + db = lancedb.connect(tmp_path / "db") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("mixed", schema=schema) + + table.add( + pd.DataFrame(rows) if as_pandas else rows, + allow_external_blob_outside_bases=True, + ) + + by_id = _row_ids_by_id(table) + assert table.fetch_blobs("image", [by_id[i] for i in range(1, 8)]).to_pylist() == [ + b"inline", + payload, + b"from-dict", + b"from-blob", + payload, + payload, + None, + ] + + +@pytest.mark.parametrize("as_pandas", [False, True]) +def test_merge_insert_mixed_python_blobs_validates_external_uri(tmp_path, as_pandas): + blob_path = tmp_path / "payload.bin" + blob_path.write_bytes(b"external-payload") + table = _blob_table("merge_mixed_uri", [{"id": 1, "image": b"before"}]) + rows = [ + {"id": 1, "image": b"updated"}, + {"id": 2, "image": blob_path.as_uri()}, + ] + + with pytest.raises(ValueError, match="allow_external_blob_outside_bases"): + ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(pd.DataFrame(rows) if as_pandas else rows) + ) + assert table.count_rows() == 1 + + +@pytest.mark.parametrize("as_pandas", [False, True]) +def test_merge_insert_python_blob_dicts_and_objects(as_pandas): + table = _blob_table("merge_python_blobs", [{"id": 1, "image": b"before"}]) + rows = [ + {"id": 1, "image": b"updated"}, + {"id": 2, "image": {"data": b"from-dict"}}, + {"id": 3, "image": Blob.from_bytes(b"from-blob")}, + ] + + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(pd.DataFrame(rows) if as_pandas else rows) + ) + + assert result.num_updated_rows == 1 + assert result.num_inserted_rows == 2 + by_id = _row_ids_by_id(table) + assert table.fetch_blobs("image", [by_id[i] for i in range(1, 4)]).to_pylist() == [ + b"updated", + b"from-dict", + b"from-blob", + ] + + +def test_create_table_with_python_blob_dict_and_object(): + db = lancedb.connect("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table( + "create_python_blobs", + data=[ + {"id": 1, "image": {"data": b"from-dict"}}, + {"id": 2, "image": Blob.from_bytes(b"from-blob")}, + ], + schema=schema, + ) + + by_id = _row_ids_by_id(table) + assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [ + b"from-dict", + b"from-blob", + ] diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index ca03940f2..81674d735 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -9,6 +9,7 @@ from typing import List, Optional, Tuple import pyarrow as pa import pydantic import pytest +from lancedb.conftest import MockTextEmbeddingFunction from lancedb.pydantic import ( LanceModel, MultiVector, @@ -425,6 +426,25 @@ def test_bare_vector_raises_clear_error(): exec("class TestModel(LanceModel):\n vector: Vector", namespace) +def test_embedding_vector_list_annotation(): + embedding = MockTextEmbeddingFunction.create() + + class StaticTypingModel(LanceModel): + text: str = embedding.SourceField() + vector: list[float] = embedding.VectorField() + + schema = pydantic_to_schema(StaticTypingModel) + assert schema == pa.schema( + [ + pa.field("text", pa.utf8(), False), + pa.field("vector", pa.list_(pa.float32(), embedding.ndims()), True), + ] + ) + + model = StaticTypingModel(text="hello", vector=[0.0] * embedding.ndims()) + assert model.vector == [0.0] * embedding.ndims() + + def test_fixed_size_list_field(): class TestModel(pydantic.BaseModel): vec: Vector(16) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index e32dbe8ca..cd4b78a1b 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2403,11 +2403,16 @@ def blob_remote_table(*, server_version=Version("0.5.0")): elif request.path.startswith("/v1/table/test/blob/image/"): path = request.path.partition("?")[0] row_id = int(path.split("/")[-2]) - payload = {10: b"alpha", 20: None, 30: b"gamma"}[row_id] + payload = {10: b"alpha", 20: None, 30: b"gamma", 40: b""}[row_id] if payload is None: request.send_response(204) request.end_headers() return + if not payload: + request.send_response(416) + request.send_header("Content-Range", "bytes */0") + request.end_headers() + return byte_range = request.headers["Range"].removeprefix("bytes=") start_text, end_text = byte_range.split("-", maxsplit=1) start = int(start_text) @@ -2416,6 +2421,9 @@ def blob_remote_table(*, server_version=Version("0.5.0")): request.send_response(206) request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}") request.send_header("Content-Length", str(len(chunk))) + request.send_header( + "x-lancedb-version", str(BLOB_DESCRIBE_RESPONSE["version"]) + ) request.end_headers() request.wfile.write(chunk) elif request.path == "/v1/table/test/query/": @@ -2459,17 +2467,21 @@ def test_remote_blob_columns_and_fetch(): def test_remote_blob_files_are_lazy_seekable_handles(): with blob_remote_table() as table: - files = table.fetch_blob_files("image", [10, 20, 30]) + files = table.fetch_blob_files("image", [10, 20, 30, 40]) - assert len(files) == 3 - alpha, null_row, gamma = files + assert len(files) == 4 + alpha, null_row, gamma, empty = files assert null_row is None assert alpha is not None assert gamma is not None + assert empty is not None assert alpha.size() == 5 assert alpha.read_range(1, 3) == b"lph" gamma.seek(2) assert gamma.read() == b"mma" + assert empty.size() == 0 + assert empty.read() == b"" + assert empty.read_range(0, 0) == b"" alpha.close() assert alpha.closed with pytest.raises(RuntimeError, match="already closed"): diff --git a/python/src/query.rs b/python/src/query.rs index 2398376d5..3a4737ee0 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -180,6 +180,9 @@ impl<'py> IntoPyObject<'py> for PyLanceDB { .expect("Failed to import namespace"); match self.0 { + FtsQuery::CombinedFields(_) => Err(PyValueError::new_err( + "Combined-fields FTS queries are not yet supported in Python", + )), FtsQuery::Match(query) => { let kwargs = PyDict::new(py); kwargs.set_item("boost", query.boost)?; diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index 616bd42d4..7dec300a2 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -13,7 +13,7 @@ mod duplicate_pairs; mod grouped; mod grouped_units; mod partitioned; -mod vector_dedup; +pub(crate) mod vector_dedup; pub use duplicate_pairs::{VectorSource, VectorSourceKind}; pub use grouped::IVF_PARTITION; /// Refreshing a view grouped by [`IVF_PARTITION`] in units, one per index diff --git a/rust/lancedb/src/materialized_view/vector_dedup.rs b/rust/lancedb/src/materialized_view/vector_dedup.rs index 38c007076..dc2759433 100644 --- a/rust/lancedb/src/materialized_view/vector_dedup.rs +++ b/rust/lancedb/src/materialized_view/vector_dedup.rs @@ -156,7 +156,7 @@ fn edge_schema() -> SchemaRef { /// Sorting canonical edges makes the first unremoved endpoint the smallest /// unassigned representative. The sort spills; only vertex masks stay in RAM. -async fn select_representatives( +pub async fn select_representatives( stream: SendableRecordBatchStream, memory: usize, ) -> Result<(RoaringTreemap, RoaringTreemap)> { diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 195c18ea5..16c546147 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -96,6 +96,8 @@ const SCHEMA_SELECTOR_CHANGED: &str = "table selector changed while fetching sch fn fts_query_requires_document_granularity_support(query: &FtsQuery) -> bool { match query { + // Combined-fields queries do not expose a document granularity option. + FtsQuery::CombinedFields(_) => false, FtsQuery::Match(query) => query .document_granularity .is_some_and(|granularity| granularity.is_list_element()), @@ -5326,6 +5328,7 @@ mod tests { "bytes=0-0" => http::Response::builder() .status(206) .header(reqwest::header::CONTENT_RANGE, "bytes 0-0/10") + .header(VERSION_HEADER, "42") .body(b"0".to_vec()) .unwrap(), "bytes=0-" => http::Response::builder() @@ -6338,6 +6341,53 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn test_query_combined_fields_uses_structured_fts() { + use lance_index::scalar::inverted::query::CombinedFieldsQuery; + + let table = + Table::new_with_handler_version("my_table", semver::Version::new(0, 3, 0), |request| { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + body["full_text_query"]["query"], + serde_json::json!({ + "combined_fields": { + "query": "hello world", + "columns": ["title", "text"], + "boost": [1.0, 1.0], + "operator": "Or" + } + }) + ); + + let data = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(write_ipc_file(&data)) + .unwrap() + }); + + table + .query() + .full_text_search(FullTextSearchQuery::new_query( + CombinedFieldsQuery::try_new( + "hello world".into(), + vec!["title".into(), "text".into()], + ) + .unwrap() + .into(), + )) + .execute() + .await + .unwrap(); + } + #[tokio::test] async fn test_query_row_document_granularity_uses_structured_fts() { let table = diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index 29023eeed..07aeeae63 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -22,7 +22,8 @@ use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; use crate::table::BaseTable; use super::{ - FreshnessHeaders, FreshnessState, ReadSnapshot, RemoteTable, freshness_headers_snapshot, + FreshnessHeaders, FreshnessState, ReadSnapshot, RemoteTable, VERSION_HEADER, + freshness_headers_snapshot, }; // The Cloud route rejects larger row-id lists. Its separate 64 MiB byte limit @@ -50,6 +51,7 @@ trait BlobRangeRequester: Send + Sync + std::fmt::Debug { &self, range_header: &str, mode: RangeRequestMode, + version: Option, ) -> Result<(String, Response)>; } @@ -71,13 +73,14 @@ impl BlobRangeRequester for TableBlobRangeRequester { &self, range_header: &str, mode: RangeRequestMode, + version: Option, ) -> Result<(String, Response)> { let freshness_request = freshness_headers_snapshot(&self.freshness, self.read_consistency_interval); let mut request = freshness_request .apply(self.client.get(&self.path)) .header(header::RANGE, range_header); - if let Some(version) = self.version { + if let Some(version) = version.or(self.version) { request = request.query(&[("version", version)]); } if let Some(branch) = &self.branch { @@ -112,21 +115,25 @@ struct RemoteBlobState { } /// Seekable Cloud blob handle over HTTP Range. +/// +/// Nonempty handles read the table version returned by their size probe. #[derive(Debug)] pub struct RemoteBlobFile { requester: Arc, state: Mutex, closed: AtomicBool, size: u64, + version: Option, } impl RemoteBlobFile { - fn new(requester: Arc, size: u64) -> Self { + fn new(requester: Arc, size: u64, version: Option) -> Self { Self { requester, state: Mutex::new(RemoteBlobState::default()), closed: AtomicBool::new(false), size, + version, } } @@ -175,7 +182,7 @@ impl RemoteBlobFile { let range_header = format!("bytes={}-{}", range.start, range.end - 1); let (request_id, response) = self .requester - .request_range(&range_header, RangeRequestMode::DataRead) + .request_range(&range_header, RangeRequestMode::DataRead, self.version) .await .map_err(remote_blob_error)?; self.ensure_open()?; @@ -251,7 +258,7 @@ impl RemoteBlobFile { let range_header = format!("bytes={cursor}-"); let (request_id, response) = self .requester - .request_range(&range_header, RangeRequestMode::DataRead) + .request_range(&range_header, RangeRequestMode::DataRead, self.version) .await .map_err(remote_blob_error)?; self.ensure_open()?; @@ -579,12 +586,25 @@ async fn probe_blob_files( const BLOB_REQUEST_CONCURRENCY: usize = 8; +fn probe_blob_version(response: &Response, request_id: &str) -> Result { + response + .headers() + .get(&VERSION_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()) + .ok_or_else(|| Error::Http { + source: "blob size probe returned a missing or invalid x-lancedb-version header".into(), + request_id: request_id.to_string(), + status_code: Some(response.status()), + }) +} + /// Probe one blob's size. /// /// `204` represents null. `416` with `bytes */0` represents an empty blob. async fn probe_blob_file(requester: Arc) -> Result> { let (request_id, response) = requester - .request_range("bytes=0-0", RangeRequestMode::SizeProbe) + .request_range("bytes=0-0", RangeRequestMode::SizeProbe, None) .await?; match response.status() { StatusCode::NO_CONTENT => { @@ -617,7 +637,8 @@ async fn probe_blob_file(requester: Arc) -> Result { let size = response @@ -633,6 +654,7 @@ async fn probe_blob_file(requester: Arc) -> Result) -> Result Err(Error::Http { source: format!("blob size probe expected HTTP 206 Partial Content, got {status}") @@ -658,7 +682,7 @@ async fn probe_blob_file(requester: Arc) -> Result, ) -> Result<(String, Response)> { assert_eq!(range_header, "bytes=0-0"); let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; @@ -1119,6 +1233,7 @@ mod tests { header::CONTENT_RANGE, format!("bytes 0-0/{}", 100 + self.index), ) + .header(VERSION_HEADER, "5") .body(vec![0u8]) .unwrap(); Ok((format!("probe-{}", self.index), Response::from(response))) @@ -1237,6 +1352,7 @@ mod tests { header::CONTENT_RANGE, format!("bytes 0-0/{}", PAYLOAD.len()), ) + .header(VERSION_HEADER, "5") .body(vec![PAYLOAD[0]]) .unwrap(); } @@ -1283,6 +1399,7 @@ mod tests { &self, range_header: &str, _mode: RangeRequestMode, + _version: Option, ) -> Result<(String, Response)> { let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; self.max_in_flight.fetch_max(now, Ordering::SeqCst); @@ -1315,7 +1432,7 @@ mod tests { in_flight, max_in_flight: max_in_flight.clone(), }); - let file = RemoteBlobFile::new(requester, PAYLOAD.len() as u64); + let file = RemoteBlobFile::new(requester, PAYLOAD.len() as u64, Some(5)); let ranges: Vec<_> = (0..16u64).map(|start| start..start + 2).collect(); let output = file.read_ranges(&ranges).await.unwrap(); @@ -1372,6 +1489,7 @@ mod tests { header::CONTENT_RANGE, format!("bytes 0-0/{}", PAYLOAD.len()), ) + .header(VERSION_HEADER, "5") .body(vec![PAYLOAD[0]]) .unwrap(); } @@ -1402,6 +1520,7 @@ mod tests { &self, range_header: &str, _mode: RangeRequestMode, + _version: Option, ) -> Result<(String, Response)> { if range_header.ends_with('-') { self.started.notify_one(); @@ -1427,7 +1546,11 @@ mod tests { release: release.clone(), started: started.clone(), }); - let file = Arc::new(RemoteBlobFile::new(requester, PAYLOAD.len() as u64)); + let file = Arc::new(RemoteBlobFile::new( + requester, + PAYLOAD.len() as u64, + Some(5), + )); let reader = { let file = file.clone(); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index da419938b..83a1eab1f 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3794,7 +3794,20 @@ impl BaseTable for NativeTable { let Some(segment) = segments.first() else { continue; }; - let params = load_segment_params(&dataset, segment).await?; + // The listing itself only needs the manifest. Missing index files must + // not hide every other index, or callers cannot find the one to repair. + let params = match load_segment_params(&dataset, segment).await { + Ok(params) => params, + Err(err) => { + log::warn!( + "Failed to read full text search configuration for index '{}': {}", + index.name, + err + ); + index.index_details = None; + continue; + } + }; let details = serde_json::to_string(¶ms).map_err(|source| Error::Other { message: format!( "Failed to serialize full text search configuration for index '{}'", diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 0a0d00da1..c5d6e072f 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -2078,6 +2078,50 @@ mod tests { assert_eq!(index.columns, vec!["tags".to_string()]); } + #[tokio::test] + async fn test_list_indices_with_missing_fts_index_files() { + let tmp_dir = tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + let conn = connect(uri).execute().await.unwrap(); + let batch = record_batch!( + ("id", Int32, [1, 2, 3]), + ("text", Utf8, ["alpha", "beta", "gamma"]) + ) + .unwrap(); + let table = conn.create_table("t", batch).execute().await.unwrap(); + table + .create_index(&["text"], Index::FTS(FtsIndexBuilder::default())) + .execute() + .await + .unwrap(); + table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .execute() + .await + .unwrap(); + + let fts_uuid = table + .list_indices() + .await + .unwrap() + .into_iter() + .find(|index| index.name == "text_idx") + .and_then(|index| index.index_uuid) + .unwrap(); + std::fs::remove_dir_all(tmp_dir.path().join("t.lance/_indices").join(fts_uuid)).unwrap(); + + // A new connection, so the index cache cannot serve the deleted files. + let conn = connect(uri).execute().await.unwrap(); + let table = conn.open_table("t").execute().await.unwrap(); + let mut indices = table.list_indices().await.unwrap(); + indices.sort_by(|a, b| a.name.cmp(&b.name)); + let names = indices.iter().map(|i| i.name.as_str()).collect::>(); + assert_eq!(names, vec!["id_idx", "text_idx"]); + assert_eq!(indices[1].index_type, crate::index::IndexType::FTS); + assert_eq!(indices[1].index_details, None); + assert!(indices[0].index_details.is_some()); + } + #[tokio::test] async fn test_create_inverted_index() { let conn = connect("memory://").execute().await.unwrap(); diff --git a/rust/lancedb/src/table/datafusion/udtf.rs b/rust/lancedb/src/table/datafusion/udtf.rs index a7a72f800..53eb55fcc 100644 --- a/rust/lancedb/src/table/datafusion/udtf.rs +++ b/rust/lancedb/src/table/datafusion/udtf.rs @@ -6,3 +6,5 @@ pub mod fts; pub mod duplicate_pairs; + +pub mod vector_query; diff --git a/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs b/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs index 47a35cbd2..20de93148 100644 --- a/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs +++ b/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs @@ -330,6 +330,96 @@ pub async fn plan_duplicate_pairs( Ok(tasks) } +/// Optional query columns. The default preserves the primitive's three-column +/// schema; SQL convenience calls include physical scope identity and may map +/// endpoints to an application key without fetching source vectors. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct DuplicatePairsOutput { + pub scope: bool, + pub id_column: Option, +} + +impl DuplicatePairsOutput { + pub fn schema(&self, dataset: &Dataset) -> Result { + let mut fields = duplicate_pair_schema().fields().to_vec(); + if self.scope { + fields.extend([ + Arc::new(Field::new("segment_id", DataType::Utf8, false)), + Arc::new(Field::new("partition_id", DataType::UInt64, false)), + Arc::new(Field::new("dataset_version", DataType::UInt64, false)), + ]); + } + if let Some(column) = &self.id_column { + let schema = Schema::from(dataset.schema()); + let field = schema.field_with_name(column)?; + if column.contains('.') + || matches!( + field.data_type(), + DataType::FixedSizeList(..) + | DataType::List(..) + | DataType::LargeList(..) + | DataType::Struct(..) + | DataType::Map(..) + ) + { + return plan_err!("id_column must be a top-level scalar column"); + } + fields.push(Arc::new(Field::new( + "id_a", + field.data_type().clone(), + field.is_nullable(), + ))); + fields.push(Arc::new(Field::new( + "id_b", + field.data_type().clone(), + field.is_nullable(), + ))); + } + Ok(Arc::new(Schema::new(fields))) + } + + async fn enrich( + &self, + batch: arrow_array::RecordBatch, + dataset: &Dataset, + task: &DuplicatePairTask, + ) -> Result { + use arrow_array::{StringArray, UInt64Array}; + let mut columns = batch.columns().to_vec(); + let n = batch.num_rows(); + if self.scope { + columns.extend([ + Arc::new(StringArray::from(vec![task.segment_id.to_string(); n])) + as arrow_array::ArrayRef, + Arc::new(UInt64Array::from(vec![task.partition_id as u64; n])), + Arc::new(UInt64Array::from(vec![dataset.version().version; n])), + ]); + } + if let Some(column) = &self.id_column { + for endpoint in 0..2 { + let ids = batch + .column(endpoint) + .as_any() + .downcast_ref::() + .expect("native row IDs"); + let projection = lance::dataset::ProjectionRequest::from_columns( + [column.as_str()], + dataset.schema(), + ); + let values = dataset.take_rows(ids.values(), projection).await?; + if values.num_rows() != n { + return plan_err!("pair row ID is missing from the pinned snapshot"); + } + columns.push(values.column(0).clone()); + } + } + Ok(arrow_array::RecordBatch::try_new( + self.schema(dataset)?, + columns, + )?) + } +} + /// Streaming native operator over partition descriptors, never source-vector /// rows. Assigned descriptors are consumed sequentially; distributed hosts send /// one descriptor per worker request. No Python or per-vector ANN calls. @@ -341,6 +431,7 @@ pub struct DuplicatePairsExec { tasks: Vec, properties: Arc, metrics: ExecutionPlanMetricsSet, + output: DuplicatePairsOutput, } impl DuplicatePairsExec { @@ -374,8 +465,23 @@ impl DuplicatePairsExec { tasks, properties, metrics: ExecutionPlanMetricsSet::new(), + output: DuplicatePairsOutput::default(), }) } + /// Configure optional output fields before dispatching this plan. + pub fn with_output(mut self, output: DuplicatePairsOutput) -> Result { + self.properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output.schema(&self.dataset)?), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + self.output = output; + Ok(self) + } + pub fn output(&self) -> &DuplicatePairsOutput { + &self.output + } pub fn dataset(&self) -> &Arc { &self.dataset } @@ -430,6 +536,7 @@ impl ExecutionPlan for DuplicatePairsExec { if partition != 0 { return plan_err!("duplicate-pair output partition out of bounds"); } + let output = self.output.clone(); let tasks = self.tasks.clone(); let dataset = self.dataset.clone(); let config = self.config.clone(); @@ -467,6 +574,7 @@ impl ExecutionPlan for DuplicatePairsExec { // into one-task PE requests before execution. let stream = stream::iter(tasks) .then(move |task| { + let output = output.clone(); let dataset = dataset.clone(); let config = config.clone(); let reservation = reservation.clone(); @@ -474,7 +582,7 @@ impl ExecutionPlan for DuplicatePairsExec { let completed = completed.clone(); async move { let reader = lance::index::vector::dedup::find_duplicate_pairs_in_partition_with_options( - dataset, + dataset.clone(), &config.column, task.segment_id, task.partition_id, @@ -506,15 +614,19 @@ impl ExecutionPlan for DuplicatePairsExec { } }, ); - Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new( - duplicate_pair_schema(), - measured, - )) as SendableRecordBatchStream) + let schema = output.schema(&dataset)?; + let measured = measured.and_then(move |batch| { + let output = output.clone(); + let dataset = dataset.clone(); + let task = task.clone(); + async move { output.enrich(batch, &dataset, &task).await } + }); + Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new(schema, measured)) as SendableRecordBatchStream) } }) .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new( - duplicate_pair_schema(), + self.schema(), stream, ))) } diff --git a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs new file mode 100644 index 000000000..0d4f23d92 --- /dev/null +++ b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs @@ -0,0 +1,352 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Composable index-backed pair and dedup relations. Hosts bind a snapshot once +//! during planning, distribute the pair child, and use their ordinary SQL sink. + +use std::sync::Arc; + +use arrow_schema::SchemaRef; +use async_trait::async_trait; +use datafusion_catalog::{Session, TableProvider}; +use datafusion_common::{Result, plan_err}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::{Expr, TableType}; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + coalesce_partitions::CoalescePartitionsExec, + empty::EmptyExec, + execution_plan::{Boundedness, EmissionType}, + projection::ProjectionExec, + stream::RecordBatchStreamAdapter, +}; +use futures::{TryStreamExt, stream}; +use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap}; +use lance::{Dataset, index::DatasetIndexInternalExt}; +use lance_core::datatypes::BlobHandling; +use lance_index::metrics::NoOpMetricsCollector; + +use super::duplicate_pairs::{ + DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsOutput, + plan_duplicate_pairs, +}; +use crate::table::{NativeTableExt, datafusion::BaseTableAdapter}; + +const REPRESENTATIVE_SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024; +const MATERIALIZE_BATCH_ROWS: usize = 1024; + +/// Select nonempty physical IVF scopes before dispatch, never by filtering the +/// resulting pairs. A count spans all segments; partition IDs are segment-local. +#[derive(Debug, Clone, Default)] +pub enum PartitionSelection { + #[default] + All, + First { + count: usize, + }, + Random { + count: usize, + seed: u64, + }, +} + +/// Query-only options. Snapshot binding is separate so multiple calls in one +/// statement can share the same immutable dataset. +#[derive(Debug, Clone)] +pub struct VectorQueryOptions { + pub column: String, + pub distance_threshold: f32, + pub selection: PartitionSelection, + pub output: DuplicatePairsOutput, + pub dedup: bool, +} + +impl BaseTableAdapter { + /// Open a retained version, or refresh and capture the latest manifest once. + /// No worker is allowed to independently resolve "latest". + pub async fn vector_query_snapshot(&self, version: Option) -> Result> { + let native = self.table.as_native().ok_or_else(|| { + datafusion_common::DataFusionError::Plan( + "vector queries require a native indexed table".into(), + ) + })?; + if version == Some(0) { + return plan_err!("dataset_version must be positive"); + } + let current = native + .dataset + .get() + .await + .map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?; + Ok(match version { + Some(version) => Arc::new(current.checkout_version(version).await?), + None => { + // Do not change a shared table handle's pinned version. + let mut latest = (*current).clone(); + latest.checkout_latest().await?; + Arc::new(latest) + } + }) + } +} + +/// A table provider bound to one source snapshot. The host retains the user's +/// unbound SQL for subsequent refreshes; this provider lives only for this run. +#[derive(Debug)] +pub struct VectorQueryTable { + dataset: Arc, + config: DuplicatePairsConfig, + tasks: Vec, + options: VectorQueryOptions, + schema: SchemaRef, +} + +impl VectorQueryTable { + pub async fn try_new(dataset: Arc, options: VectorQueryOptions) -> Result { + if options.dedup + && (!matches!(options.selection, PartitionSelection::All) + || options.output.id_column.is_some()) + { + return plan_err!( + "vector_dedup requires all partitions and returns original rows; sampling and id_column are pair-query options" + ); + } + let config = DuplicatePairsConfig { + dataset_version: dataset.version().version, + column: options.column.clone(), + distance_threshold: options.distance_threshold, + }; + let mut tasks = plan_duplicate_pairs(dataset.clone(), &config).await?; + // Canonical scope order is independent of metadata enumeration order. + tasks.sort_by_key(|task| (task.segment_id, task.partition_id)); + let mut nonempty = Vec::new(); + for group in tasks.chunk_by(|a, b| a.segment_id == b.segment_id) { + let index = dataset + .open_vector_index(&config.column, &group[0].segment_id, &NoOpMetricsCollector) + .await?; + nonempty.extend( + group + .iter() + .filter(|task| index.partition_size(task.partition_id) > 0) + .cloned(), + ); + } + match options.selection { + PartitionSelection::All => {} + PartitionSelection::First { count } | PartitionSelection::Random { count, .. } + if count == 0 => + { + return plan_err!("partition count must be positive"); + } + PartitionSelection::First { count } => nonempty.truncate(count), + PartitionSelection::Random { count, seed } => { + // A single segment matches ivf_partition's SplitMix64 sampling. + // With multiple segments the canonical scope ordinal extends the + // same policy without conflating equal partition IDs. + let segments = tasks + .iter() + .map(|t| t.segment_id) + .collect::>() + .len(); + let mut ranked = nonempty + .into_iter() + .enumerate() + .map(|(ordinal, task)| { + let key = if segments == 1 { + task.partition_id as u64 + } else { + ordinal as u64 + }; + (sample_rank(seed, key), task) + }) + .collect::>(); + ranked.sort_by_key(|(rank, task)| (*rank, task.segment_id, task.partition_id)); + ranked.truncate(count); + nonempty = ranked.into_iter().map(|(_, task)| task).collect(); + nonempty.sort_by_key(|task| (task.segment_id, task.partition_id)); + } + } + let schema = if options.dedup { + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::AllBinary); + scanner.schema().await? + } else { + options.output.schema(&dataset)? + }; + Ok(Self { + dataset, + config, + tasks: nonempty, + options, + schema, + }) + } +} + +fn sample_rank(seed: u64, ordinal: u64) -> u64 { + let mut value = seed.wrapping_add(ordinal).wrapping_add(0x9E3779B97F4A7C15); + value = (value ^ (value >> 30)).wrapping_mul(0xBF58476D1CE4E5B9); + value = (value ^ (value >> 27)).wrapping_mul(0x94D049BB133111EB); + value ^ (value >> 31) +} + +#[async_trait] +impl TableProvider for VectorQueryTable { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn table_type(&self) -> TableType { + TableType::View + } + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + let output = if self.options.dedup { + DuplicatePairsOutput::default() + } else { + self.options.output.clone() + }; + let pairs: Arc = if self.tasks.is_empty() { + Arc::new(EmptyExec::new(output.schema(&self.dataset)?)) + } else { + Arc::new( + DuplicatePairsExec::try_new( + self.dataset.clone(), + self.config.clone(), + self.tasks.clone(), + )? + .with_output(output)?, + ) + }; + let plan: Arc = if self.options.dedup { + Arc::new(VectorDedupExec::new( + self.dataset.clone(), + pairs, + self.schema.clone(), + )) + } else { + pairs + }; + if let Some(indices) = projection { + let expressions: Vec<_> = indices + .iter() + .map(|&i| { + ( + Arc::new(datafusion_physical_expr::expressions::Column::new( + self.schema.field(i).name(), + i, + )) + as Arc, + self.schema.field(i).name().clone(), + ) + }) + .collect(); + Ok(Arc::new(ProjectionExec::try_new(expressions, plan)?)) + } else { + Ok(plan) + } + } +} + +/// Greedy direct representatives over the complete pair relation. Sorting spills +/// to disk; only vertex masks are retained. This has no durable checkpoint and +/// publishes nothing: failure is handled by the enclosing query/MV transaction. +#[derive(Debug)] +pub struct VectorDedupExec { + dataset: Arc, + pairs: Arc, + properties: Arc, +} + +impl VectorDedupExec { + fn new(dataset: Arc, pairs: Arc, schema: SchemaRef) -> Self { + Self { + dataset, + pairs, + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )), + } + } +} + +impl DisplayAs for VectorDedupExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "VectorDedupExec: source_version={}, policy=direct_representative", + self.dataset.version().version + ) + } +} + +impl ExecutionPlan for VectorDedupExec { + fn name(&self) -> &str { + "VectorDedupExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.pairs] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return plan_err!("VectorDedupExec requires one pair input"); + } + Ok(Arc::new(Self::new( + self.dataset.clone(), + children[0].clone(), + self.schema(), + ))) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + if partition != 0 { + return plan_err!("invalid dedup output partition"); + } + let dataset = self.dataset.clone(); + let pairs = CoalescePartitionsExec::new(self.pairs.clone()).execute(0, context)?; + let stream = stream::once(async move { + let (removed, _) = crate::materialized_view::vector_dedup::select_representatives( + pairs, + REPRESENTATIVE_SORT_MEMORY_BYTES, + ) + .await + .map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?; + let mut scanner = dataset.scan(); + scanner + .with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter( + removed, + ))) + .blob_handling(BlobHandling::AllBinary) + .batch_size(MATERIALIZE_BATCH_ROWS) + .strict_batch_size(true); + Ok::<_, datafusion_common::DataFusionError>( + scanner + .try_into_stream() + .await? + .map_err(datafusion_common::DataFusionError::from), + ) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} diff --git a/rust/lancedb/tests/duplicate_pairs.rs b/rust/lancedb/tests/duplicate_pairs.rs index 910088bc8..f7d7ed3a2 100644 --- a/rust/lancedb/tests/duplicate_pairs.rs +++ b/rust/lancedb/tests/duplicate_pairs.rs @@ -76,7 +76,7 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result .to_str() .unwrap() .to_string(); - let n = 1025; + let n = 8193; let vectors = (0..n * 2) .flat_map(|i| { let i = i % n; @@ -196,6 +196,89 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result ); let actual = ctx.sql(&sql).await?.collect().await?; assert_eq!(pairs(&actual), pairs(&expected)); + use lancedb::table::datafusion::udtf::{ + duplicate_pairs::DuplicatePairsOutput, + vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable}, + }; + let query_options = VectorQueryOptions { + column: "vector".into(), + distance_threshold: threshold, + selection: PartitionSelection::All, + output: DuplicatePairsOutput { + scope: true, + id_column: Some("id".into()), + }, + dedup: false, + }; + let provider = VectorQueryTable::try_new(ds.clone(), query_options.clone()).await?; + let extended = ctx.read_table(Arc::new(provider))?.collect().await?; + assert_eq!(pairs(&extended), pairs(&expected)); + for batch in &extended { + let rowids = batch + .column_by_name("row_id_a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let ids = batch + .column_by_name("id_a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + let expected_id = (rowids.value(i) >> 32) as i64 * n as i64 + + (rowids.value(i) & 0xffffffff) as i64; + assert_eq!(ids.value(i), expected_id); + } + } + let mut sampled_options = query_options.clone(); + sampled_options.selection = PartitionSelection::Random { count: 1, seed: 42 }; + let sample = ctx + .read_table(Arc::new( + VectorQueryTable::try_new(ds.clone(), sampled_options.clone()).await?, + ))? + .collect() + .await?; + let replay = ctx + .read_table(Arc::new( + VectorQueryTable::try_new(ds.clone(), sampled_options).await?, + ))? + .collect() + .await?; + assert_eq!(sample, replay); + let scopes = sample + .iter() + .flat_map(|batch| { + let segments = batch + .column_by_name("segment_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let partitions = batch + .column_by_name("partition_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|i| (segments.value(i).to_string(), partitions.value(i))) + .collect::>() + }) + .collect::>(); + assert!(scopes.len() <= 1); + let mut dedup_options = query_options; + dedup_options.dedup = true; + dedup_options.output = DuplicatePairsOutput::default(); + let clean_query = ctx + .read_table(Arc::new( + VectorQueryTable::try_new(ds.clone(), dedup_options).await?, + ))? + .select(vec![datafusion::prelude::col("id")])? + .collect() + .await?; + assert_eq!( pairs(&actual).len(), match threshold { @@ -388,6 +471,21 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result .collect::>(); actual_ids.sort_unstable(); expected_ids.sort_unstable(); + let mut query_ids = clean_query + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + query_ids.sort_unstable(); + assert_eq!(query_ids, expected_ids); + assert_eq!(actual_ids, expected_ids); if threshold == 1.0 { assert!(actual_ids.contains(&0) && actual_ids.contains(&2)); @@ -510,6 +608,90 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() ) .await?; let version = ds.version().version; + use lancedb::table::datafusion::udtf::{ + duplicate_pairs::DuplicatePairsOutput, + vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable}, + }; + let options = VectorQueryOptions { + column: "vector".into(), + distance_threshold: 1.0, + selection: PartitionSelection::All, + output: DuplicatePairsOutput::default(), + dedup: true, + }; + let mapped = VectorQueryTable::try_new( + Arc::new(ds.clone()), + VectorQueryOptions { + dedup: false, + output: DuplicatePairsOutput { + scope: true, + id_column: Some("id".into()), + }, + ..options.clone() + }, + ) + .await?; + let mapped = SessionContext::new() + .read_table(Arc::new(mapped))? + .collect() + .await?; + let mut mapped_edges = Vec::new(); + for batch in mapped { + let a = batch + .column_by_name("id_a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let b = batch + .column_by_name("id_b") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + mapped_edges.push((a.value(i).min(b.value(i)), a.value(i).max(b.value(i)))); + } + } + mapped_edges.sort_unstable(); + assert_eq!(mapped_edges, [(0, 1), (1, 2)]); + let provider = VectorQueryTable::try_new(Arc::new(ds.clone()), options).await?; + let query_rows = SessionContext::new() + .read_table(Arc::new(provider))? + .collect() + .await?; + let mut query_ids = query_rows + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + query_ids.sort_unstable(); + assert_eq!(query_ids, [0, 2, 3, 4]); + for batch in &query_rows { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let images = batch + .column_by_name("image") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + assert_eq!(images.value(i), vec![ids.value(i) as u8; 4096]); + } + } let definition = MaterializedViewDefinition::from_sql(&format!( "SELECT * FROM vector_dedup('source', {version}, 'vector', 1)" ))?; @@ -519,6 +701,13 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() .await?; // A later source delete does not change the declared result snapshot. table.delete("id = 0").await?; + let adapter = BaseTableAdapter::try_new(table.base_table().clone()).await?; + let explicit = adapter.vector_query_snapshot(Some(version)).await?; + let latest = adapter.vector_query_snapshot(None).await?; + assert_eq!(explicit.version().version, version); + assert!(latest.version().version > version); + assert_eq!(explicit.count_rows(None).await?, 5); + assert_eq!(latest.count_rows(None).await?, 4); clean.refresh().execute().await?; assert_eq!(clean.table().count_rows(None).await?, 4); let output = Dataset::open(&clean.table().uri().await?).await?;