From 4a03859b2a95f38dbef8fab962570de93d70eded Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:44:00 -0700 Subject: [PATCH] fix!: normalize blob layouts and support external byte ranges (#4589) --- docs/src/js/functions/blob.md | 9 + docs/src/js/type-aliases/BlobInput.md | 4 + nodejs/__test__/blob.test.ts | 188 ++- nodejs/__test__/remote.test.ts | 135 ++ nodejs/__test__/table.test.ts | 187 +++ nodejs/lancedb/arrow.ts | 34 +- nodejs/lancedb/blob.ts | 99 +- python/python/lancedb/table.py | 36 +- python/python/tests/test_blob.py | 146 +++ python/python/tests/test_remote_db.py | 83 ++ rust/lancedb/src/blob.rs | 63 +- rust/lancedb/src/remote/table.rs | 302 +++++ rust/lancedb/src/table/add_data.rs | 112 +- .../src/table/datafusion/blob_coerce.rs | 1151 ++++++++++++++++- rust/lancedb/src/table/datafusion/cast.rs | 61 +- rust/lancedb/src/table/merge.rs | 71 +- rust/lancedb/src/table/merge/lsm.rs | 9 +- rust/lancedb/src/table/refresh.rs | 6 +- rust/lancedb/tests/blob_integration.rs | 270 +++- rust/lancedb/tests/duplicate_pairs.rs | 2 + 20 files changed, 2837 insertions(+), 131 deletions(-) diff --git a/docs/src/js/functions/blob.md b/docs/src/js/functions/blob.md index 1a65b5133..d56735aca 100644 --- a/docs/src/js/functions/blob.md +++ b/docs/src/js/functions/blob.md @@ -11,6 +11,15 @@ function blob(name, options): Field ``` Declares a `lance.blob.v2` column. +Uses nullable `data`, `uri`, `position`, and `size` fields for inline bytes, +whole external files, and external byte ranges. Set both `position` and +positive `size` for a range, or omit both for a whole file. + +Existing two-field columns remain writable. Appends and merge inserts fill +missing range fields with nulls for four-field targets and drop all-null +range fields for two-field targets. Non-null ranges require a four-field +target. Overwrite with Arrow input replaces the table's schema; row objects +are converted using the current table schema. Query results are descriptors, not payload bytes. Use [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles) to read bytes. diff --git a/docs/src/js/type-aliases/BlobInput.md b/docs/src/js/type-aliases/BlobInput.md index d1738a760..53a47e5fd 100644 --- a/docs/src/js/type-aliases/BlobInput.md +++ b/docs/src/js/type-aliases/BlobInput.md @@ -12,6 +12,7 @@ type BlobInput: | BlobUri | object | object + | object | null | undefined; ``` @@ -20,6 +21,9 @@ A value accepted for a `lance.blob.v2` column (see [blob](../functions/blob.md)) Either inline bytes, a URI pointing at external bytes, or a struct that sets exactly one of `data` or `uri`. `null` and `undefined` write a null blob. +External byte ranges require both +`position` (non-negative) and `size` (positive). Use safe integer numbers or +`bigint` values for these unsigned 64-bit byte counts. ## Example diff --git a/nodejs/__test__/blob.test.ts b/nodejs/__test__/blob.test.ts index b41d5716f..12ff56b83 100644 --- a/nodejs/__test__/blob.test.ts +++ b/nodejs/__test__/blob.test.ts @@ -21,7 +21,41 @@ import { BlobFile, blob, coerceBlobValue, isBlobField } from "../lancedb/blob"; import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding"; import { getRegistry } from "../lancedb/embedding/registry"; +function legacyBlob(name: string): Field { + const field = blob(name); + return new Field( + name, + new Struct((field.type as Struct).children.slice(0, 2)), + field.nullable, + field.metadata, + ); +} + describe("blob()", () => { + it("declares the four-field range layout", () => { + const field = blob("image", { + nullable: false, + inlineSizeThreshold: 12, + }); + expect(field.nullable).toBe(false); + expect(isBlobField(field)).toBe(true); + expect( + field.metadata.get("lance-encoding:blob-inline-size-threshold"), + ).toBe("12"); + expect( + (field.type as Struct).children.map((f) => [ + f.name, + f.type.toString(), + f.nullable, + ]), + ).toEqual([ + ["data", "LargeBinary", true], + ["uri", "Utf8", true], + ["position", "Uint64", true], + ["size", "Uint64", true], + ]); + }); + it("marks the field as lance.blob.v2", () => { const field = blob("image", { nullable: false }); expect(field.nullable).toBe(false); @@ -156,7 +190,7 @@ describe("coerceBlobValue", () => { [ "unknown struct field", { data: Buffer.from("y"), mime: "image/png" }, - /Blob struct values only support 'data' and 'uri', got 'mime'/, + /Blob struct values only support .*got 'mime'/, ], [ "misspelled struct field", @@ -174,6 +208,158 @@ describe("coerceBlobValue", () => { }); }); +describe("external blob ranges", () => { + const uri = "s3://bucket/key"; + + it.each([0, 4, 0n, 4n, 9007199254740993n])( + "preserves position %s", + (position) => { + expect(coerceBlobValue({ uri, position, size: 6 })).toEqual({ + data: null, + uri, + position: BigInt(position), + size: 6n, + }); + }, + ); + + it.each([ + [{ uri, position: 0 }, /both position and size/], + [{ uri, size: 6 }, /both position and size/], + [{ uri, position: null, size: 6 }, /both position and size/], + [{ uri, position: 0, size: 0 }, /size must be positive/], + [{ uri, position: -1, size: 6 }, /position must be non-negative/], + [{ uri, position: 0, size: -1n }, /size must be positive/], + [{ uri, position: 0.5, size: 6 }, /safe integer or bigint/], + [ + { uri, position: Number.MAX_SAFE_INTEGER + 1, size: 6 }, + /safe integer or bigint/, + ], + [{ uri, position: 0, size: Number.NaN }, /safe integer or bigint/], + [{ uri, position: "4", size: 6 }, /safe integer or bigint/], + [{ uri, position: false, size: 6 }, /safe integer or bigint/], + [{ uri, position: 1n << 64n, size: 6 }, /fit in uint64/], + [{ uri, position: 0xffffffffffffffffn, size: 1 }, /overflows uint64/], + [{ data: Buffer.from("x"), position: 0, size: 6 }, /require a uri/], + ])("rejects malformed input %p", (value, message) => { + expect(() => coerceBlobValue(value)).toThrow(message); + }); + + it("preserves ranges while normalizing a URL held by getters", () => { + class Range { + get uri() { + return new URL(uri); + } + get position() { + return 4; + } + get size() { + return 6; + } + } + expect(coerceBlobValue(new Range())).toEqual({ + data: null, + uri, + position: 4n, + size: 6n, + }); + }); + + it("writes ranges and mixed whole-file, inline, and null values to Arrow", () => { + const schema = new Schema([blob("image")]); + const rows = [ + null, + Buffer.from("x"), + uri, + { uri, position: 4, size: 6 }, + ].map((image) => ({ image })); + const column = makeArrowTable(rows, { schema }).getChild("image")!; + expect(column.get(0)).toBeNull(); + expect(Buffer.from(column.getChild("data")!.get(1))).toEqual( + Buffer.from("x"), + ); + expect(column.getChild("uri")!.get(2)).toBe(uri); + expect(column.getChild("position")!.toArray()).toEqual( + new BigUint64Array([0n, 0n, 0n, 4n]), + ); + expect(column.getChild("position")!.get(2)).toBeNull(); + expect(column.getChild("size")!.get(3)).toBe(6n); + }); + + it.each([ + Buffer.from("inline"), + uri, + { uri, position: null, size: null }, + null, + ])("retains the full declared range layout when seeded with %p", (image) => { + const schema = new Schema([blob("image")]); + const table = makeArrowTable([{ image }], { schema }); + expect(table.schema).toEqual(schema); + expect(table.getChild("image")!.getChild("position")!.get(0)).toBeNull(); + }); + + it.each([0, 4])( + "rejects position %s on a minimal schema before dropping keys", + (position) => { + expect(() => + makeArrowTable( + [ + { image: Buffer.from("inline") }, + { image: { uri, position, size: 6 } }, + ], + { schema: new Schema([legacyBlob("image")]) }, + ), + ).toThrow( + /field image at row 1: Blob storage schema cannot store byte ranges/, + ); + }, + ); + + it("accepts omitted or null range fields on the minimal schema", () => { + const column = makeArrowTable( + [{ image: { uri, position: null, size: null } }], + { + schema: new Schema([legacyBlob("image")]), + }, + ).getChild("image")!; + expect(column.getChild("uri")!.get(0)).toBe(uri); + }); + + it("validates ranges nested in structs and lists against their storage fields", () => { + const value = { uri, position: 4, size: 6 }; + for (const field of [legacyBlob("image"), blob("image")]) { + const schema = new Schema([ + new Field("info", new Struct([field]), true), + new Field("images", new List(field), true), + ]); + if ((field.type as Struct).children.length === 2) { + expect(() => + makeArrowTable([{ info: { image: value } }], { schema }), + ).toThrow(/field info.image.*storage schema/); + expect(() => makeArrowTable([{ images: [value] }], { schema })).toThrow( + /field images\[0\].*storage schema/, + ); + } else { + const table = makeArrowTable( + [{ info: { image: value }, images: [value] }], + { schema }, + ); + const info = table.getChild("info")!.get(0) as { + image: { position: bigint; size: bigint }; + }; + expect(info.image.position).toBe(4n); + const images = Array.from( + table.getChild("images")!.get(0) as Iterable<{ + position: bigint; + size: bigint; + }>, + ); + expect(images[0].size).toBe(6n); + } + } + }); +}); + describe("resolveBlobInputs", () => { it("reads Blob and File values in blob columns only", async () => { const schema = new Schema([ diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index c06cb737e..7e9479f94 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -12,6 +12,7 @@ import { OAuthConfig, OAuthFlowType, TlsConfig, + blob, connect, } from "../lancedb"; import { @@ -19,7 +20,10 @@ import { Field, Int64, RecordBatch, + RecordBatchReader, Schema, + Struct, + makeArrowTable, tableToIPC, } from "../lancedb/arrow"; import { @@ -94,6 +98,137 @@ async function withMockDatabase( } describe("remote connection", () => { + it.each([false, true])( + "normalizes remote blob writes (range target: %s)", + async (wide) => { + const field = (full: boolean) => { + const f = blob("image"); + return full + ? f + : new Field( + f.name, + new Struct((f.type as Struct).children.slice(0, 2)), + f.nullable, + f.metadata, + ); + }; + const schema = (full: boolean) => + new Schema([new Field("id", new Int64(), true), field(full)]); + const writes: { path: string; schema: Schema }[] = []; + let storedWide = wide; + const describe = () => ({ + version: writes.length + 1, + schema: { + fields: [ + { name: "id", type: { type: "int64" }, nullable: true }, + { + name: "image", + nullable: true, + metadata: { "ARROW:extension:name": "lance.blob.v2" }, + type: { + type: "struct", + fields: [ + { + name: "data", + nullable: true, + type: { type: "large_binary" }, + }, + { name: "uri", nullable: true, type: { type: "string" } }, + ...(storedWide + ? ["position", "size"].map((name) => ({ + name, + nullable: true, + type: { type: "uint64" }, + })) + : []), + ], + }, + }, + ], + }, + }); + await withMockDatabase( + (req, res) => { + if (req.url?.includes("/describe/")) { + res + .writeHead(200, { "content-type": "application/json" }) + .end(JSON.stringify(describe())); + return; + } + const chunks: Buffer[] = []; + // Failed streaming input aborts before end; only complete bodies commit. + req.on("error", () => { + /* Aborted inputs must not commit. */ + }); + req.on("data", (chunk) => chunks.push(chunk)); + req.on("end", () => { + const reader = RecordBatchReader.from(Buffer.concat(chunks)).open(); + // Only read the schema: Rust separately checks compressed wire values. + writes.push({ path: req.url ?? "", schema: reader.schema }); + storedWide = + (reader.schema.fields[1].type as Struct).children.length === 4; + res.writeHead(200, { "content-type": "application/json" }).end( + JSON.stringify({ + version: writes.length + 1, + // biome-ignore lint/style/useNamingConvention: server wire format + num_inserted_rows: 1, + }), + ); + }); + }, + async (db) => { + const table = await db.openTable("test"); + const input = makeArrowTable( + [{ id: 1n, image: Buffer.from("payload") }], + { schema: schema(!wide) }, + ); + await table.add(input); + await table + .mergeInsert("id") + .whenMatchedUpdateAll() + .whenNotMatchedInsertAll() + .execute(input); + expect(writes).toHaveLength(2); + for (const write of writes) + expect(write.schema).toEqual(schema(wide)); + expect(writes[0].path).toContain("/insert/"); + expect(writes[1].path).toContain("/merge_insert/"); + if (!wide) { + const invalid = makeArrowTable( + [ + { + id: 1n, + image: { + uri: "s3://registered/payload", + position: 4, + size: 6, + }, + }, + ], + { schema: schema(true) }, + ); + await expect(table.add(invalid)).rejects.toThrow( + /cannot store byte ranges/, + ); + await expect( + table + .mergeInsert("id") + .whenMatchedUpdateAll() + .whenNotMatchedInsertAll() + .execute(invalid), + ).rejects.toThrow(/cannot store byte ranges/); + expect(writes).toHaveLength(2); + } + await table.add(input, { mode: "overwrite" }); + expect(writes).toHaveLength(3); + expect(writes[2].path).toContain("mode=overwrite"); + expect(writes[2].schema).toEqual(schema(!wide)); + expect(await table.schema()).toEqual(schema(!wide)); + }, + ); + }, + ); + it("reports unsupported operations with backend-neutral errors", async () => { await withMockDatabase( (req, res) => { diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index b30121725..1c4592a14 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2583,6 +2583,16 @@ describe("when dealing with versioning", () => { }); }); +function legacyBlob(name: string): Field { + const field = blob(name); + return new Field( + name, + new Struct((field.type as Struct).children.slice(0, 2)), + field.nullable, + field.metadata, + ); +} + describe("when dealing with blob columns", () => { let tmpDir: tmp.DirResult; beforeEach(() => { @@ -2592,6 +2602,89 @@ describe("when dealing with blob columns", () => { tmpDir.removeCallback(); }); + it.each([false, true])( + "normalizes Arrow blob layouts for append and merge (range target: %s)", + async (wide) => { + const db = await connect(tmpDir.name); + const schema = (full: boolean) => + new Schema([ + new Field("id", new Int64(), true), + full ? blob("image") : legacyBlob("image"), + ]); + const table = await db.createTable( + "layouts", + [{ id: 1n, image: Buffer.from("old") }], + { schema: schema(wide) }, + ); + await table.add( + makeArrowTable([{ id: 2n, image: Buffer.from("appended") }], { + schema: schema(!wide), + }), + ); + await table + .mergeInsert("id") + .whenMatchedUpdateAll() + .whenNotMatchedInsertAll() + .execute( + makeArrowTable( + [ + { id: 1n, image: Buffer.from("updated") }, + { id: 3n, image: Buffer.from("inserted") }, + ], + { schema: schema(!wide) }, + ), + ); + expect(await table.schema()).toEqual(schema(wide)); + const rows = await table.query().withRowId().toArray(); + rows.sort((a, b) => Number(a.id - b.id)); + expect( + await table.fetchBlobs( + "image", + rows.map((row) => row._rowid as bigint), + ), + ).toEqual([ + Buffer.from("updated"), + Buffer.from("appended"), + Buffer.from("inserted"), + ]); + if (!wide) { + const bad = makeArrowTable( + [ + { + id: 1n, + image: { uri: "s3://registered/file", position: 0, size: 6 }, + }, + ], + { schema: schema(true) }, + ); + const version = await table.version(); + await expect(table.add(bad)).rejects.toThrow( + /cannot store byte ranges/, + ); + await expect( + table + .mergeInsert("id") + .whenMatchedUpdateAll() + .whenNotMatchedInsertAll() + .execute(bad), + ).rejects.toThrow(/cannot store byte ranges/); + expect(await table.version()).toBe(version); + expect(await table.countRows()).toBe(3); + } + const replacement = makeArrowTable( + [{ id: 4n, image: Buffer.from("replacement") }], + { schema: schema(!wide) }, + ); + await table.add(replacement, { mode: "overwrite" }); + expect(await table.schema()).toEqual(schema(!wide)); + expect(await table.countRows()).toBe(1); + const [row] = await table.query().withRowId().toArray(); + expect(await table.fetchBlobs("image", [row._rowid as bigint])).toEqual([ + Buffer.from("replacement"), + ]); + }, + ); + it("rejects unsupported blob values in createTable and add", async () => { const db = await connect(tmpDir.name); const schema = new Schema([ @@ -2668,6 +2761,100 @@ describe("when dealing with blob columns", () => { expect(bytes[0]).toEqual(payload); }); + it.each([0, 4])( + "fetches the exact external range at position %s after reopening", + async (position) => { + const { db, uri, payload } = await openExternalBlobTable(); + const schema = new Schema([ + new Field("id", new Int64(), true), + blob("payload"), + ]); + const rows: { id: bigint; payload: BlobInput }[] = [ + { id: 1n, payload: Buffer.from("inline") }, + { id: 2n, payload: { uri: new URL(uri), position, size: 6 } }, + { id: 3n, payload: uri }, + { id: 4n, payload: null }, + ]; + const table = await db.createTable("ranges", rows.slice(0, 1), { + schema, + }); + await table.add(rows.slice(1), { allowExternalBlobOutsideBases: true }); + table.close(); + const reopened = await db.openTable("ranges"); + const results = await reopened.query().withRowId().toArray(); + results.sort((a, b) => Number(a.id - b.id)); + expect(results[1].payload.position).toBe(BigInt(position)); + expect(results[1].payload.size).toBe(6n); + expect( + await reopened.fetchBlobs( + "payload", + results.map((r) => r._rowid as bigint), + ), + ).toEqual([ + Buffer.from("inline"), + payload.subarray(position, position + 6), + payload, + null, + ]); + }, + ); + + it.each([0, 4])( + "rejects ranges at position %s on reopened minimal tables without writing", + async (position) => { + const { db, uri, payload } = await openExternalBlobTable(); + const table = await db.createEmptyTable( + "legacy_blobs", + new Schema([new Field("id", new Int64(), true), legacyBlob("payload")]), + ); + table.close(); + const reopened = await db.openTable("legacy_blobs"); + await expect( + reopened.add( + [ + { id: 1n, payload: Buffer.from("inline") }, + { id: 2n, payload: { uri, position, size: 6 } }, + ], + { allowExternalBlobOutsideBases: true }, + ), + ).rejects.toThrow(/storage schema cannot store byte ranges/); + expect(await reopened.countRows()).toBe(0); + await reopened.add( + [ + { id: 3n, payload: uri }, + { id: 4n, payload: Buffer.from("inline") }, + ], + { allowExternalBlobOutsideBases: true }, + ); + const rows = await reopened.query().withRowId().toArray(); + rows.sort((a, b) => Number(a.id - b.id)); + expect( + await reopened.fetchBlobs( + "payload", + rows.map((r) => r._rowid as bigint), + ), + ).toEqual([payload, Buffer.from("inline")]); + }, + ); + + it.each([{ position: 0 }, { size: 6 }, { position: 0, size: 0 }])( + "rejects partial or empty ranges without writing: %p", + async (range) => { + const { db, uri } = await openExternalBlobTable(); + const table = await db.createEmptyTable( + "invalid_ranges", + new Schema([blob("payload")]), + ); + await expect( + table.add( + [{ payload: Buffer.from("inline") }, { payload: { uri, ...range } }], + { allowExternalBlobOutsideBases: true }, + ), + ).rejects.toThrow(/both position and size|size must be positive/); + expect(await table.countRows()).toBe(0); + }, + ); + it("applies the external blob opt-in only to the current add", async () => { const { table, uri } = await openExternalBlobTable(); diff --git a/nodejs/lancedb/arrow.ts b/nodejs/lancedb/arrow.ts index 3130da8eb..68b35e256 100644 --- a/nodejs/lancedb/arrow.ts +++ b/nodejs/lancedb/arrow.ts @@ -456,14 +456,29 @@ export function makeArrowTable( if (schema !== undefined) { // Validate blob values up front and give every one the full - // `{ data, uri }` shape, so inference sees the same struct whether a row + // storage shape, so inference sees the same struct whether a row // passed bytes, a URI, or a partial struct. - data = mapBlobInputs(data, schema, (value, field, row) => { + data = mapBlobInputs(data, schema, (value, field, row, storageField) => { if (value === undefined) { return value; } try { - return coerceBlobValue(value); + const normalized = coerceBlobValue(value, storageField); + if ( + normalized === null || + !(storageField.type instanceof Struct) || + storageField.type.children.length === 2 + ) { + return normalized; + } + // Keep even unused range children during schema inference. Otherwise + // creating a range-capable table with inline data would lose them. + return Object.fromEntries( + storageField.type.children.map((child) => [ + child.name, + normalized[child.name as keyof typeof normalized] ?? null, + ]), + ); } catch (e) { throw new Error( `Invalid value for blob field ${field} at row ${row}: ${(e as Error).message}`, @@ -533,7 +548,12 @@ function containsBlobField(field: Field): boolean { ); } -type BlobInputVisitor = (value: unknown, field: string, row: number) => unknown; +type BlobInputVisitor = ( + value: unknown, + field: string, + row: number, + storageField: Field, +) => unknown; /** * Calls `visit` on every value that lands in a blob column of `schema`, @@ -592,7 +612,7 @@ function mapBlobField( visit: BlobInputVisitor, ): unknown { if (isBlobField(field)) { - return visit(value, label, row); + return visit(value, label, row, field); } if (field.type instanceof Struct && isObject(value)) { return mapBlobFields(value, field.type.children, `${label}.`, row, visit); @@ -697,11 +717,11 @@ function transposeData( const valuesPath = [...path, field.name]; if (isBlobField(field) && field.type instanceof Struct) { const blobRows = data.map((datum) => - coerceBlobValue(valueAtPath(datum, valuesPath)), + coerceBlobValue(valueAtPath(datum, valuesPath), field), ); const childVectors = field.type.children.map((child) => { const values = blobRows.map((row) => - row == null ? null : (row[child.name as "data" | "uri"] ?? null), + row == null ? null : (row[child.name as keyof typeof row] ?? null), ); return makeVector(values, child.type, undefined, child.nullable); }); diff --git a/nodejs/lancedb/blob.ts b/nodejs/lancedb/blob.ts index 7b0276f14..82729644c 100644 --- a/nodejs/lancedb/blob.ts +++ b/nodejs/lancedb/blob.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -import { Field, LargeBinary, Struct, Utf8 } from "apache-arrow"; +import { Field, LargeBinary, Struct, Uint64, Utf8 } from "apache-arrow"; import { BlobFile as NativeBlobFile } from "./native"; const BLOB_V2_EXTENSION_NAME = "lance.blob.v2"; @@ -28,6 +28,9 @@ export type BlobUri = string | URL; * * Either inline bytes, a URI pointing at external bytes, or a struct that sets * exactly one of `data` or `uri`. `null` and `undefined` write a null blob. + * External byte ranges require both + * `position` (non-negative) and `size` (positive). Use safe integer numbers or + * `bigint` values for these unsigned 64-bit byte counts. * * @example * ```ts @@ -45,8 +48,14 @@ export type BlobUri = string | URL; export type BlobInput = | BlobData | BlobUri - | { data: BlobData; uri?: null } - | { data?: null; uri: BlobUri } + | { data: BlobData; uri?: null; position?: null; size?: null } + | { data?: null; uri: BlobUri; position?: null; size?: null } + | { + data?: null; + uri: BlobUri; + position: number | bigint; + size: number | bigint; + } | null | undefined; @@ -54,6 +63,8 @@ export type BlobInput = export type BlobValue = { data: Buffer | Uint8Array | null; uri: string | null; + position?: bigint; + size?: bigint; }; export type BlobOptions = { @@ -78,6 +89,15 @@ export type BlobOptions = { /** * Declares a `lance.blob.v2` column. + * Uses nullable `data`, `uri`, `position`, and `size` fields for inline bytes, + * whole external files, and external byte ranges. Set both `position` and + * positive `size` for a range, or omit both for a whole file. + * + * Existing two-field columns remain writable. Appends and merge inserts fill + * missing range fields with nulls for four-field targets and drop all-null + * range fields for two-field targets. Non-null ranges require a four-field + * target. Overwrite with Arrow input replaces the table's schema; row objects + * are converted using the current table schema. * * Query results are descriptors, not payload bytes. Use {@link Table.fetchBlobs} * or {@link Table.fetchBlobFiles} to read bytes. @@ -142,6 +162,8 @@ export function blob(name: string, options: BlobOptions = {}): Field { new Struct([ new Field("data", new LargeBinary(), true), new Field("uri", new Utf8(), true), + new Field("position", new Uint64(), true), + new Field("size", new Uint64(), true), ]), options.nullable ?? true, metadata, @@ -317,6 +339,8 @@ export function normalizeBlobInput(value: unknown): unknown { ...value, data: value.data, uri: value.uri, + position: value.position, + size: value.size, }; let changed = false; if (value.data instanceof ArrayBuffer) { @@ -355,10 +379,19 @@ export function withBlobBytes(value: unknown, bytes: Uint8Array): unknown { return bytes; } const record = value as Record; - return { ...record, uri: record.uri, data: bytes }; + return { + ...record, + uri: record.uri, + position: record.position, + size: record.size, + data: bytes, + }; } -export function coerceBlobValue(input: unknown): BlobValue | null { +export function coerceBlobValue( + input: unknown, + field?: Field, +): BlobValue | null { const value = normalizeBlobInput(input); if (value == null) { return null; @@ -385,11 +418,11 @@ export function coerceBlobValue(input: unknown): BlobValue | null { ); } const unknownKeys = Object.keys(record).filter( - (key) => key !== "data" && key !== "uri", + (key) => !["data", "uri", "position", "size"].includes(key), ); if (unknownKeys.length > 0) { throw new Error( - `Blob struct values only support 'data' and 'uri', got ${unknownKeys + `Blob struct values only support 'data', 'uri', 'position', and 'size', got ${unknownKeys .map((key) => `'${key}'`) .join(", ")}`, ); @@ -416,13 +449,63 @@ export function coerceBlobValue(input: unknown): BlobValue | null { "Blob struct values must set exactly one of 'data' or 'uri'", ); } + const hasPosition = record.position != null; + const hasSize = record.size != null; + if (hasPosition !== hasSize) { + throw new Error( + "External blob must set both position and size, or neither", + ); + } + if (hasPosition) { + if (uriValue == null) { + throw new Error("External blob byte ranges require a uri"); + } + const position = blobRangeInteger(record.position, "position", 0n); + const size = blobRangeInteger(record.size, "size", 1n); + if (position + size > 0xffffffffffffffffn) { + throw new Error("External blob range overflows uint64"); + } + if ( + field !== undefined && + (!(field.type instanceof Struct) || + ["position", "size"].some( + (name) => + !field.type.children.some((child: Field) => child.name === name), + )) + ) { + throw new Error( + "Blob storage schema cannot store byte ranges. Use a four-field blob() schema to store position and size", + ); + } + return { data: null, uri: uriValue, position, size }; + } return { data: bytes, uri: uriValue }; } throw new Error( - "Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob, a URI string or URL, null, or { data?, uri? }", + "Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob, a URI string or URL, null, or { data?, uri?, position?, size? }", ); } +function blobRangeInteger( + value: unknown, + name: string, + minimum: bigint, +): bigint { + if ( + typeof value !== "bigint" && + (typeof value !== "number" || !Number.isSafeInteger(value)) + ) { + throw new Error(`Blob ${name} must be a safe integer or bigint`); + } + const integer = BigInt(value); + if (integer < minimum || integer > 0xffffffffffffffffn) { + throw new Error( + `Blob ${name} must be ${minimum === 0n ? "non-negative" : "positive"} and fit in uint64`, + ); + } + return integer; +} + function isBlobLike(value: unknown): value is Blob { return typeof Blob !== "undefined" && value instanceof Blob; } diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index a9cf13b65..e0a93f33e 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -588,7 +588,7 @@ def _cast_to_target_schema( def _preserve_blob_v2_fields( inferred_schema: pa.Schema, existing_schema: pa.Schema ) -> pa.Schema: - """Keep existing blob fields when an overwrite infers a replacement schema.""" + """Keep inferred blob semantics without overriding explicit replacement types.""" existing_fields = {field.name: field for field in existing_schema} fields = [ _preserve_blob_v2_field(field, existing_fields[field.name]) @@ -600,6 +600,9 @@ def _preserve_blob_v2_fields( def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field: + if is_blob_v2_field(inferred): + # An explicitly declared blob layout is part of the replacement schema. + return inferred if is_blob_v2_field(existing): return existing @@ -637,12 +640,10 @@ def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field: def _coerce_blob_write_columns( batch: pa.RecordBatch, target_schema: pa.Schema ) -> pa.RecordBatch: - """Materialize blob storage structs before the stream leaves Python. + """Materialize Python bytes and URIs as blob structs before Arrow casting. - merge_insert requires its source reader to already match the table's - physical schema. Unlike add and insert, it does not pass through - LanceDB's Rust blob coercion, so preserving binary input here would - reach Lance as binary and fail the schema check. + Existing blob structs retain their input layout. Rust normalizes them to + the stored schema and rejects populated ranges on two-field targets. """ columns = [] fields = [] @@ -973,6 +974,17 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field: "(Utf8, LargeUtf8, Utf8View), or a Struct with a " "'data' or 'uri' child" ) + if pa.types.is_struct(source_type) or ( + isinstance(source_type, pa.ExtensionType) + and ( + not isinstance(target_field.type, pa.ExtensionType) + or type(source_type) is type(target_field.type) + ) + ): + # Keep the source layout until Rust can fill missing range fields or + # validate every batch before dropping all-null range fields. Arrow + # casting here would reject differing layouts before that validation. + return field.with_metadata(target_field.metadata) # LanceDB exposes stored JSON columns as lance.json (JSONB-backed LargeBinary), but # casting the input to that storage type here merely relabels the raw JSON bytes as # JSONB. Lance must see arrow.json so it can perform the JSONB encoding. @@ -1832,8 +1844,8 @@ class Table(ABC): - pyarrow.Table or pyarrow.RecordBatch mode: str The mode to use when writing the data. Valid values are - "append" and "overwrite". Overwrite infers a replacement schema - but preserves existing blob v2 fields present in the new data. + "append" and "overwrite". Overwrite adopts explicitly declared blob + layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" What to do if any of the vectors are not the same size or contains NaNs. One of "error", "drop", "fill". @@ -4111,8 +4123,8 @@ class LanceTable(Table): The data to insert into the table. mode: str The mode to use when writing the data. Valid values are - "append" and "overwrite". Overwrite infers a replacement schema - but preserves existing blob v2 fields present in the new data. + "append" and "overwrite". Overwrite adopts explicitly declared blob + layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" What to do if any of the vectors are not the same size or contains NaNs. One of "error", "drop", "fill", "null". @@ -6145,8 +6157,8 @@ class AsyncTable: - pyarrow.Table or pyarrow.RecordBatch mode: str The mode to use when writing the data. Valid values are - "append" and "overwrite". Overwrite infers a replacement schema - but preserves existing blob v2 fields present in the new data. + "append" and "overwrite". Overwrite adopts explicitly declared blob + layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" What to do if any of the vectors are not the same size or contains NaNs. One of "error", "drop", "fill", "null". diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 130526eaf..ec1b8b27e 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -63,6 +63,25 @@ def _blob_array(name, values): return pa.ExtensionArray.from_storage(blob_type, storage) +def _blob_layout_data(rows, with_ranges, *, as_extension=True): + field = lancedb.blob("image") + storage_type = field.type.storage_type + if not with_ranges or not as_extension: + storage_type = pa.struct(list(storage_type)[: (4 if with_ranges else 2)]) + field = pa.field( + "image", + storage_type, + metadata={"ARROW:extension:name": "lance.blob.v2"}, + ) + images = pa.array([row["image"] for row in rows], type=storage_type) + if isinstance(field.type, pa.ExtensionType): + images = pa.ExtensionArray.from_storage(field.type, images) + return pa.Table.from_arrays( + [pa.array([row["id"] for row in rows], type=pa.int64()), images], + schema=pa.schema([pa.field("id", pa.int64()), field]), + ) + + def _row_ids_by_id(table): hits = table.search().with_row_id(True).limit(1000).to_arrow() assert "_rowid" in hits.column_names @@ -998,6 +1017,133 @@ def test_merge_insert_writes_python_bytes(): assert blobs.to_pylist() == [b"updated", b"inserted"] +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("input_type", ["extension", "struct", "reader"]) +@pytest.mark.parametrize("target_with_ranges", [False, True]) +async def test_merge_insert_normalizes_blob_layouts( + use_async, input_type, target_with_ranges +): + data = _blob_layout_data( + [{"id": 1, "image": {"data": b"before"}}], target_with_ranges + ) + table = lancedb.connect("memory:///").create_table("merge_layouts", data=data) + original_schema = table.schema + updates = _blob_layout_data( + [ + {"id": 1, "image": {"data": b"updated"}}, + {"id": 2, "image": {"data": b"inserted"}}, + {"id": 3, "image": None}, + ], + not target_with_ranges, + as_extension=input_type != "struct", + ) + if input_type == "reader": + updates = updates.to_reader(max_chunksize=1) + writer = table._table if use_async else table + result = ( + writer.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(updates) + ) + if use_async: + result = await result + + assert result.num_updated_rows == 1 + assert result.num_inserted_rows == 2 + assert table.schema == original_schema + by_id = _row_ids_by_id(table) + assert table.fetch_blobs("image", [by_id[i] for i in (1, 2, 3)]).to_pylist() == [ + b"updated", + b"inserted", + None, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + "byte_range", [{"position": 0}, {"size": 6}, {"position": 4, "size": 6}] +) +async def test_merge_insert_legacy_blob_rejects_ranges_atomically( + use_async, byte_range +): + data = _blob_layout_data([{"id": 1, "image": {"data": b"before"}}], False) + table = lancedb.connect("memory:///").create_table("reject_ranges", data=data) + original_schema, original_version = table.schema, table.version + updates = _blob_layout_data( + [ + {"id": 1, "image": {"data": b"updated"}}, + {"id": 2, "image": {"uri": "file:///payload", **byte_range}}, + ], + True, + ).to_reader(max_chunksize=1) + writer = table._table if use_async else table + with pytest.raises(RuntimeError, match="cannot store byte ranges"): + result = ( + writer.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(updates) + ) + if use_async: + await result + + assert table.version == original_version + assert table.schema == original_schema + assert table.count_rows() == 1 + assert table.fetch_blobs("image", [_row_ids_by_id(table)[1]]).to_pylist() == [ + b"before" + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("input_type", ["extension", "struct", "reader"]) +@pytest.mark.parametrize("source_with_ranges", [False, True]) +async def test_add_overwrite_adopts_explicit_blob_layout( + tmp_path, use_async, input_type, source_with_ranges +): + data = _blob_layout_data( + [{"id": 1, "image": {"data": b"before"}}], not source_with_ranges + ) + table = lancedb.connect("memory:///").create_table("overwrite_layouts", data=data) + payload = b"0123456789abcdef" + path = tmp_path / "payload.bin" + path.write_bytes(payload) + value = ( + {"uri": path.as_uri(), "position": 4, "size": 6} + if source_with_ranges + else {"data": b"updated"} + ) + updates = _blob_layout_data( + [{"id": 2, "image": value}, {"id": 3, "image": None}], + source_with_ranges, + as_extension=input_type != "struct", + ) + if input_type == "reader": + updates = updates.to_reader(max_chunksize=1) + writer = table._table if use_async else table + result = writer.add( + updates, mode="overwrite", allow_external_blob_outside_bases=True + ) + if use_async: + await result + + blob_type = table.schema.field("image").type + if isinstance(blob_type, pa.ExtensionType): + blob_type = blob_type.storage_type + assert blob_type.num_fields == (4 if source_with_ranges else 2) + assert table.blob_columns() == ["image"] + assert table.count_rows() == 2 + by_id = _row_ids_by_id(table) + assert table.fetch_blobs("image", [by_id[2], by_id[3]]).to_pylist() == [ + payload[4:10] if source_with_ranges else b"updated", + None, + ] + + @pytest.mark.parametrize("write", ["merge_insert", "create_table"]) def test_uri_string_blob_write_reaches_external_base_validation(tmp_path, write): blob_path = tmp_path / "payload.bin" diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index e9768b0bf..73fb6ab12 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -3254,6 +3254,89 @@ BLOB_DESCRIBE_RESPONSE = { } +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("mode", ["merge", "overwrite"]) +@pytest.mark.parametrize("target_with_ranges", [False, True]) +async def test_remote_writes_normalize_blob_layouts( + use_async, mode, target_with_ranges +): + description = json.loads(json.dumps(BLOB_DESCRIBE_RESPONSE)) + if target_with_ranges: + description["schema"]["fields"][1]["type"]["fields"].extend( + {"name": name, "type": {"type": "uint64"}, "nullable": True} + for name in ("position", "size") + ) + received = [] + + def handler(request): + if request.path == "/v1/table/test/describe/": + response = description + else: + expected = "merge_insert" if mode == "merge" else "insert" + assert urlparse(request.path).path == f"/v1/table/test/{expected}/" + if request.headers.get("Transfer-Encoding") == "chunked": + body = bytearray() + while size := int(request.rfile.readline().split(b";")[0], 16): + body.extend(request.rfile.read(size)) + assert request.rfile.read(2) == b"\r\n" + assert request.rfile.read(2) == b"\r\n" + else: + body = request.rfile.read(int(request.headers["Content-Length"])) + received.append(pa.ipc.open_stream(body).read_all()) + response = { + "version": 2, + "num_updated_rows": 0, + "num_inserted_rows": 2, + "num_deleted_rows": 0, + } + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.send_header("phalanx-version", "0.5.0") + request.end_headers() + request.wfile.write(json.dumps(response).encode()) + + field = lancedb.blob("image") + storage_type = field.type.storage_type + if target_with_ranges: + storage_type = pa.struct(list(storage_type)[:2]) + field = pa.field( + "image", + storage_type, + metadata={"ARROW:extension:name": "lance.blob.v2"}, + ) + images = pa.array([{"data": b"payload"}, None], type=storage_type) + if isinstance(field.type, pa.ExtensionType): + images = pa.ExtensionArray.from_storage(field.type, images) + data = pa.Table.from_arrays( + [pa.array([1, 2], type=pa.int64()), images], + schema=pa.schema([pa.field("id", pa.int64()), field]), + ) + + with mock_lancedb_connection(handler) as db: + remote_table = db.open_table("test") + table = remote_table._table if use_async else remote_table + if mode == "merge": + result = ( + table.merge_insert("id").when_not_matched_insert_all().execute(data) + ) + else: + result = table.add(data, mode="overwrite") + if use_async: + await result + + assert len(received) == 1 + images = received[0]["image"].combine_chunks() + if isinstance(images.type, pa.ExtensionType): + images = images.storage + with_ranges = target_with_ranges if mode == "merge" else not target_with_ranges + assert images.type.num_fields == (4 if with_ranges else 2) + expected = {"data": b"payload", "uri": None} + if with_ranges: + expected.update(position=None, size=None) + assert images.to_pylist() == [expected, None] + + def blob_query_response_table(): image_field = pa.field( "image", diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index ca68164e4..05650d7c2 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -244,8 +244,34 @@ impl BlobRangeRequest { /// Creates an Arrow field for a Lance blob v2 column. /// -/// `Struct` with the `lance.blob.v2` marker. Same layout Lance -/// expects on write. +/// `Struct` with nullable children and the +/// `lance.blob.v2` marker. Set both range fields to read a byte range from a +/// URI, or leave both null for a whole file. Inline bytes are also supported. +/// +/// # Migration +/// +/// This helper previously returned only `data` and `uri`. Rust callers building +/// a `StructArray` from this field must now supply four children: append typed +/// null UInt64 arrays for `position` and `size` when no range is needed: +/// +/// ``` +/// use std::sync::Arc; +/// use arrow_array::{ArrayRef, LargeBinaryArray, StringArray, StructArray, UInt64Array}; +/// use arrow_schema::DataType; +/// let field = lancedb::blob("image", true); +/// let DataType::Struct(children) = field.data_type() else { unreachable!() }; +/// let image = StructArray::new(children.clone(), vec![ +/// Arc::new(LargeBinaryArray::from(vec![Some(b"image".as_slice())])) as ArrayRef, +/// Arc::new(StringArray::from(vec![None::<&str>])), +/// Arc::new(UInt64Array::new_null(1)), +/// Arc::new(UInt64Array::new_null(1)), +/// ], None); +/// ``` +/// +/// Existing two-field tables remain writable: append and merge-insert fill +/// missing range fields or drop all-null range fields to match the table. +/// Non-null ranges cannot be written to a two-field table. Overwrite replaces +/// the table schema with the input schema rather than normalizing its layout. /// /// A blob column may be top-level or nested inside a struct or list. Blobs /// nested inside structs use dotted paths (e.g. `info.blob`) in the read APIs. @@ -261,9 +287,21 @@ impl BlobRangeRequest { /// ]); /// ``` pub fn blob(name: impl AsRef, nullable: bool) -> Field { - lance::blob::blob_field(name.as_ref(), nullable) + blob_with_options(name, nullable, BlobOptions::default()) } +static BLOB_STORAGE_TYPE: std::sync::LazyLock = std::sync::LazyLock::new(|| { + DataType::Struct( + vec![ + Field::new("data", DataType::LargeBinary, true), + Field::new("uri", DataType::Utf8, true), + Field::new("position", DataType::UInt64, true), + Field::new("size", DataType::UInt64, true), + ] + .into(), + ) +}); + /// Options for creating a blob v2 column field with tier thresholds. /// /// Re-exported from lance so the thresholds are built exactly the way @@ -293,6 +331,7 @@ pub use lance::blob::BlobFieldOptions as BlobOptions; /// ``` pub fn blob_with_options(name: impl AsRef, nullable: bool, options: BlobOptions) -> Field { lance::blob::blob_field_with_options(name.as_ref(), nullable, options) + .with_data_type(BLOB_STORAGE_TYPE.clone()) } /// Returns true if `field` is a blob v2 column. @@ -306,7 +345,7 @@ pub fn is_blob(field: &Field) -> bool { } /// Returns true if `field`, or any field nested under it, is a blob v2 column. -fn field_tree_has_blob_v2(field: &Field) -> bool { +pub(crate) fn field_tree_has_blob_v2(field: &Field) -> bool { if field.is_blob_v2() { return true; } @@ -673,7 +712,21 @@ mod tests { field.metadata().get(ARROW_EXT_NAME_KEY).map(String::as_str), Some("lance.blob.v2") ); - assert!(matches!(field.data_type(), DataType::Struct(_))); + let DataType::Struct(children) = field.data_type() else { + panic!("expected struct") + }; + assert_eq!( + children + .iter() + .map(|f| (f.name().as_str(), f.data_type(), f.is_nullable())) + .collect::>(), + vec![ + ("data", &DataType::LargeBinary, true), + ("uri", &DataType::Utf8, true), + ("position", &DataType::UInt64, true), + ("size", &DataType::UInt64, true), + ] + ); } #[test] diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index d21583c91..272ac83f7 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3154,6 +3154,17 @@ impl BaseTable for RemoteTable { let rescannable = source.rescannable(); let input: Arc = Arc::new(crate::table::datafusion::scannable_exec::ScannableExec::new(source, None)); + let table_schema = self.schema().await?; + let input = if crate::blob::has_blob_columns(&table_schema) { + crate::table::merge::validate_merge_source_fields( + input.schema().fields(), + table_schema.fields(), + "", + )?; + crate::table::datafusion::cast::cast_to_table_schema(input, &table_schema)? + } else { + input + }; let freshness_request = self.snapshot_freshness_headers(); let mut merge: Arc = Arc::new( @@ -4499,6 +4510,253 @@ mod tests { assert!(image.column_by_name("data").unwrap().is_null(0)); } + #[tokio::test] + async fn add_external_blob_ranges_preserves_remote_wire_values() { + let field = crate::blob("image", true); + let DataType::Struct(children) = field.data_type().clone() else { + unreachable!(); + }; + let schema = Arc::new(Schema::new(vec![field])); + // Registered external bases are enforced by the server, so remote + // clients send the URI and range without the local-only opt-in. + let uri = "s3://registered/payload.bin"; + let image = StructArray::new( + children, + vec![ + arrow_array::new_null_array(&DataType::LargeBinary, 2), + Arc::new(StringArray::from(vec![uri, uri])), + Arc::new(arrow_array::UInt64Array::from(vec![0, 4])), + Arc::new(arrow_array::UInt64Array::from(vec![6, 6])), + ], + None, + ); + let input = RecordBatch::try_new(schema.clone(), vec![Arc::new(image)]).unwrap(); + let describe_body = describe_response(&schema).into_bytes(); + let payload = b"0123456789abcdef"; + let fetched = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "image", + DataType::LargeBinary, + true, + )])), + vec![Arc::new(LargeBinaryArray::from(vec![ + &payload[..6], + &payload[4..10], + ]))], + ) + .unwrap(); + let fetched_body = write_ipc_stream_uncompressed(&fetched); + let (sender, receiver) = std::sync::mpsc::channel(); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |mut request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_body.clone()) + .unwrap(), + "/v1/table/my_table/insert/" => { + sender.send(request.body_mut().take().unwrap()).unwrap(); + http::Response::builder() + .status(200) + .body(br#"{"version": 2}"#.to_vec()) + .unwrap() + } + "/v1/table/my_table/fetch_blobs/" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["column"], "image"); + assert_eq!(body["row_ids"], json!([0, 1])); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .body(fetched_body.clone()) + .unwrap() + } + path => panic!("unexpected path: {path}"), + }, + ); + table.add(input).execute().await.unwrap(); + let body = collect_body(receiver.recv().unwrap()).await; + let batch = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None) + .unwrap() + .next() + .unwrap() + .unwrap(); + assert_eq!(batch.schema(), schema); + let image = batch.column(0).as_struct(); + let positions = image + .column(2) + .as_primitive::(); + let sizes = image + .column(3) + .as_primitive::(); + assert_eq!(positions.values().as_ref(), &[0, 4]); + assert_eq!(sizes.values().as_ref(), &[6, 6]); + assert_eq!(image.column(1).as_string::().value(0), uri); + assert_eq!(image.column(1).as_string::().value(1), uri); + let bytes = table.fetch_blobs("image", &[0, 1]).await.unwrap(); + assert_eq!(bytes.value(0), &payload[..6]); + assert_eq!(bytes.value(1), &payload[4..10]); + } + + #[tokio::test] + async fn remote_blob_layout_normalization_and_overwrite() { + use arrow_array::{ArrayRef, UInt64Array}; + for wide_target in [false, true] { + let field = |wide| { + if wide { + crate::blob("image", true) + } else { + lance::blob::blob_field("image", true) + } + }; + let target = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + field(wide_target), + ]); + for operation in ["add", "merge", "overwrite"] { + for range in [ + None, + Some((Some(4), None)), + Some((None, Some(6))), + Some((Some(4), Some(6))), + ] { + if range.is_some() && (wide_target || operation == "overwrite") { + continue; + } + // Range input always uses the full layout. Otherwise use the + // opposite layout to exercise both normalization directions. + let wide_input = range.is_some() || !wide_target; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + field(wide_input), + ])); + let DataType::Struct(children) = schema.field(1).data_type().clone() else { + unreachable!() + }; + let mut arrays: Vec = vec![ + Arc::new(LargeBinaryArray::from(vec![Some(b"inline".as_slice())])), + Arc::new(StringArray::from(vec![None::<&str>])), + ]; + if wide_input { + arrays.push(Arc::new(UInt64Array::from(vec![range.and_then(|r| r.0)]))); + arrays.push(Arc::new(UInt64Array::from(vec![range.and_then(|r| r.1)]))); + } + let image = Arc::new(StructArray::new(children, arrays, None)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(vec![1])), image], + ) + .unwrap(); + let describe = describe_response(&target); + let (sender, receiver) = std::sync::mpsc::channel(); + let table = Table::new_with_handler("my_table", move |mut request| { + if request.url().path().ends_with("/describe/") { + return http::Response::builder() + .status(200) + .body(describe.clone()) + .unwrap(); + } + assert!( + request.url().path().ends_with("/insert/") + || request.url().path().ends_with("/merge_insert/") + ); + if operation == "overwrite" { + assert!( + request + .url() + .query_pairs() + .any(|(k, v)| k == "mode" && v == "overwrite") + ); + } + sender.send(request.body_mut().take().unwrap()).unwrap(); + http::Response::builder() + .status(200) + .body(r#"{"version": 2, "num_inserted_rows": 1}"#.to_string()) + .unwrap() + }); + // A bad later batch must abort the request, including after + // an earlier valid batch was encoded. No write handler runs. + let mut batches = Vec::new(); + if range.is_some() { + let empty_range_arrays: Vec = vec![ + Arc::new(LargeBinaryArray::from(vec![Some(b"first".as_slice())])), + Arc::new(StringArray::from(vec![None::<&str>])), + Arc::new(UInt64Array::new_null(1)), + Arc::new(UInt64Array::new_null(1)), + ]; + let DataType::Struct(children) = schema.field(1).data_type().clone() else { + unreachable!() + }; + batches.push(Ok(RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![0])), + Arc::new(StructArray::new(children, empty_range_arrays, None)), + ], + ) + .unwrap())); + } + batches.push(Ok(batch)); + let reader: Box = + Box::new(RecordBatchIterator::new(batches, schema.clone())); + let result = match operation { + "merge" => { + let mut merge = table.merge_insert(&["id"]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all(); + merge.execute(reader).await.map(|_| ()) + } + "overwrite" => table + .add(reader) + .mode(AddDataMode::Overwrite) + .execute() + .await + .map(|_| ()), + _ => table.add(reader).execute().await.map(|_| ()), + }; + if !wide_target && range.is_some() && operation != "overwrite" { + assert!( + result + .unwrap_err() + .to_string() + .contains("cannot store byte ranges") + ); + assert!( + receiver.try_recv().is_err(), + "invalid stream must not reach a committed write" + ); + continue; + } + result.unwrap(); + let body = collect_body(receiver.recv().unwrap()).await; + let reader = StreamReader::try_new(Cursor::new(body), None).unwrap(); + let expected = if operation == "overwrite" { + schema.as_ref() + } else { + &target + }; + assert_eq!(reader.schema().as_ref(), expected); + for batch in reader { + let batch = batch.unwrap(); + let image = batch.column(1).as_struct(); + assert!( + image.column(0).as_binary::().value(0) == b"inline" + || image.column(0).as_binary::().value(0) == b"first" + ); + if image.num_columns() == 4 && range.is_none() { + assert!(image.column(2).is_null(0)); + assert!(image.column(3).is_null(0)); + } + } + } + } + } + } + #[rstest] #[case(true)] #[case(false)] @@ -4852,6 +5110,12 @@ mod tests { )); let table = Table::new_with_handler("my_table", move |request| { + if request.url().path().ends_with("/describe/") { + return http::Response::builder() + .status(200) + .body(r#"{"version": 1, "schema": {"fields": []}}"#) + .unwrap(); + } if request.url().path() == "/v1/table/my_table/merge_insert/" { assert_eq!(request.method(), "POST"); @@ -4908,6 +5172,12 @@ mod tests { )); let table = Table::new_with_handler("my_table", move |request| { + if request.url().path().ends_with("/describe/") { + return http::Response::builder() + .status(200) + .body(r#"{"version": 1, "schema": {"fields": []}}"#) + .unwrap(); + } assert_eq!(request.method(), "POST"); assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); let params = request.url().query_pairs().collect::>(); @@ -4993,6 +5263,12 @@ mod tests { let num_deleted_rows = if delete_unmatched { 3 } else { 0 }; let table = Table::new_with_handler("my_table", move |request| { + if request.url().path().ends_with("/describe/") { + return http::Response::builder() + .status(200) + .body(describe_response(&Schema::empty())) + .unwrap(); + } assert_eq!(request.method(), "POST"); assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); assert_eq!(request.headers()[CONTENT_TYPE], ARROW_STREAM_CONTENT_TYPE); @@ -5063,6 +5339,12 @@ mod tests { )); let table = Table::new_with_handler("my_table", move |request| { + if request.url().path().ends_with("/describe/") { + return http::Response::builder() + .status(200) + .body(r#"{"version": 1, "schema": {"fields": []}}"#) + .unwrap(); + } assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); // One repeated `on` per column, in the order the caller gave them. @@ -5134,6 +5416,12 @@ mod tests { // Default parameters let table = Table::new_with_handler("my_table", |request| { + if request.url().path().ends_with("/describe/") { + return http::Response::builder() + .status(200) + .body(r#"{"version": 1, "schema": {"fields": []}}"#) + .unwrap(); + } assert_eq!(request.method(), "POST"); assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); @@ -14405,6 +14693,20 @@ mod tests { )); let table = Table::new_with_handler("my_table", move |request| { match request.url().path() { + "/v1/table/my_table/describe/" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["branch"], "exp"); + http::Response::builder() + .status(200) + .body(describe_response(&Schema::new(vec![Field::new( + "id", + DataType::Int32, + false, + )]))) + .unwrap() + } "/v1/table/my_table/branches/create/" => http::Response::builder() .status(200) .body("{}".to_string()) diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index c82b85971..12404ec73 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use std::collections::HashSet; use std::sync::Arc; -use arrow_schema::{DataType, Fields, Schema}; +use arrow_schema::{DataType, Field, Fields, Schema}; use lance::dataset::WriteMode; +use lance_arrow::FieldExt; use serde::{Deserialize, Serialize}; use crate::data::scannable::Scannable; @@ -233,6 +235,7 @@ pub struct PreprocessingOutput { /// Fields can be in different orders, so match by name. /// /// If a column exists in input but not in table, error (no extra columns allowed). +/// Input field names must be unique within each struct. /// /// If a column exists in table but not in input, that is okay - it may be filled with nulls. /// @@ -244,7 +247,13 @@ fn validate_schema(input: &Schema, table: &Schema) -> Result<()> { } fn validate_fields(input: &Fields, table: &Fields) -> Result<()> { + let mut names = HashSet::with_capacity(input.len()); for field in input { + if !names.insert(field.name()) { + return Err(Error::InvalidInput { + message: format!("field '{}' is specified more than once", field.name()), + }); + } match table.iter().find(|f| f.name() == field.name()) { None => { return Err(Error::InvalidInput { @@ -252,17 +261,33 @@ fn validate_fields(input: &Fields, table: &Fields) -> Result<()> { }); } Some(table_field) => { - if let (DataType::Struct(in_children), DataType::Struct(tbl_children)) = - (field.data_type(), table_field.data_type()) - { - validate_fields(in_children, tbl_children)?; - } + validate_field(field, table_field)?; } } } Ok(()) } +fn validate_field(input: &Field, table: &Field) -> Result<()> { + // Blob coercion validates its children and permits dropping only + // all-null range fields when targeting a legacy two-field layout. + if table.is_blob_v2() { + return Ok(()); + } + match (input.data_type(), table.data_type()) { + (DataType::Struct(input), DataType::Struct(table)) => validate_fields(input, table), + ( + DataType::List(input) | DataType::LargeList(input) | DataType::FixedSizeList(input, _), + DataType::List(table) | DataType::LargeList(table) | DataType::FixedSizeList(table, _), + ) => { + // Preserve the item's extension metadata, and match nested struct + // children by name rather than comparing collection item names. + validate_field(input, table) + } + _ => Ok(()), + } +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -299,6 +324,81 @@ mod tests { conn.create_table("test", batch).execute().await.unwrap() } + #[rstest::rstest] + fn test_validate_append_collection_fields( + #[values("list", "large_list", "fixed_size_list")] source_kind: &str, + #[values("list", "large_list", "fixed_size_list")] target_kind: &str, + ) { + fn schema(kind: &str, item: Field) -> Schema { + let item = Arc::new(item); + let data_type = match kind { + "list" => DataType::List(item), + "large_list" => DataType::LargeList(item), + "fixed_size_list" => DataType::FixedSizeList(item, 1), + _ => unreachable!(), + }; + Schema::new(vec![Field::new("items", data_type, true)]) + } + let image = crate::blob("image", true); + let legacy_image = lance::blob::blob_field("image", true); + let caption = Field::new("caption", DataType::Utf8, true); + let input = schema( + source_kind, + Field::new( + "source_item", + DataType::Struct(vec![image.clone(), caption.clone()].into()), + true, + ), + ); + let target = schema( + target_kind, + Field::new( + "target_item", + DataType::Struct(vec![legacy_image.clone()].into()), + true, + ), + ); + let error = super::validate_schema(&input, &target) + .unwrap_err() + .to_string(); + assert!( + error.contains("caption") && error.contains("does not exist"), + "{error}" + ); + + let target = schema( + target_kind, + Field::new( + "target_item", + DataType::Struct(vec![caption.clone(), legacy_image.clone()].into()), + true, + ), + ); + // Declared siblings can be reordered, but must not occur more than once. + super::validate_schema(&input, &target).unwrap(); + let duplicate = schema( + source_kind, + Field::new( + "source_item", + DataType::Struct(vec![image.clone(), caption.clone(), caption].into()), + true, + ), + ); + let error = super::validate_schema(&duplicate, &target) + .unwrap_err() + .to_string(); + assert!( + error.contains("caption") && error.contains("specified more than once"), + "{error}" + ); + // Item names need not match, and blob children still reach range coercion. + super::validate_schema( + &schema(source_kind, image.with_name("source_item")), + &schema(target_kind, legacy_image.with_name("target_item")), + ) + .unwrap(); + } + async fn test_add_with_data(data: T) where T: Scannable + 'static, diff --git a/rust/lancedb/src/table/datafusion/blob_coerce.rs b/rust/lancedb/src/table/datafusion/blob_coerce.rs index 80b6cd6fc..07a200d16 100644 --- a/rust/lancedb/src/table/datafusion/blob_coerce.rs +++ b/rust/lancedb/src/table/datafusion/blob_coerce.rs @@ -5,6 +5,7 @@ //! //! [`super::cast::cast_to_table_schema`] calls [`coerce_blob_expr`]. +use std::collections::HashSet; use std::fmt; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -12,9 +13,9 @@ use std::sync::Arc; use arrow_array::{ Array, ArrayRef, BooleanArray, FixedSizeListArray, LargeListArray, ListArray, RecordBatch, }; -use arrow_buffer::OffsetBuffer; +use arrow_buffer::{BooleanBuffer, NullBuffer, NullBufferBuilder, OffsetBuffer}; use arrow_cast::cast; -use arrow_schema::{DataType, Field, FieldRef, Fields, Schema}; +use arrow_schema::{DataType, Field, FieldRef, Fields, Schema, SchemaRef}; use arrow_select::nullif::nullif; use datafusion::functions::core::{get_field, named_struct}; use datafusion_common::config::ConfigOptions; @@ -24,7 +25,6 @@ use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::expressions::{CastExpr, Literal}; use datafusion_physical_plan::PhysicalExpr; use datafusion_physical_plan::expressions::Column; -use lance_arrow::FieldExt; use crate::error::{Error, Result}; @@ -35,6 +35,7 @@ pub(super) fn coerce_blob_expr( input_field: &Field, table_field: &FieldRef, config: &Arc, + ancestors: &[Arc], ) -> Result<(Arc, FieldRef)> { let DataType::Struct(declared_fields) = table_field.data_type() else { return Err(Error::InvalidInput { @@ -69,6 +70,33 @@ pub(super) fn coerce_blob_expr( ), }); } + let mut names = HashSet::with_capacity(children.len()); + for child in children { + // Name-based extraction must not discard an ambiguous payload + // or silently ignore a duplicate range child during narrowing. + if !names.insert(child.name()) { + return Err(Error::InvalidInput { + message: format!( + "blob child '{}' in column '{}' is specified more than once", + child.name(), + table_field.name() + ), + }); + } + if !declared_fields + .iter() + .any(|field| field.name() == child.name()) + && !matches!(child.name().as_str(), "position" | "size") + { + return Err(Error::InvalidInput { + message: format!( + "unknown blob child '{}' in column '{}'", + child.name(), + table_field.name() + ), + }); + } + } BlobInputShape::Struct(children) } other => { @@ -155,22 +183,38 @@ pub(super) fn coerce_blob_expr( // `named_struct` always yields a valid struct, so a null input would land // as a row that set neither `data` nor `uri` -- not an absent blob but a // malformed one, which Lance rejects on write. + let dropped_ranges = match &input_shape { + BlobInputShape::Struct(children) => children + .iter() + .enumerate() + .filter(|(_, child)| { + matches!(child.name().as_str(), "position" | "size") + && !declared_fields + .iter() + .any(|field| field.name() == child.name()) + }) + .map(|(index, _)| index) + .collect(), + _ => Vec::new(), + }; let expr: Arc = Arc::new(AbsentBlobIsNull { source: input_expr, built, field: table_field.clone(), + dropped_ranges, + ancestors: ancestors.to_vec(), }); Ok((expr, table_field.clone())) } -/// Apply the normal blob coercion to list values while retaining the list's -/// offsets and validity. A child batch lets the existing blob expression handle -/// binary, null, and descriptor inputs exactly as it does for top-level fields. +/// Recurse through list items while retaining the list's offsets and validity. +/// A child batch reuses the field projection, including nested structs and lists, +/// so every blob leaf receives the same coercion as a top-level blob column. pub(super) fn coerce_blob_list_expr( input_expr: Arc, input_field: &Field, table_field: &FieldRef, - config: &Arc, + ancestors: &[Arc], ) -> Result, FieldRef)>> { let (DataType::List(input_item) | DataType::LargeList(input_item) @@ -184,23 +228,45 @@ pub(super) fn coerce_blob_list_expr( else { return Ok(None); }; - if !table_item.is_blob_v2() || input_item == table_item { + if !crate::blob::field_tree_has_blob_v2(table_item) || input_item == table_item { return Ok(None); } - let (item_expr, _) = coerce_blob_expr( - Arc::new(Column::new(input_item.name(), 0)), - input_item, - table_item, - config, - )?; + // List item names need not match across SDKs. Their position identifies the + // item, while named children inside a struct still match by name. + let input_fields = vec![Arc::new( + input_item.as_ref().clone().with_name(table_item.name()), + )] + .into(); + let validity_field = Field::new("__blob_parent_validity", DataType::Boolean, true); + let (item_expr, output_item) = super::cast::build_field_exprs( + &input_fields, + &vec![table_item.clone()].into(), + &|_| Arc::new(Column::new(input_item.name(), 0)), + &[Arc::new(Column::new(validity_field.name(), 1))], + )? + .pop() + .expect("one list item projection"); + let item_schema = Arc::new(Schema::new(vec![ + input_item.clone(), + Arc::new(validity_field), + ])); + let output_type = match table_field.data_type() { + DataType::List(_) => DataType::List(output_item), + DataType::LargeList(_) => DataType::LargeList(output_item), + DataType::FixedSizeList(_, size) => DataType::FixedSizeList(output_item, *size), + _ => unreachable!(), + }; + let output_field = Arc::new(table_field.as_ref().clone().with_data_type(output_type)); let expr: Arc = Arc::new(CoerceBlobList { source: input_expr, input_item: input_item.clone(), item_expr, - field: table_field.clone(), + field: output_field.clone(), + ancestors: ancestors.to_vec(), + item_schema, }); - Ok(Some((expr, table_field.clone()))) + Ok(Some((expr, output_field))) } #[derive(Debug, Clone)] @@ -209,6 +275,8 @@ struct CoerceBlobList { input_item: FieldRef, item_expr: Arc, field: FieldRef, + ancestors: Vec>, + item_schema: SchemaRef, } impl fmt::Display for CoerceBlobList { @@ -223,6 +291,7 @@ impl PartialEq for CoerceBlobList { && self.input_item == other.input_item && self.item_expr.eq(&other.item_expr) && self.field == other.field + && self.ancestors == other.ancestors } } @@ -234,6 +303,7 @@ impl Hash for CoerceBlobList { self.input_item.hash(state); self.item_expr.hash(state); self.field.hash(state); + self.ancestors.hash(state); } } @@ -263,34 +333,56 @@ impl PhysicalExpr for CoerceBlobList { } else { cast(source.as_ref(), &intermediate_type)? }; - let values = match source.data_type() { + let row_nulls = NullBuffer::union( + source.logical_nulls().as_ref(), + ancestor_nulls(&self.ancestors, batch)?.as_ref(), + ); + let (values, item_nulls) = match source.data_type() { DataType::List(_) => { let list = source.as_any().downcast_ref::().unwrap(); let first = list.offsets()[0] as usize; let last = *list.offsets().last().unwrap() as usize; - list.values().slice(first, last - first) + ( + list.values().slice(first, last - first), + expand_list_nulls(row_nulls.as_ref(), list.offsets().lengths(), last - first), + ) } DataType::LargeList(_) => { let list = source.as_any().downcast_ref::().unwrap(); let first = list.offsets()[0] as usize; let last = *list.offsets().last().unwrap() as usize; - list.values().slice(first, last - first) + ( + list.values().slice(first, last - first), + expand_list_nulls(row_nulls.as_ref(), list.offsets().lengths(), last - first), + ) } - DataType::FixedSizeList(_, _) => source - .as_any() - .downcast_ref::() - .unwrap() - .values() - .clone(), + DataType::FixedSizeList(_, size) => ( + source + .as_any() + .downcast_ref::() + .unwrap() + .values() + .clone(), + row_nulls.as_ref().map(|nulls| nulls.expand(*size as usize)), + ), other => { return Err(DataFusionError::Internal(format!( "coerce_blob_list expected a list, got {other}" ))); } }; + // Carry list/ancestor validity in a separate column. Applying it to the + // items themselves would introduce nulls into non-nullable child fields. + let validity = BooleanArray::new( + item_nulls + .as_ref() + .map(|nulls| nulls.inner().clone()) + .unwrap_or_else(|| BooleanBuffer::new_set(values.len())), + item_nulls, + ); let item_batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![self.input_item.clone()])), - vec![values.clone()], + self.item_schema.clone(), + vec![values.clone(), Arc::new(validity)], )?; let coerced = self .item_expr @@ -335,7 +427,9 @@ impl PhysicalExpr for CoerceBlobList { } fn children(&self) -> Vec<&Arc> { - vec![&self.source] + std::iter::once(&self.source) + .chain(self.ancestors.iter()) + .collect() } fn with_new_children( @@ -347,6 +441,8 @@ impl PhysicalExpr for CoerceBlobList { input_item: self.input_item.clone(), item_expr: self.item_expr.clone(), field: self.field.clone(), + ancestors: children[1..].to_vec(), + item_schema: self.item_schema.clone(), })) } @@ -355,6 +451,35 @@ impl PhysicalExpr for CoerceBlobList { } } +fn ancestor_nulls( + ancestors: &[Arc], + batch: &RecordBatch, +) -> datafusion_common::Result> { + let mut nulls = None; + for ancestor in ancestors { + let array = ancestor.evaluate(batch)?.into_array(batch.num_rows())?; + nulls = NullBuffer::union(nulls.as_ref(), array.logical_nulls().as_ref()); + } + Ok(nulls) +} + +fn expand_list_nulls( + nulls: Option<&NullBuffer>, + lengths: impl Iterator, + item_count: usize, +) -> Option { + let nulls = nulls?; + let mut builder = NullBufferBuilder::new(item_count); + for (length, valid) in lengths.zip(nulls.iter()) { + if valid { + builder.append_n_non_nulls(length); + } else { + builder.append_n_nulls(length); + } + } + builder.finish() +} + /// Carries the source column's nullity onto the struct built for it. /// /// This is its own expression rather than a `CASE` because the projection @@ -366,6 +491,8 @@ struct AbsentBlobIsNull { source: Arc, built: Arc, field: FieldRef, + dropped_ranges: Vec, + ancestors: Vec>, } impl fmt::Display for AbsentBlobIsNull { @@ -376,7 +503,11 @@ impl fmt::Display for AbsentBlobIsNull { impl PartialEq for AbsentBlobIsNull { fn eq(&self, other: &Self) -> bool { - self.source.eq(&other.source) && self.built.eq(&other.built) && self.field == other.field + self.source.eq(&other.source) + && self.built.eq(&other.built) + && self.field == other.field + && self.dropped_ranges == other.dropped_ranges + && self.ancestors == other.ancestors } } @@ -387,6 +518,8 @@ impl Hash for AbsentBlobIsNull { self.source.hash(state); self.built.hash(state); self.field.hash(state); + self.dropped_ranges.hash(state); + self.ancestors.hash(state); } } @@ -401,8 +534,30 @@ impl PhysicalExpr for AbsentBlobIsNull { fn evaluate(&self, batch: &RecordBatch) -> datafusion_common::Result { let rows = batch.num_rows(); - let built = self.built.evaluate(batch)?.into_array(rows)?; let source = self.source.evaluate(batch)?.into_array(rows)?; + if !self.dropped_ranges.is_empty() { + let ancestor_nulls = ancestor_nulls(&self.ancestors, batch)?; + let values = source + .as_any() + .downcast_ref::() + .unwrap(); + let effective_nulls = NullBuffer::union(values.nulls(), ancestor_nulls.as_ref()); + for index in &self.dropped_ranges { + let nulls = NullBuffer::union( + effective_nulls.as_ref(), + values.column(*index).logical_nulls().as_ref(), + ); + if nulls.as_ref().map_or(values.len() != 0, |nulls| { + nulls.null_count() != values.len() + }) { + return Err(DataFusionError::Execution(format!( + "blob storage schema cannot store byte ranges in column '{}': position and size must both be null", + self.field.name() + ))); + } + } + } + let built = self.built.evaluate(batch)?.into_array(rows)?; let Some(nulls) = source.logical_nulls() else { return Ok(ColumnarValue::Array(built)); }; @@ -413,7 +568,10 @@ impl PhysicalExpr for AbsentBlobIsNull { } fn children(&self) -> Vec<&Arc> { - vec![&self.source, &self.built] + [&self.source, &self.built] + .into_iter() + .chain(self.ancestors.iter()) + .collect() } fn with_new_children( @@ -424,6 +582,8 @@ impl PhysicalExpr for AbsentBlobIsNull { source: children[0].clone(), built: children[1].clone(), field: self.field.clone(), + dropped_ranges: self.dropped_ranges.clone(), + ancestors: children[2..].to_vec(), })) } @@ -449,7 +609,6 @@ fn typed_null(data_type: &DataType) -> Result> { mod tests { use super::super::cast::cast_to_table_schema; use super::*; - use crate::blob::blob; use arrow_array::{ Array, ArrayRef, BinaryArray, BinaryViewArray, FixedSizeListArray, Int32Array, Int64Array, LargeBinaryArray, LargeListArray, ListArray, NullArray, RecordBatch, StringArray, @@ -461,6 +620,7 @@ mod tests { use datafusion_catalog::MemTable; use datafusion_physical_plan::ExecutionPlan; use futures::TryStreamExt; + use lance::blob::blob_field as blob; use lance_arrow::FieldExt; use std::collections::HashMap; @@ -946,6 +1106,47 @@ mod tests { assert!(image.column_by_name("data").unwrap().is_null(0)); } + #[tokio::test] + async fn dropping_ranges_respects_parent_validity_and_slices() { + for parent_valid in [false, true] { + let DataType::Struct(children) = wide_blob_field("image").data_type().clone() else { + unreachable!() + }; + let input = StructArray::new( + children, + vec![ + Arc::new(LargeBinaryArray::from(vec![ + Some(b"ignored".as_slice()), + Some(b"payload".as_slice()), + ])), + Arc::new(StringArray::from(vec![None::<&str>, None])), + Arc::new(UInt64Array::from(vec![None, Some(4)])), + Arc::new(UInt64Array::new_null(2)), + ], + Some(NullBuffer::from(vec![true, parent_valid])), + ) + .slice(1, 1); + let batch = batch_with_image( + Field::new("image", input.data_type().clone(), true), + Arc::new(input), + ); + if parent_valid { + let plan = cast_to_table_schema(plan_from_batch(batch).await, &blob_table_schema()) + .unwrap(); + let error = plan + .execute(0, SessionContext::new().task_ctx()) + .unwrap() + .try_collect::>() + .await + .unwrap_err(); + assert!(error.to_string().contains("cannot store byte ranges")); + } else { + let output = coerce(batch, &blob_table_schema()).await; + assert!(image_struct(&output).is_null(0)); + } + } + } + #[tokio::test] async fn descriptor_struct_without_value_child_is_rejected() { let descriptor = StructArray::new( @@ -1094,4 +1295,890 @@ mod tests { let plan = cast_to_table_schema(input, &table_schema).unwrap(); assert_eq!(Arc::as_ptr(&plan), input_ptr, "no projection inserted"); } + + fn nested_layout_batch( + wide: bool, + id: i64, + range: Option, + nesting: &[&str], + ) -> RecordBatch { + nested_layout_batch_with_null(wide, id, range, nesting, None) + } + + fn nested_layout_batch_with_null( + wide: bool, + id: i64, + range: Option, + nesting: &[&str], + null_depth: Option, + ) -> RecordBatch { + nested_layout_batch_with_options(wide, id, range, nesting, null_depth, &[]) + } + + fn nested_layout_batch_with_options( + wide: bool, + id: i64, + range: Option, + nesting: &[&str], + null_depth: Option, + captions: &[&str], + ) -> RecordBatch { + // Ancestor validity must not introduce nulls into a required blob child. + let field = if wide { + crate::blob("image", false) + } else { + blob("image", false) + }; + let DataType::Struct(children) = field.data_type().clone() else { + unreachable!() + }; + let mut columns: Vec = vec![ + Arc::new(LargeBinaryArray::from(vec![Some(b"payload".as_slice())])), + Arc::new(StringArray::from(vec![None::<&str>])), + ]; + if wide { + columns.push(Arc::new(UInt64Array::from(vec![range]))); + columns.push(Arc::new(UInt64Array::from(vec![range]))); + } + let array: ArrayRef = Arc::new(StructArray::new(children, columns, None)); + nest_blob_batch(field, array, id, nesting, null_depth, captions) + } + + fn nest_blob_batch( + mut field: Field, + mut array: ArrayRef, + id: i64, + nesting: &[&str], + null_depth: Option, + captions: &[&str], + ) -> RecordBatch { + for (depth, kind) in nesting.iter().enumerate().rev() { + let nulls = (null_depth == Some(depth)).then(|| NullBuffer::from(vec![false])); + let data_type = match *kind { + "struct" => { + let mut fields = vec![Arc::new(field)]; + let mut columns = vec![array]; + if depth + 1 == nesting.len() { + for caption in captions { + fields.push(Arc::new(Field::new("caption", DataType::Utf8, true))); + columns.push(Arc::new(StringArray::from(vec![*caption]))); + } + } + let fields: Fields = fields.into(); + array = Arc::new(StructArray::new(fields.clone(), columns, nulls)); + DataType::Struct(fields) + } + "list" => { + let item = Arc::new(field.with_name("item")); + array = Arc::new(ListArray::new( + item.clone(), + OffsetBuffer::new(vec![0, 1].into()), + array, + nulls, + )); + DataType::List(item) + } + "large_list" => { + let item = Arc::new(field.with_name("item")); + array = Arc::new(LargeListArray::new( + item.clone(), + OffsetBuffer::new(vec![0_i64, 1].into()), + array, + nulls, + )); + DataType::LargeList(item) + } + "fixed_size_list" => { + let item = Arc::new(field.with_name("item")); + array = Arc::new(FixedSizeListArray::new(item.clone(), 1, array, nulls)); + DataType::FixedSizeList(item, 1) + } + _ => unreachable!(), + }; + field = Field::new("value", data_type, true); + } + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + field, + ])), + vec![Arc::new(Int64Array::from(vec![id])), array], + ) + .unwrap() + } + + fn duplicate_blob_child_batch(wide: bool, nesting: &[&str], child_index: usize) -> RecordBatch { + let field = if wide { + crate::blob("image", false) + } else { + blob("image", false) + }; + let DataType::Struct(children) = field.data_type() else { + unreachable!() + }; + let duplicate = children[child_index].clone(); + let mut fields = children.to_vec(); + fields.push(duplicate.clone()); + let mut columns = children + .iter() + .map(|child| arrow_array::new_null_array(child.data_type(), 1)) + .collect::>(); + columns[0] = Arc::new(LargeBinaryArray::from(vec![b"first".as_slice()])); + // Distinct duplicate values expose silent name-based selection. Range + // duplicates also must not bypass null-only narrowing to old layouts. + columns.push(match duplicate.name().as_str() { + "data" => Arc::new(LargeBinaryArray::from(vec![b"second".as_slice()])) as ArrayRef, + "uri" => Arc::new(StringArray::from(vec!["file:///discarded"])), + "position" | "size" => Arc::new(UInt64Array::from(vec![4])), + _ => unreachable!(), + }); + let fields: Fields = fields.into(); + let array = Arc::new(StructArray::new(fields.clone(), columns, None)); + let field = field.with_data_type(DataType::Struct(fields)); + nest_blob_batch(field, array, 1, nesting, None, &[]) + } + + #[rstest::rstest] + #[case(&[])] + #[case(&["struct"])] + #[case(&["list"])] + #[case(&["list", "struct"])] + #[case(&["list", "list", "struct"])] + #[case(&["struct", "list", "struct"])] + #[case(&["large_list", "struct"])] + #[case(&["fixed_size_list", "struct", "large_list"])] + #[tokio::test] + async fn projection_rejects_duplicate_blob_children( + #[case] nesting: &[&str], + #[values(false, true)] wide_source: bool, + #[values(false, true)] wide_target: bool, + ) { + let target = nested_layout_batch(wide_target, 0, None, nesting); + for child_index in 0..if wide_source { 4 } else { 2 } { + let input = duplicate_blob_child_batch(wide_source, nesting, child_index); + let error = coerce_err(input, target.schema().as_ref()) + .await + .to_string(); + assert!(error.contains("specified more than once"), "{error}"); + } + } + + #[cfg(feature = "remote")] + #[rstest::rstest] + #[case(&[])] + #[case(&["struct"])] + #[case(&["list"])] + #[case(&["list", "struct"])] + #[case(&["list", "list", "struct"])] + #[case(&["struct", "list", "struct"])] + #[case(&["large_list", "struct"])] + #[tokio::test] + async fn writes_reject_duplicate_blob_children( + #[case] nesting: &[&str], + #[values(false, true)] wide_source: bool, + #[values(false, true)] wide_target: bool, + #[values(false, true)] remote: bool, + #[values(false, true)] merge: bool, + ) { + use crate::query::ExecutableQuery; + + let initial = nested_layout_batch(wide_target, 0, None, nesting); + let schema = initial.schema(); + let db = crate::connect("memory://").execute().await.unwrap(); + let (table, receiver) = if remote { + let (table, receiver) = remote_layout_table(schema.as_ref()); + (table, Some(receiver)) + } else { + ( + db.create_table("descriptors", initial) + .execute() + .await + .unwrap(), + None, + ) + }; + let snapshot = if remote { + None + } else { + Some(( + table.version().await.unwrap(), + table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(), + )) + }; + for child_index in 0..if wide_source { 4 } else { 2 } { + let input = duplicate_blob_child_batch(wide_source, nesting, child_index); + let result = if merge { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let schema = input.schema(); + builder + .execute(Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(input)], + schema, + ))) + .await + .map(|_| ()) + } else { + table.add(input).execute().await.map(|_| ()) + }; + let error = result.unwrap_err().to_string(); + assert!(error.contains("specified more than once"), "{error}"); + if let Some(receiver) = &receiver { + assert!( + receiver.try_recv().is_err(), + "ambiguous write reached server" + ); + } else { + let (version, before) = snapshot.as_ref().unwrap(); + assert_eq!(table.version().await.unwrap(), *version); + assert_eq!(table.count_rows(None).await.unwrap(), 1); + assert_eq!(table.schema().await.unwrap(), schema); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(&batches, before); + } + } + } + + #[rstest::rstest] + #[case(&["struct"])] + #[case(&["struct", "struct"])] + #[case(&["list", "struct"])] + #[case(&["struct", "list", "struct"])] + #[case(&["large_list", "struct", "list"])] + #[case(&["fixed_size_list", "struct", "large_list"])] + #[tokio::test] + async fn nullable_nested_blob_layout_projection( + #[case] nesting: &[&str], + #[values(false, true)] wide_target: bool, + #[values(false, true)] sliced: bool, + ) { + for depth in 0..nesting.len() { + // Physical range slots beneath a null ancestor are not logical values. + let input = + nested_layout_batch_with_null(!wide_target, 1, Some(0), nesting, Some(depth)); + let input = if sliced { + let outside = nested_layout_batch(!wide_target, 0, Some(4), nesting); + arrow_select::concat::concat_batches( + &input.schema(), + &[outside.clone(), input, outside], + ) + .unwrap() + .slice(1, 1) + } else { + input + }; + let expected = + nested_layout_batch_with_null(wide_target, 1, None, nesting, Some(depth)); + let output = coerce(input, expected.schema().as_ref()).await; + assert_eq!(output, expected, "null ancestor at {depth} in {nesting:?}"); + } + } + + #[rstest::rstest] + #[case(&["list"])] + #[case(&["list", "struct"])] + #[case(&["list", "list"])] + #[case(&["large_list", "struct", "list"])] + #[case(&["fixed_size_list", "struct", "large_list"])] + #[case(&["struct", "list", "struct"])] + #[tokio::test] + async fn nested_blob_layout_projection( + #[case] nesting: &[&str], + #[values(false, true)] wide_target: bool, + #[values(false, true)] sliced: bool, + ) { + let input = nested_layout_batch(!wide_target, 1, None, nesting); + let input = if sliced { + // Populated ranges outside a sliced collection must not be examined. + let before = nested_layout_batch(!wide_target, 0, Some(4), nesting); + let after = nested_layout_batch(!wide_target, 2, Some(4), nesting); + arrow_select::concat::concat_batches(&input.schema(), &[before, input, after]) + .unwrap() + .slice(1, 1) + } else { + input + }; + let expected = nested_layout_batch(wide_target, 1, None, nesting); + let output = coerce(input, expected.schema().as_ref()).await; + assert_eq!(output, expected); + } + + #[cfg(feature = "remote")] + fn remote_layout_table( + schema: &Schema, + ) -> (crate::Table, std::sync::mpsc::Receiver) { + let describe = serde_json::json!({ + "version": 1, + "schema": lance::arrow::json::JsonSchema::try_from(schema).unwrap(), + }) + .to_string(); + let (sender, receiver) = std::sync::mpsc::channel(); + let table = crate::Table::new_with_handler("layouts", move |mut request| { + if request.url().path().ends_with("/describe/") { + return http::Response::builder() + .status(200) + .body(describe.clone()) + .unwrap(); + } + assert!( + request.url().path().ends_with("/insert/") + || request.url().path().ends_with("/merge_insert/") + ); + sender.send(request.body_mut().take().unwrap()).unwrap(); + http::Response::builder() + .status(200) + .body(r#"{"version":2,"num_inserted_rows":1}"#.to_string()) + .unwrap() + }); + (table, receiver) + } + + #[cfg(feature = "remote")] + #[rstest::rstest] + #[case(&["list", "struct"])] + #[case(&["list", "list", "struct"])] + #[case(&["struct", "list", "struct"])] + #[case(&["large_list", "struct"])] + #[case(&["large_list", "struct", "list", "struct"])] + #[tokio::test] + async fn nested_blob_append_validates_siblings( + #[case] nesting: &[&str], + #[values(false, true)] wide_source: bool, + #[values(false, true)] wide_target: bool, + #[values(false, true)] remote: bool, + ) { + use crate::query::ExecutableQuery; + use futures::StreamExt; + use http_body::Body; + + let db = crate::connect("memory://").execute().await.unwrap(); + let initial = nested_layout_batch(wide_target, 0, None, nesting); + let schema = initial.schema(); + let (table, receiver) = if remote { + let (table, receiver) = remote_layout_table(schema.as_ref()); + (table, Some(receiver)) + } else { + ( + db.create_table("no_caption", initial) + .execute() + .await + .unwrap(), + None, + ) + }; + let input = + nested_layout_batch_with_options(wide_source, 1, None, nesting, None, &["keep me"]); + let version = if remote { + None + } else { + Some(table.version().await.unwrap()) + }; + let error = table + .add(input.clone()) + .execute() + .await + .unwrap_err() + .to_string(); + assert!( + error.contains("caption") && error.contains("does not exist"), + "{error}" + ); + if let Some(receiver) = receiver { + assert!( + receiver.try_recv().is_err(), + "invalid append reached server" + ); + } else { + assert_eq!(table.version().await.unwrap(), version.unwrap()); + assert_eq!(table.count_rows(None).await.unwrap(), 1); + assert_eq!(table.schema().await.unwrap(), schema); + } + + // A declared sibling must survive normalization in every layout direction. + let initial = + nested_layout_batch_with_options(wide_target, 0, None, nesting, None, &["keep me"]); + let schema = initial.schema(); + let (table, receiver) = if remote { + let (table, receiver) = remote_layout_table(schema.as_ref()); + (table, Some(receiver)) + } else { + ( + db.create_table("with_caption", initial) + .execute() + .await + .unwrap(), + None, + ) + }; + table.add(input).execute().await.unwrap(); + if let Some(receiver) = receiver { + let mut body = receiver.recv().unwrap(); + let mut bytes = Vec::new(); + while let Some(frame) = + futures::stream::poll_fn(|cx| std::pin::Pin::new(&mut body).poll_frame(cx)) + .next() + .await + { + bytes.extend_from_slice(frame.unwrap().data_ref().unwrap()); + } + let batches = + arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!( + batches, + vec![nested_layout_batch_with_options( + wide_target, + 1, + None, + nesting, + None, + &["keep me"] + )] + ); + } else { + assert_eq!(table.count_rows(None).await.unwrap(), 2); + assert_eq!(table.schema().await.unwrap(), schema); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut found = false; + for batch in batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let Some(row) = (0..ids.len()).find(|row| ids.value(*row) == 1) else { + continue; + }; + let mut array = batch.column(1).slice(row, 1); + for kind in nesting { + array = match *kind { + "list" => array.as_any().downcast_ref::().unwrap().value(0), + "large_list" => array + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "struct" => { + let values = array.as_any().downcast_ref::().unwrap(); + if let Some(caption) = values.column_by_name("caption") { + assert_eq!( + caption + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "keep me" + ); + found = true; + break; + } + values.column(0).clone() + } + _ => unreachable!(), + }; + } + } + assert!(found, "declared caption was lost"); + } + } + + #[cfg(feature = "remote")] + #[rstest::rstest] + #[case(&["list", "struct"])] + #[case(&["list", "list", "struct"])] + #[case(&["struct", "list", "struct"])] + #[case(&["large_list", "struct"])] + #[case(&["large_list", "struct", "list", "struct"])] + #[tokio::test] + async fn nested_blob_append_rejects_duplicate_siblings( + #[case] nesting: &[&str], + #[values(false, true)] wide_source: bool, + #[values(false, true)] wide_target: bool, + #[values(false, true)] remote: bool, + ) { + use crate::query::ExecutableQuery; + + let initial = + nested_layout_batch_with_options(wide_target, 0, None, nesting, None, &["keep me"]); + let schema = initial.schema(); + let db = crate::connect("memory://").execute().await.unwrap(); + let (table, receiver) = if remote { + let (table, receiver) = remote_layout_table(schema.as_ref()); + (table, Some(receiver)) + } else { + ( + db.create_table("captions", initial) + .execute() + .await + .unwrap(), + None, + ) + }; + let snapshot = if remote { + None + } else { + Some(( + table.version().await.unwrap(), + table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(), + )) + }; + // Construct distinct arrays for the duplicate names: selecting only the + // first matching field must not silently discard the second value. + let input = nested_layout_batch_with_options( + wide_source, + 1, + None, + nesting, + None, + &["first", "second"], + ); + let error = table.add(input).execute().await.unwrap_err().to_string(); + assert!( + error.contains("caption") && error.contains("specified more than once"), + "{error}" + ); + if let Some(receiver) = receiver { + assert!( + receiver.try_recv().is_err(), + "ambiguous append reached server" + ); + } else { + let (version, before) = snapshot.unwrap(); + assert_eq!(table.version().await.unwrap(), version); + assert_eq!(table.count_rows(None).await.unwrap(), 1); + assert_eq!(table.schema().await.unwrap(), schema); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches, before); + } + } + + #[cfg(feature = "remote")] + #[rstest::rstest] + #[case(&["struct"])] + #[case(&["list", "struct"])] + #[case(&["struct", "list", "struct"])] + #[case(&["large_list", "list"])] + #[tokio::test] + async fn nullable_nested_blob_layout_writes( + #[case] nesting: &[&str], + #[values(false, true)] merge: bool, + #[values(false, true)] remote: bool, + ) { + use crate::query::ExecutableQuery; + use futures::StreamExt; + use http_body::Body; + + for depth in 0..nesting.len() { + let initial = nested_layout_batch(false, 0, None, nesting); + let schema = initial.schema(); + let db = crate::connect("memory://").execute().await.unwrap(); + let (table, receiver) = if remote { + let (table, receiver) = remote_layout_table(schema.as_ref()); + (table, Some(receiver)) + } else { + ( + db.create_table("layouts", initial).execute().await.unwrap(), + None, + ) + }; + let visible_id = if merge { 0 } else { 1 }; + let visible = nested_layout_batch(true, visible_id, None, nesting); + let hidden = nested_layout_batch_with_null(true, 2, Some(0), nesting, Some(depth)); + let outside = nested_layout_batch(true, 99, Some(4), nesting); + let input = arrow_select::concat::concat_batches( + &visible.schema(), + &[outside.clone(), visible, hidden, outside], + ) + .unwrap() + .slice(1, 2); + let expected = arrow_select::concat::concat_batches( + &schema, + &[ + nested_layout_batch(false, visible_id, None, nesting), + nested_layout_batch_with_null(false, 2, None, nesting, Some(depth)), + ], + ) + .unwrap(); + if merge { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let input_schema = input.schema(); + builder + .execute(Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(input)], + input_schema, + ))) + .await + .unwrap(); + } else { + table.add(input).execute().await.unwrap(); + } + let count = if merge { 2 } else { 3 }; + if let Some(receiver) = &receiver { + let mut body = receiver.recv().unwrap(); + let mut bytes = Vec::new(); + while let Some(frame) = + futures::stream::poll_fn(|cx| std::pin::Pin::new(&mut body).poll_frame(cx)) + .next() + .await + { + bytes.extend_from_slice(frame.unwrap().data_ref().unwrap()); + } + let batches = + arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches, vec![expected]); + } else { + assert_eq!(table.schema().await.unwrap(), schema); + assert_eq!(table.count_rows(None).await.unwrap(), count); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut found = false; + for batch in batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + if let Some(row) = (0..ids.len()).find(|row| ids.value(*row) == 2) { + let mut array = batch.column(1).slice(row, 1); + for kind in &nesting[..depth] { + array = match *kind { + "struct" => array + .as_any() + .downcast_ref::() + .unwrap() + .column(0) + .clone(), + "list" => { + array.as_any().downcast_ref::().unwrap().value(0) + } + "large_list" => array + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + _ => unreachable!(), + }; + } + assert!(array.is_null(0), "lost null ancestor at {depth}"); + found = true; + } + } + assert!(found); + } + + // A visible populated range in a later batch must still abort everything. + let version = if remote { + None + } else { + Some(table.version().await.unwrap()) + }; + let hidden = nested_layout_batch_with_null(true, 3, Some(0), nesting, Some(depth)); + let bad = nested_layout_batch(true, 4, Some(4), nesting); + let reader: Box = + Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(hidden), Ok(bad.clone())], + bad.schema(), + )); + let result = if merge { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder.execute(reader).await.map(|_| ()) + } else { + table.add(reader).execute().await.map(|_| ()) + }; + let error = result.unwrap_err().to_string(); + assert!(error.contains("cannot store byte ranges"), "{error}"); + if let Some(receiver) = receiver { + assert!( + receiver.try_recv().is_err(), + "rejected write reached server" + ); + } else { + assert_eq!(table.version().await.unwrap(), version.unwrap()); + assert_eq!(table.count_rows(None).await.unwrap(), count); + } + } + } + + #[cfg(feature = "remote")] + #[rstest::rstest] + #[case(&["list"])] + #[case(&["list", "struct"])] + #[case(&["list", "list"])] + #[case(&["large_list", "struct", "list"])] + #[case(&["struct", "list", "struct"])] + #[tokio::test] + async fn nested_blob_layout_writes( + #[case] nesting: &[&str], + #[values(false, true)] wide_target: bool, + #[values(false, true)] merge: bool, + #[values(false, true)] remote: bool, + ) { + use crate::query::ExecutableQuery; + use futures::StreamExt; + use http_body::Body; + + let initial = nested_layout_batch(wide_target, 0, None, nesting); + let schema = initial.schema(); + let db = crate::connect("memory://").execute().await.unwrap(); + let (table, receiver) = if remote { + let (table, receiver) = remote_layout_table(schema.as_ref()); + (table, Some(receiver)) + } else { + ( + db.create_table("layouts", initial).execute().await.unwrap(), + None, + ) + }; + let input = nested_layout_batch(!wide_target, 1, None, nesting); + let expected = nested_layout_batch(wide_target, 1, None, nesting); + let (input, expected) = if merge { + ( + arrow_select::concat::concat_batches( + &input.schema(), + &[nested_layout_batch(!wide_target, 0, None, nesting), input], + ) + .unwrap(), + arrow_select::concat::concat_batches( + &schema, + &[nested_layout_batch(wide_target, 0, None, nesting), expected], + ) + .unwrap(), + ) + } else { + (input, expected) + }; + if merge { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let schema = input.schema(); + builder + .execute(Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(input)], + schema, + ))) + .await + .unwrap(); + } else { + table.add(input).execute().await.unwrap(); + } + if let Some(receiver) = &receiver { + let mut body = receiver.recv().unwrap(); + let mut bytes = Vec::new(); + while let Some(frame) = + futures::stream::poll_fn(|cx| std::pin::Pin::new(&mut body).poll_frame(cx)) + .next() + .await + { + bytes.extend_from_slice(frame.unwrap().data_ref().unwrap()); + } + let batches = + arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches, vec![expected]); + } else { + assert_eq!(table.schema().await.unwrap(), schema); + assert_eq!(table.count_rows(None).await.unwrap(), 2); + // Read the stored nested descriptors as well as the row count. + table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + } + if !wide_target { + let version = if remote { + None + } else { + Some(table.version().await.unwrap()) + }; + let first = nested_layout_batch(true, 2, None, nesting); + let bad = nested_layout_batch(true, 3, Some(4), nesting); + let reader: Box = + Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(first), Ok(bad.clone())], + bad.schema(), + )); + let result = if merge { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder.execute(reader).await.map(|_| ()) + } else { + table.add(reader).execute().await.map(|_| ()) + }; + let error = result.unwrap_err().to_string(); + assert!(error.contains("cannot store byte ranges"), "{error}"); + if let Some(receiver) = receiver { + assert!( + receiver.try_recv().is_err(), + "a rejected range must not reach the server" + ); + } else { + assert_eq!(table.version().await.unwrap(), version.unwrap()); + assert_eq!(table.count_rows(None).await.unwrap(), 2); + } + } + } } diff --git a/rust/lancedb/src/table/datafusion/cast.rs b/rust/lancedb/src/table/datafusion/cast.rs index 063572cbe..b6e815d5e 100644 --- a/rust/lancedb/src/table/datafusion/cast.rs +++ b/rust/lancedb/src/table/datafusion/cast.rs @@ -33,9 +33,12 @@ pub fn cast_to_table_schema( return Ok(input); } - let exprs = build_field_exprs(input_schema.fields(), table_schema.fields(), &|idx| { - Arc::new(Column::new(input_schema.field(idx).name(), idx)) as Arc - })?; + let exprs = build_field_exprs( + input_schema.fields(), + table_schema.fields(), + &|idx| Arc::new(Column::new(input_schema.field(idx).name(), idx)) as Arc, + &[], + )?; let exprs: Vec<(Arc, String)> = exprs .into_iter() @@ -52,10 +55,13 @@ pub fn cast_to_table_schema( /// For each table field that exists in the input, produce an expression that /// reads from the input and casts if needed. Fields in the table but not in the /// input are omitted (the storage layer handles missing columns). -fn build_field_exprs( +/// `ancestors` carries validity for blob range validation without nullifying +/// non-nullable children hidden beneath a null parent. +pub(super) fn build_field_exprs( input_fields: &Fields, table_fields: &Fields, get_input_expr: &dyn Fn(usize) -> Arc, + ancestors: &[Arc], ) -> Result, FieldRef)>> { let config = Arc::new(ConfigOptions::default()); let mut result = Vec::new(); @@ -109,13 +115,14 @@ fn build_field_exprs( input_field, table_field, &config, + ancestors, )?); continue; } // Blob values in lists need the same coercion as top-level blob columns. if let Some(coerced) = - coerce_blob_list_expr(input_expr.clone(), input_field, table_field, &config)? + coerce_blob_list_expr(input_expr.clone(), input_field, table_field, ancestors)? { result.push(coerced); continue; @@ -138,30 +145,38 @@ fn build_field_exprs( (DataType::Struct(in_children), DataType::Struct(tbl_children)) if in_children != tbl_children => { - let sub_exprs = build_field_exprs(in_children, tbl_children, &|child_idx| { - let child_name = in_children[child_idx].name(); - Arc::new(ScalarFunctionExpr::new( - &format!("get_field({child_name})"), - get_field(), - vec![ - input_expr.clone(), - Arc::new(Literal::new(ScalarValue::from(child_name.as_str()))), - ], - Arc::new(in_children[child_idx].as_ref().clone()), - config.clone(), - )) as Arc - })?; + let mut child_ancestors = ancestors.to_vec(); + child_ancestors.push(input_expr.clone()); + let sub_exprs = build_field_exprs( + in_children, + tbl_children, + &|child_idx| { + let child_name = in_children[child_idx].name(); + Arc::new(ScalarFunctionExpr::new( + &format!("get_field({child_name})"), + get_field(), + vec![ + input_expr.clone(), + Arc::new(Literal::new(ScalarValue::from(child_name.as_str()))), + ], + Arc::new(in_children[child_idx].as_ref().clone()), + config.clone(), + )) as Arc + }, + &child_ancestors, + )?; let output_struct_fields: Fields = sub_exprs .iter() .map(|(_, f)| f.clone()) .collect::>() .into(); - let output_field: FieldRef = Arc::new(Field::new( - table_field.name(), - DataType::Struct(output_struct_fields), - table_field.is_nullable(), - )); + let output_field: FieldRef = Arc::new( + table_field + .as_ref() + .clone() + .with_data_type(DataType::Struct(output_struct_fields)), + ); // Build named_struct(lit("a"), expr_a, lit("b"), expr_b, ...) let mut ns_args: Vec> = Vec::new(); diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 66ec48baf..589353efa 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -254,7 +254,11 @@ fn canonicalize_merge_filter(filter: Option) -> Result Result<()> { +pub(crate) fn validate_merge_source_fields( + input: &Fields, + target: &Fields, + parent: &str, +) -> Result<()> { let mut names = HashSet::with_capacity(input.len()); for field in input { let path = if parent.is_empty() { @@ -273,7 +277,9 @@ fn validate_merge_source_fields(input: &Fields, target: &Fields, parent: &str) - .ok_or_else(|| Error::InvalidInput { message: format!("merge source field '{path}' is not present in the table schema"), })?; - validate_merge_source_type(field.data_type(), target_field.data_type(), &path)?; + if !crate::blob::is_blob(target_field) { + validate_merge_source_type(field.data_type(), target_field.data_type(), &path)?; + } } Ok(()) } @@ -296,7 +302,13 @@ fn validate_merge_source_type(input: &DataType, target: &DataType, path: &str) - | DataType::LargeListView(target), ) | (DataType::Map(input, _), DataType::Map(target, _)) => { - validate_merge_source_type(input.data_type(), target.data_type(), path) + // Collection items carry extension metadata just like named fields. + // Blob coercion, rather than struct validation, checks range narrowing. + if crate::blob::is_blob(target) { + Ok(()) + } else { + validate_merge_source_type(input.data_type(), target.data_type(), path) + } } (DataType::Dictionary(_, input), _) => validate_merge_source_type(input, target, path), (_, DataType::Dictionary(_, target)) => validate_merge_source_type(input, target, path), @@ -313,35 +325,18 @@ pub(crate) async fn execute_merge_insert( new_data: Box, ) -> Result { params.canonicalize_filters()?; + let schema = table.schema().await?; super::computed_columns::ensure_no_function_bindings_for_mutation( - table.schema().await?.as_ref(), + schema.as_ref(), "merge_insert", )?; - match lsm::lsm_dispatch_decision(table, ¶ms).await? { - lsm::LsmDispatch::Lsm(plan) => { - let future = - lsm::execute_lsm_merge_insert(table, plan, params.validate_single_shard, new_data); - return match params.timeout { - Some(timeout) => match tokio::time::timeout(timeout, future).await { - Ok(result) => result, - Err(_) => Err(Error::Runtime { - message: "merge insert timed out".to_string(), - }), - }, - None => future.await, - }; - } - lsm::LsmDispatch::Standard => {} - } - - let dataset = table.dataset.get().await?; - let schema = arrow_schema::Schema::from(dataset.schema()); - // JSON source fields must carry the stored extension metadata just as on - // append. Keep arrow.json text labelled until Lance encodes it as JSONB. - let source = if schema - .fields() - .iter() - .any(|field| lance_arrow::json::has_json_fields(field)) + // Normalize blob layouts and JSON extension metadata just as on append. + // Keep arrow.json text labelled until Lance encodes it as JSONB. + let source = if crate::blob::has_blob_columns(&schema) + || schema + .fields() + .iter() + .any(|field| lance_arrow::json::has_json_fields(field)) { validate_merge_source_fields(new_data.schema().fields(), schema.fields(), "")?; let plan = Arc::new(super::datafusion::scannable_exec::ScannableExec::new( @@ -356,6 +351,24 @@ pub(crate) async fn execute_merge_insert( } else { new_data.into_stream() }; + match lsm::lsm_dispatch_decision(table, ¶ms).await? { + lsm::LsmDispatch::Lsm(plan) => { + let future = + lsm::execute_lsm_merge_insert(table, plan, params.validate_single_shard, source); + return match params.timeout { + Some(timeout) => match tokio::time::timeout(timeout, future).await { + Ok(result) => result, + Err(_) => Err(Error::Runtime { + message: "merge insert timed out".to_string(), + }), + }, + None => future.await, + }; + } + lsm::LsmDispatch::Standard => {} + } + + let dataset = table.dataset.get().await?; let mut builder = LanceMergeInsertBuilder::try_new(dataset.clone(), params.on)?; match ( params.when_matched_update_all, diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index 39e4fbf28..d8877428d 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -24,8 +24,10 @@ use arrow_array::cast::AsArray; use arrow_array::types::{ Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; -use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, RecordBatchReader}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch}; use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef}; +use datafusion_physical_plan::SendableRecordBatchStream; +use futures::TryStreamExt; use lance::Dataset; use lance::dataset::mem_wal::{ DatasetMemWalExt, ShardWriter, ShardWriterConfig, evaluate_sharding_spec, @@ -707,7 +709,7 @@ pub(crate) async fn execute_lsm_merge_insert( table: &NativeTable, plan: LsmPlan, validate_single_shard: bool, - new_data: Box, + mut new_data: SendableRecordBatchStream, ) -> Result { let dataset = table.dataset.get().await?; let target_schema: SchemaRef = Arc::new(ArrowSchema::from(dataset.schema())); @@ -718,8 +720,7 @@ pub(crate) async fn execute_lsm_merge_insert( let mut batches: Vec = Vec::new(); let mut total_rows: u64 = 0; - for batch in new_data { - let batch = batch.map_err(|e| Error::Arrow { source: e })?; + while let Some(batch) = new_data.try_next().await? { if batch.num_rows() == 0 { continue; } diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 4274aaa6f..0aebcaf2d 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -1747,7 +1747,7 @@ mod tests { RecordBatch::try_new( Arc::new(Schema::new(vec![ Field::new("id", arrow_schema::DataType::Int32, false), - crate::blob("image", true), + lance::blob::blob_field("image", true), ])), vec![Arc::new(Int32Array::from(ids)), builder.finish().unwrap()], ) @@ -1901,7 +1901,7 @@ mod tests { let mut blob_builder = lance::blob::BlobArrayBuilder::new(2); blob_builder.push_bytes(b"nested").unwrap(); blob_builder.push_null().unwrap(); - let blob_field = crate::blob("image", true); + let blob_field = lance::blob::blob_field("image", true); let metadata_fields = Fields::from(vec![blob_field.clone()]); let metadata = StructArray::new( metadata_fields.clone(), @@ -1968,7 +1968,7 @@ mod tests { blob_builder.push_bytes(b"a").unwrap(); blob_builder.push_bytes(b"bb").unwrap(); blob_builder.push_null().unwrap(); - let item = Arc::new(crate::blob("item", true)); + let item = Arc::new(lance::blob::blob_field("item", true)); let images = ListArray::new( item.clone(), OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])), diff --git a/rust/lancedb/tests/blob_integration.rs b/rust/lancedb/tests/blob_integration.rs index 063690129..91f3a978c 100644 --- a/rust/lancedb/tests/blob_integration.rs +++ b/rust/lancedb/tests/blob_integration.rs @@ -10,12 +10,13 @@ use arrow_array::{ use arrow_schema::{DataType, Field, Fields, Schema}; use futures::TryStreamExt; use lance::Dataset; +use lance::blob::blob_field as blob; use lance::dataset::WriteParams; use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; use lance_table::format::BasePath; use lancedb::{ Connection, Error, Result, Table, - blob::{BlobRangeRequest, blob}, + blob::BlobRangeRequest, connect, connect_namespace, database::listing::{ ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS, @@ -1610,6 +1611,88 @@ async fn external_uri_struct_round_trips_with_flag() -> Result<()> { Ok(()) } +#[tokio::test] +async fn external_ranges_round_trip_with_registered_base_and_opt_in() -> Result<()> { + let tmp = tempdir().unwrap(); + let payload = b"0123456789abcdef"; + let uri = write_payload_file_uri(tmp.path(), "payload.bin", payload); + let base_uri = url::Url::from_directory_path(tmp.path()) + .unwrap() + .to_string(); + let db = connect(tmp.path().join("db").to_str().unwrap()) + .execute() + .await?; + + for registered_base in [false, true] { + let field = lancedb::blob("image", true); + let DataType::Struct(children) = field.data_type().clone() else { + unreachable!("blob field is a struct"); + }; + let schema = Arc::new(Schema::new(vec![field])); + let name = if registered_base { + "registered" + } else { + "opt_in" + }; + let table = db + .create_empty_table(name, schema.clone()) + .write_options(WriteOptions { + lance_write_params: Some(WriteParams { + initial_bases: registered_base.then(|| { + vec![BasePath { + id: 1, + name: Some("external".to_string()), + path: base_uri.clone(), + is_dataset_root: false, + }] + }), + ..Default::default() + }), + }) + .execute() + .await?; + let values = StructArray::new( + children, + vec![ + Arc::new(LargeBinaryArray::from(vec![ + None, + None, + None, + Some(b"inline".as_slice()), + ])), + Arc::new(StringArray::from(vec![ + Some(uri.as_str()), + Some(uri.as_str()), + Some(uri.as_str()), + None, + ])), + Arc::new(UInt64Array::from(vec![Some(0), Some(4), None, None])), + Arc::new(UInt64Array::from(vec![Some(6), Some(6), None, None])), + ], + None, + ); + table + .add(RecordBatch::try_new(schema, vec![Arc::new(values)])?) + .allow_external_blob_outside_bases(!registered_base) + .execute() + .await?; + let reopened = db.open_table(name).execute().await?; + let ids = collect_row_ids(&reopened).await?; + let bytes = reopened.fetch_blobs("image", &ids).await?; + assert_eq!(bytes.value(0), &payload[..6]); + assert_eq!(bytes.value(1), &payload[4..10]); + assert_eq!(bytes.value(2), payload); + assert_eq!(bytes.value(3), b"inline"); + let files = reopened.fetch_blob_files("image", &ids).await?; + assert_eq!(files[1].as_ref().unwrap().size(), 6); + assert_eq!( + files[1].as_ref().unwrap().read_range(1..3).await?, + &payload[5..7] + ); + } + Ok(()) +} + #[tokio::test] async fn external_uri_add_requires_opt_in() -> Result<()> { let tmp = tempdir().unwrap(); @@ -1761,3 +1844,188 @@ async fn malformed_string_uri_is_rejected_at_write() -> Result<()> { assert_eq!(table.count_rows(None).await?, 0); Ok(()) } + +fn layout_batch(wide: bool, id: i64, position: Option, size: Option) -> RecordBatch { + let field = if wide { + lancedb::blob("image", true) + } else { + blob("image", true) + }; + let DataType::Struct(children) = field.data_type().clone() else { + unreachable!() + }; + let mut columns: Vec = vec![ + Arc::new(LargeBinaryArray::from(vec![Some(b"payload".as_slice())])), + new_null_array(&DataType::Utf8, 1), + ]; + if wide { + columns.push(Arc::new(UInt64Array::from(vec![position]))); + columns.push(Arc::new(UInt64Array::from(vec![size]))); + } + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + field, + ])), + vec![ + Arc::new(Int64Array::from(vec![id])), + Arc::new(StructArray::new(children, columns, None)), + ], + ) + .unwrap() +} + +fn layout_reader(batches: Vec) -> Box { + let schema = batches[0].schema(); + Box::new(arrow_array::RecordBatchIterator::new( + batches.into_iter().map(Ok), + schema, + )) +} + +#[tokio::test] +async fn blob_layout_append_and_merge_normalize_both_directions() -> Result<()> { + for wide in [false, true] { + for merge in [false, true] { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let initial = layout_batch(wide, 1, None, None); + let schema = initial.schema(); + let table = db.create_table("layouts", initial).execute().await?; + let input = layout_batch(!wide, if merge { 1 } else { 2 }, None, None); + if merge { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let result = builder + .execute(layout_reader(vec![ + input, + layout_batch(!wide, 2, None, None), + ])) + .await?; + assert_eq!(result.num_updated_rows, 1); + assert_eq!(result.num_inserted_rows, 1); + } else { + table.add(input).execute().await?; + } + assert_eq!(table.schema().await?, schema); + assert_eq!(table.count_rows(None).await?, 2); + let ids = collect_row_ids(&table).await?; + let values = table.fetch_blobs("image", &ids).await?; + assert_eq!(values.value(0), b"payload"); + assert_eq!(values.value(1), b"payload"); + } + } + Ok(()) +} + +#[tokio::test] +async fn blob_layout_rejects_non_null_ranges_atomically_in_later_batches() -> Result<()> { + for merge in [false, true] { + for (position, size) in [(Some(0), None), (None, Some(6)), (Some(0), Some(6))] { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = db + .create_table("layouts", layout_batch(false, 1, None, None)) + .execute() + .await?; + let version = table.version().await?; + // A successful first batch must not commit when a later batch has a range. + let input = layout_reader(vec![ + layout_batch(true, 2, None, None), + layout_batch(true, 3, position, size), + ]); + let result = if merge { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder.execute(input).await.map(|_| ()) + } else { + table.add(input).execute().await.map(|_| ()) + }; + let error = result.unwrap_err().to_string(); + assert!(error.contains("cannot store byte ranges"), "{error}"); + assert_eq!(table.version().await?, version); + assert_eq!(table.count_rows(None).await?, 1); + let ids = collect_row_ids(&table).await?; + assert_eq!(table.fetch_blobs("image", &ids).await?.value(0), b"payload"); + } + } + Ok(()) +} + +#[tokio::test] +async fn blob_layout_overwrite_replaces_schema_in_both_directions() -> Result<()> { + for wide in [false, true] { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = db + .create_table("layouts", layout_batch(wide, 1, None, None)) + .execute() + .await?; + let input = layout_batch(!wide, 2, None, None); + let expected_schema = input.schema(); + table + .add(input) + .mode(AddDataMode::Overwrite) + .execute() + .await?; + assert_eq!(table.schema().await?, expected_schema); + assert_eq!(table.count_rows(None).await?, 1); + let ids = collect_row_ids(&table).await?; + assert_eq!(table.fetch_blobs("image", &ids).await?.value(0), b"payload"); + } + Ok(()) +} + +#[tokio::test] +async fn blob_layout_memwal_merge_normalizes_before_writing() -> Result<()> { + use lancedb::{query::Select, table::MemWalWriteSpec}; + for wide in [false, true] { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = db + .create_table("layouts", layout_batch(wide, 1, None, None)) + .execute() + .await?; + table.set_unenforced_primary_key(["id"]).await?; + table.enable_memwal(MemWalWriteSpec::unsharded()).await?; + let mut merge = table.merge_insert(&["id"]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all(); + merge + .clone() + .execute(layout_reader(vec![layout_batch(!wide, 2, None, None)])) + .await?; + let batches = table + .query() + .select(Select::columns(&["id"])) + .execute() + .await? + .try_collect::>() + .await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + if !wide { + let error = merge + .execute(layout_reader(vec![ + layout_batch(true, 3, None, None), + layout_batch(true, 4, Some(0), Some(6)), + ])) + .await + .unwrap_err(); + assert!(error.to_string().contains("cannot store byte ranges")); + let batches = table + .query() + .select(Select::columns(&["id"])) + .execute() + .await? + .try_collect::>() + .await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + } + } + Ok(()) +} diff --git a/rust/lancedb/tests/duplicate_pairs.rs b/rust/lancedb/tests/duplicate_pairs.rs index 63f1c9f5a..6b602ac4f 100644 --- a/rust/lancedb/tests/duplicate_pairs.rs +++ b/rust/lancedb/tests/duplicate_pairs.rs @@ -410,6 +410,8 @@ async fn dedup_query_materializes_blob_payloads_and_null_vectors() -> anyhow::Re (0..6).map(|id| vec![id; 4096]), )), Arc::new(StringArray::from(vec![None::<&str>; 6])), + Arc::new(UInt64Array::new_null(6)), + Arc::new(UInt64Array::new_null(6)), ], None, );