diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 431796843..3150d94ec 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -650,6 +650,48 @@ tables; no partial handles are returned. *** +### fetchBlobRanges() + +```ts +abstract fetchBlobRanges( + column, + requests, + options?): Promise<(null | Buffer)[]> +``` + +Materializes row-specific byte ranges at the table's current checkout. +Preserves request order, duplicates, and nulls. Zero length returns an empty +buffer for a non-null blob. Invalid ranges or missing rows reject the entire +call. See [Table.fetchBlobs](Table.md#fetchblobs) for row-ID validity across versions. + +Remote tables use the batched endpoint when the server advertises support, +splitting at 1024 ranges and 64 MiB, including individual larger ranges. +Older servers use descriptor queries and BlobFile range reads. Both protocols +use one exact branch and version. Errors from an advertised endpoint propagate +without switching protocols. All output bytes occupy memory. + +#### Parameters + +* **column**: `string` + +* **requests**: readonly [`BlobRangeRequest`](../type-aliases/BlobRangeRequest.md)[] + +* **options?**: [`BlobReadOptions`](../type-aliases/BlobReadOptions.md) + +#### Returns + +`Promise`<(`null` \| `Buffer`)[]> + +#### Example + +```ts +const bytes = await table.fetchBlobRanges("image", [ + { rowId: 10n, offset: 0n, length: 1024n }, +]); +``` + +*** + ### fetchBlobs() ```ts diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 19d7ceb97..beed7c8af 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -169,6 +169,7 @@ - [BlobMode](type-aliases/BlobMode.md) - [BlobOptions](type-aliases/BlobOptions.md) - [BlobRange](type-aliases/BlobRange.md) +- [BlobRangeRequest](type-aliases/BlobRangeRequest.md) - [BlobReadOptions](type-aliases/BlobReadOptions.md) - [BlobUri](type-aliases/BlobUri.md) - [Data](type-aliases/Data.md) diff --git a/docs/src/js/type-aliases/BlobRangeRequest.md b/docs/src/js/type-aliases/BlobRangeRequest.md new file mode 100644 index 000000000..daa231260 --- /dev/null +++ b/docs/src/js/type-aliases/BlobRangeRequest.md @@ -0,0 +1,39 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BlobRangeRequest + +# Type Alias: BlobRangeRequest + +```ts +type BlobRangeRequest: object; +``` + +One row-specific range for [Table.fetchBlobRanges](../classes/Table.md#fetchblobranges). + +## Type declaration + +### length + +```ts +length: bigint; +``` + +Number of bytes to read, as a uint64. Zero returns an empty non-null value. + +### offset + +```ts +offset: bigint; +``` + +Byte offset from the beginning of the logical blob, as a uint64. + +### rowId + +```ts +rowId: bigint | number; +``` + +A row ID from a query with row IDs enabled. diff --git a/docs/web/management/table.mdx b/docs/web/management/table.mdx index 031342fbc..bf6bf3edd 100644 --- a/docs/web/management/table.mdx +++ b/docs/web/management/table.mdx @@ -736,6 +736,10 @@ The other fetch methods load bytes directly: | Whole payloads for a set of rows | `fetch_blobs` | `fetchBlobs` | `fetch_blobs` | | Byte slices across many rows in one call | `fetch_blob_ranges` | Not available. Use `fetchBlobFiles` and `readRanges`. | `fetch_blob_ranges` | +{/* Draft for the TypeScript cross-row Blob range release. Replace the row above: +| Byte slices across many rows in one call | `fetch_blob_ranges` | `fetchBlobRanges` | `fetch_blob_ranges` | +*/} + - `fetch_blobs` returns the full bytes for each row ID in order. Null blobs produce null entries, and empty blobs produce empty byte strings. - `fetch_blob_files` loads bytes when you call `read`, `read_range`, or @@ -759,6 +763,21 @@ The other fetch methods load bytes directly: +{/* Draft for the TypeScript cross-row Blob range release. + +Upcoming-release draft: TypeScript adds `fetchBlobRanges` for cross-row slices. +Until that SDK release is available, use `fetchBlobFiles` and `readRanges`. +Server-side batching requires a deployment with the validated range capability; +older servers continue using the descriptor and Range path selected before the +read. An error from the selected batched endpoint is returned without fallback. +A batched call keeps the native snapshot identity from its first range response +across every partition and requires each serving node to verify it. Missing identity +support or a replaced/deleted snapshot fails the whole call; upgrade all serving +nodes before selecting the batched endpoint with a new SDK. +Custom TypeScript `Table` subclasses must implement the new abstract method. + +*/} + Python and Rust read slices across rows with `fetch_blob_ranges`. On a handle, `read_ranges` reads more than one slice of one blob. Python takes `(offset, length)` pairs, TypeScript takes `{ start, end }` ranges with an @@ -770,6 +789,14 @@ exclusive `end`, and Rust takes `start..end` ranges: + {/* Draft for the TypeScript cross-row Blob range release. + {`// Upcoming-release draft: slices from more than one row. +const slices = await blobTable.fetchBlobRanges("image", [ + { rowId: rowIds[0]!, offset: 0n, length: 6n }, + { rowId: rowIds[1]!, offset: 7n, length: 5n }, +]); +`}{"\n"} + */} {TsBlobApiRanges} diff --git a/nodejs/__test__/blob_ranges.test.ts b/nodejs/__test__/blob_ranges.test.ts new file mode 100644 index 000000000..7d2fab566 --- /dev/null +++ b/nodejs/__test__/blob_ranges.test.ts @@ -0,0 +1,589 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import { createHash } from "crypto"; +import { getEventListeners } from "events"; +import { createReadStream } from "fs"; +import * as http from "http"; +import { AddressInfo } from "net"; +import { + Table as ArrowTable, + Binary, + Field, + LargeBinary, + Struct, + Uint8, + Uint32, + Uint64, + Utf8, + tableToIPC, + vectorFromArray, +} from "apache-arrow"; +import { BlobRangeRequest, Connection, Table, connect } from "../lancedb"; + +jest.setTimeout(30000); + +const IDENTITY = `v1:${"a".repeat(64)}`; +const MIB = 1024 * 1024; +type Range = { + // biome-ignore lint/style/useNamingConvention: Matches the server JSON contract. + row_id: number; + offset: number; + length: number; +}; +type Request = { + column: string; + branch: string | null; + version: number | null; + ranges: Range[]; + identity?: string; +}; + +function payload(row: number, offset: number, length: number, version: number) { + const bytes = Buffer.alloc(length); + for (let i = 0; i < length; i++) { + bytes[i] = (row * 17 + ((offset + i) % 251) + version) % 251; + } + return bytes; +} + +async function withRangeServer( + run: (db: Connection, calls: Request[], paths: string[]) => Promise, + fault?: string | number, + capability: string | null = "1", + onStall?: () => void, +) { + let takes = 0; + const calls: Request[] = []; + const paths: string[] = []; + const server = http.createServer((req, res) => { + req.setTimeout(15000); + paths.push(req.url ?? ""); + if ( + req.url === "/v1/table/test/describe/" || + req.url === "/v1/table/test/branches/list/" + ) { + res.writeHead(200, { + "Content-Type": "application/json", + "phalanx-version": "1.13.0-beta.0", + ...(req.url.endsWith("/describe/") && capability !== null + ? { "x-lancedb-blob-ranges": capability } + : {}), + }); + const response = req.url.endsWith("branches/list/") + ? { + branches: { + experiment: { + parentBranch: null, + parentVersion: 1, + createAt: 1, + manifestSize: 1, + }, + }, + } + : { version: 8, schema: { fields: [] } }; + res.end(JSON.stringify(response)); + return; + } + if (capability !== "1" && req.method === "GET") { + const url = new URL(req.url!, "http://localhost"); + expect(url.searchParams.get("branch")).toBe("experiment"); + expect(url.searchParams.get("version")).toBe("3"); + const [start, end] = req.headers.range!.slice(6).split("-").map(Number); + const bytes = Buffer.from("alpha").subarray(start, end + 1); + res + .writeHead(206, { + "Content-Range": `bytes ${start}-${end}/5`, + "Content-Length": bytes.length, + "x-lancedb-version": "3", + }) + .end(bytes); + return; + } + expect(req.method).toBe("POST"); + expect(req.url).toBe( + capability === "1" + ? "/v1/table/test/fetch_blob_ranges/" + : "/v1/table/test/query/", + ); + const chunks: Buffer[] = []; + let size = 0; + req.on("data", (chunk: Buffer) => { + size += chunk.length; + expect(size).toBeLessThanOrEqual(256 * 1024); + chunks.push(chunk); + }); + req.on("end", () => { + if (capability !== "1") { + const body = JSON.parse(Buffer.concat(chunks).toString()) as { + branch: string; + version: number; + filter: string; + }; + expect(body.branch).toBe("experiment"); + expect(body.version).toBe(3); + takes++; + const ids = (body.filter.match(/\d+/g) ?? []) + .map(Number) + .filter((id) => [10, 20, 40].includes(id)) + .reverse(); + const descriptor = new Struct([ + new Field("kind", new Uint8(), false), + new Field("position", new Uint64(), false), + new Field("size", new Uint64(), false), + new Field("blob_id", new Uint32(), false), + new Field("blob_uri", new Utf8(), false), + ]); + const arrow = new ArrowTable({ + image: vectorFromArray( + ids.map((id) => + id === 20 + ? null + : { + kind: 0, + position: 0n, + size: id === 40 ? 0n : 5n, + // biome-ignore lint/style/useNamingConvention: Arrow descriptor field. + blob_id: 0, + // biome-ignore lint/style/useNamingConvention: Arrow descriptor field. + blob_uri: "", + }, + ), + descriptor, + ), + _rowid: vectorFromArray(ids.map(BigInt), new Uint64()), + }); + expect(takes).toBeGreaterThan(0); + res + .writeHead(200, { + "Content-Type": "application/vnd.apache.arrow.stream", + "x-lancedb-version": "3", + }) + .end(tableToIPC(arrow, "stream")); + return; + } + const body = JSON.parse(Buffer.concat(chunks).toString()) as Request; + body.identity = req.headers["x-lancedb-query-snapshot"] as + | string + | undefined; + calls.push(body); + expect(body.column).toBe("image"); + expect(body.branch).toBe("experiment"); + expect(body.ranges.length).toBeLessThanOrEqual(1024); + expect( + body.ranges.reduce((sum, range) => sum + range.length, 0), + ).toBeLessThanOrEqual(64 * MIB); + const version = body.version ?? (calls.length === 1 ? 7 : 8); + const failure = + calls.length > 1 || String(fault).startsWith("first-") + ? fault + : undefined; + if (failure === "recreated" || failure === "gone") { + res + .writeHead(failure === "recreated" ? 400 : 404) + .end("query snapshot identity mismatch: source gone or replaced"); + return; + } + if (failure === "stall" && calls.length === 2) { + res.writeHead(200, { + "Content-Type": "application/vnd.apache.arrow.stream", + "x-lancedb-version": "7", + "x-lancedb-query-snapshot": IDENTITY, + }); + res.flushHeaders(); + onStall?.(); + return; + } + if (typeof failure === "number") { + res.writeHead(failure).end("ordinary server failure"); + return; + } + if (failure === "disconnect") { + req.socket.destroy(); + return; + } + let values = body.ranges.map((range) => + range.row_id === 20 + ? null + : payload(range.row_id, range.offset, range.length, version), + ); + if (failure === "count") values = []; + if (failure === "length") values = [Buffer.alloc(0)]; + const arrow = new ArrowTable({ + [failure === "column" ? "wrong" : "value"]: vectorFromArray( + values, + failure === "type" ? new Binary() : new LargeBinary(), + ), + }); + let ipc = tableToIPC(arrow, "stream"); + if (failure === "eos") ipc = ipc.subarray(0, ipc.length - 8); + if (failure === "body") ipc = ipc.subarray(0, Math.floor(ipc.length / 2)); + res.writeHead(200, { + "Content-Type": "application/vnd.apache.arrow.stream", + connection: "close", + "Content-Length": ipc.length + (failure === "transport" ? 10 : 0), + ...(failure === "version" ? {} : { "x-lancedb-version": version }), + ...(["missing-identity", "first-missing-identity"].includes( + String(failure), + ) + ? {} + : { + "x-lancedb-query-snapshot": + failure === "wrong-identity" + ? `v1:${"b".repeat(64)}` + : failure === "first-invalid-identity" + ? "broken" + : IDENTITY, + }), + }); + // Fragment the generated IPC payload across actual HTTP writes. + for (let i = 0; i < ipc.length; i += 8192) + res.write(ipc.subarray(i, i + 8192)); + res.end(); + }); + }); + server.listen(0, "127.0.0.1"); + await new Promise((resolve) => server.once("listening", resolve)); + const port = (server.address() as AddressInfo).port; + const db = await connect("db://dev", { + apiKey: "fake", + hostOverride: `http://127.0.0.1:${port}`, + clientConfig: { + retryConfig: { retries: 2, connectRetries: 2, readRetries: 2 }, + blobRequestConcurrency: 1, + timeoutConfig: { timeout: 20 }, + }, + }); + try { + await run(db, calls, paths); + } finally { + await new Promise((resolve) => { + server.closeAllConnections(); + server.close(() => resolve()); + }); + } +} + +async function open(db: Connection, version?: number): Promise { + return db.openTable("test", [], { branch: "experiment", version }); +} + +describe("fetchBlobRanges", () => { + it.each([null, "0", "2"])( + "uses old routes before data requests for capability %s", + async (capability) => { + await withRangeServer( + async (db, calls, paths) => { + // Reopening exercises both the initial and cached database describe paths. + for (let i = 0; i < 2; i++) { + const table = await open(db, 3); + const requests = [ + { rowId: 10n, offset: 1n, length: 3n }, + { rowId: 20n, offset: 0n, length: 1n }, + { rowId: 40n, offset: 0n, length: 0n }, + { rowId: 10n, offset: 5n, length: 0n }, + { rowId: 10n, offset: 1n, length: 3n }, + ]; + expect(await table.fetchBlobRanges("image", requests)).toEqual([ + Buffer.from("lph"), + null, + Buffer.alloc(0), + Buffer.alloc(0), + Buffer.from("lph"), + ]); + for (const row of [ + { rowId: 10n, offset: 4n, length: 2n }, + { rowId: 99n, offset: 0n, length: 1n }, + ]) + await expect( + table.fetchBlobRanges("image", [row]), + ).rejects.toThrow(); + const before = paths.length; + expect(await table.fetchBlobRanges("image", [])).toEqual([]); + expect(paths).toHaveLength(before); + } + expect(calls).toHaveLength(0); + expect(paths.some((path) => path.includes("fetch_blob_ranges"))).toBe( + false, + ); + }, + undefined, + capability, + ); + }, + ); + + it("renegotiates advertised support when reopening a cached table", async () => { + await withRangeServer(async (db, calls) => { + for (let i = 0; i < 2; i++) { + const table = await open(db); + expect( + await table.fetchBlobRanges("image", [ + { rowId: 10n, offset: 0n, length: 1n }, + ]), + ).toEqual([payload(10, 0, 1, i === 0 ? 7 : 8)]); + } + expect(calls).toHaveLength(2); + }); + }); + + it("cancels while waiting for a later response body and releases the only permit", async () => { + let markStalled!: () => void; + const stalled = new Promise((resolve) => { + markStalled = resolve; + }); + await withRangeServer( + async (db, calls) => { + const table = await open(db); + const controller = new AbortController(); + const pending = table.fetchBlobRanges( + "image", + Array.from({ length: 2049 }, () => ({ + rowId: 10n, + offset: 0n, + length: 1n, + })), + { signal: controller.signal }, + ); + const rejected = expect(pending).rejects.toMatchObject({ + name: "AbortError", + }); + await stalled; + controller.abort(); + await rejected; + expect(getEventListeners(controller.signal, "abort")).toHaveLength(0); + expect(calls).toHaveLength(2); + expect( + await table.fetchBlobRanges("image", [ + { rowId: 10n, offset: 0n, length: 1n }, + ]), + ).toEqual([payload(10, 0, 1, 8)]); + expect(calls).toHaveLength(3); + }, + "stall", + "1", + () => markStalled(), + ); + }); + it.each([undefined, 3])( + "preserves order, nulls, duplicates and snapshot (checkout %s)", + async (version) => { + await withRangeServer(async (db, calls, paths) => { + const table = await open(db, version); + const requests: BlobRangeRequest[] = Array.from( + { length: 1025 }, + (_, i) => ({ + rowId: i % 9 === 0 ? 20 : 10 + (i % 7), + offset: BigInt(i), + length: BigInt(i % 4), + }), + ); + requests.push( + { rowId: 10n, offset: 0n, length: 0n }, + { rowId: 10n, offset: 2n, length: 3n }, + { rowId: 10n, offset: 2n, length: 3n }, + ); + const result = await table.fetchBlobRanges("image", requests); + expect(result).toEqual( + requests.map((range) => + Number(range.rowId) === 20 + ? null + : payload( + Number(range.rowId), + Number(range.offset), + Number(range.length), + version ?? 7, + ), + ), + ); + expect(calls.map((body) => body.version)).toEqual([ + version ?? null, + version ?? 7, + ]); + expect( + await table.fetchBlobRanges("image", [ + { rowId: 10n, offset: 0n, length: 1n }, + ]), + ).toEqual([payload(10, 0, 1, version ?? 8)]); + expect(calls[2].version).toBe(version ?? null); + expect( + paths.filter((path) => path.endsWith("fetch_blob_ranges/")), + ).toHaveLength(3); + }); + }, + ); + + it("partitions the byte budget and reassembles one range larger than 64 MiB", async () => { + await withRangeServer(async (db, calls) => { + const table = await open(db); + const requests = [ + { rowId: 10n, offset: 3n, length: BigInt(65 * MIB) }, + { rowId: 30n, offset: 7n, length: 2n }, + ]; + const result = await table.fetchBlobRanges("image", requests); + for (const [i, range] of requests.entries()) { + expect(result[i]?.length).toBe(Number(range.length)); + expect( + createHash("sha256") + .update(result[i] as Buffer) + .digest("hex"), + ).toBe( + createHash("sha256") + .update( + payload( + Number(range.rowId), + Number(range.offset), + Number(range.length), + 7, + ), + ) + .digest("hex"), + ); + } + expect(calls).toHaveLength(2); + expect(calls.map((body) => body.version)).toEqual([null, 7]); + }); + }); + + it.each([ + "type", + "column", + "count", + "length", + "eos", + "body", + "transport", + "version", + "disconnect", + 401, + 403, + 404, + 429, + 500, + ])("rejects the whole call after a later %s failure", async (fault) => { + await withRangeServer(async (db, calls, paths) => { + const table = await open(db); + await expect( + table.fetchBlobRanges( + "image", + Array.from({ length: 1025 }, () => ({ + rowId: 10n, + offset: 0n, + length: 1n, + })), + ), + ).rejects.toThrow(); + expect(calls).toHaveLength(2); + expect( + paths.every( + (path) => + path.endsWith("describe/") || + path.endsWith("branches/list/") || + path.endsWith("fetch_blob_ranges/"), + ), + ).toBe(true); + }, fault); + }); + + it.each(["missing-identity", "wrong-identity", "recreated", "gone"])( + "retains original identity through row and byte partitions (%s)", + async (fault) => { + for (const huge of [false, true]) { + await withRangeServer(async (db, calls) => { + const table = await open(db); + const requests = huge + ? [{ rowId: 10n, offset: 0n, length: BigInt(64 * MIB + 1) }] + : Array.from({ length: 1025 }, () => ({ + rowId: 10n, + offset: 0n, + length: 1n, + })); + await expect( + table.fetchBlobRanges("image", requests), + ).rejects.toThrow(); + expect(calls.map((call) => call.identity)).toEqual([ + undefined, + IDENTITY, + ]); + expect(calls.map((call) => call.version)).toEqual([null, 7]); + }, fault); + } + }, + ); + + it.each(["first-missing-identity", "first-invalid-identity"])( + "requires a first authoritative bulk identity (%s)", + async (fault) => { + await withRangeServer(async (db, calls, paths) => { + const table = await open(db); + const before = paths.length; + await expect( + table.fetchBlobRanges( + "image", + Array.from({ length: 1025 }, () => ({ + rowId: 10n, + offset: 0n, + length: 1n, + })), + ), + ).rejects.toThrow(/snapshot identity/); + expect(calls).toHaveLength(1); + expect(paths).toHaveLength(before + 1); + }, fault); + }, + ); + + it("validates integers, overflow and cancellation and leaves empty calls at zero I/O", async () => { + await withRangeServer(async (db, calls, paths) => { + const table = await open(db); + const before = paths.length; + expect(await table.fetchBlobRanges("image", [])).toEqual([]); + for (const request of [ + { rowId: -1, offset: 0n, length: 1n }, + { rowId: 1.5, offset: 0n, length: 1n }, + { rowId: Number.MAX_SAFE_INTEGER + 1, offset: 0n, length: 1n }, + { rowId: -1n, offset: 0n, length: 1n }, + { rowId: 1n << 64n, offset: 0n, length: 1n }, + { rowId: 10n, offset: -1n, length: 1n }, + { rowId: 10n, offset: 0n, length: -1n }, + { rowId: 10n, offset: 0n, length: 1n << 64n }, + { rowId: 20n, offset: (1n << 64n) - 1n, length: 1n }, + ]) { + await expect( + table.fetchBlobRanges("image", [request]), + ).rejects.toThrow(); + } + // Runtime type validation must also reject a number passed as an offset. + await expect( + table.fetchBlobRanges("image", [ + { rowId: 10n, offset: 0 as unknown as bigint, length: 1n }, + ]), + ).rejects.toThrow(); + const controller = new AbortController(); + controller.abort(new Error("cancelled by caller")); + await expect( + table.fetchBlobRanges( + "image", + [{ rowId: 10n, offset: 0n, length: 1n }], + { signal: controller.signal }, + ), + ).rejects.toThrow("cancelled by caller"); + expect(paths).toHaveLength(before); + expect(calls).toHaveLength(0); + }); + }); + + it("records the loaded native artifact", async () => { + const artifacts = Object.keys(require.cache).filter( + (path) => path.endsWith(".node") && path.includes("lancedb"), + ); + expect(artifacts.length).toBeGreaterThan(0); + for (const path of artifacts) { + const hash = createHash("sha256"); + for await (const chunk of createReadStream(path)) hash.update(chunk); + console.log( + `loaded native artifact=${path} sha256=${hash.digest("hex")}`, + ); + } + }); +}); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index b803e9043..59d6929b4 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3019,6 +3019,43 @@ describe("when dealing with blob columns", () => { expect(bytes.map((b) => b?.toString())).toEqual(["new", "hi"]); }); + it("fetchBlobRanges reads local ranges with order, duplicates, nulls and validation", async () => { + const { table, rowIds } = await openBlobTable(); + const [alphaId, betaId, nullId] = rowIds; + const requests = [ + { rowId: betaId, offset: 1n, length: 3n }, + { rowId: alphaId, offset: 0n, length: 2n }, + { rowId: nullId, offset: 0n, length: 1n }, + { rowId: alphaId, offset: 5n, length: 0n }, + { rowId: betaId, offset: 1n, length: 3n }, + ]; + expect(await table.fetchBlobRanges("image", requests)).toEqual([ + Buffer.from("eta"), + Buffer.from("al"), + null, + Buffer.alloc(0), + Buffer.from("eta"), + ]); + expect(await table.fetchBlobRanges("image", [])).toEqual([]); + await expect( + table.fetchBlobRanges("image", [ + { rowId: alphaId, offset: 4n, length: 2n }, + ]), + ).rejects.toThrow(); + await expect( + table.fetchBlobRanges("image", [ + { rowId: (1n << 63n) - 1n, offset: 0n, length: 1n }, + ]), + ).rejects.toThrow(); + const controller = new AbortController(); + expect( + await table.fetchBlobRanges("image", requests, { + signal: controller.signal, + }), + ).toHaveLength(5); + expect(getEventListeners(controller.signal, "abort")).toHaveLength(0); + }); + it("discovers blob columns", async () => { const { table } = await openBlobTable(); expect(await table.blobColumns()).toEqual(["image"]); diff --git a/nodejs/lancedb/blob.ts b/nodejs/lancedb/blob.ts index 67e662175..a46484314 100644 --- a/nodejs/lancedb/blob.ts +++ b/nodejs/lancedb/blob.ts @@ -305,6 +305,16 @@ export async function runWithSignal( } } +/** One row-specific range for {@link Table.fetchBlobRanges}. */ +export type BlobRangeRequest = { + /** A row ID from a query with row IDs enabled. */ + rowId: bigint | number; + /** Byte offset from the beginning of the logical blob, as a uint64. */ + offset: bigint; + /** Number of bytes to read, as a uint64. Zero returns an empty non-null value. */ + length: bigint; +}; + /** A half-open byte range `[start, end)` for {@link BlobFile.readRanges}. */ export type BlobRange = { start: bigint; diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 35697dc8c..08a254bdb 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -91,6 +91,7 @@ export type { BlobInput, BlobOptions, BlobRange, + BlobRangeRequest, BlobReadOptions, BlobUri, } from "./blob"; diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 185088424..136eb1d6c 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -17,7 +17,12 @@ import { tableFromIPC, } from "./arrow"; -import { BlobFile, BlobReadOptions, runWithSignal } from "./blob"; +import { + BlobFile, + BlobRangeRequest, + BlobReadOptions, + runWithSignal, +} from "./blob"; import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry"; import { IndexOptions } from "./indices"; import { Job } from "./job"; @@ -683,6 +688,31 @@ export abstract class Table { options?: BlobReadOptions, ): Promise<(Buffer | null)[]>; + /** + * Materializes row-specific byte ranges at the table's current checkout. + * Preserves request order, duplicates, and nulls. Zero length returns an empty + * buffer for a non-null blob. Invalid ranges or missing rows reject the entire + * call. See {@link Table.fetchBlobs} for row-ID validity across versions. + * + * Remote tables use the batched endpoint when the server advertises support, + * splitting at 1024 ranges and 64 MiB, including individual larger ranges. + * Older servers use descriptor queries and BlobFile range reads. Both protocols + * use one exact branch and version. Errors from an advertised endpoint propagate + * without switching protocols. All output bytes occupy memory. + * + * @example + * ```ts + * const bytes = await table.fetchBlobRanges("image", [ + * { rowId: 10n, offset: 0n, length: 1024n }, + * ]); + * ``` + */ + abstract fetchBlobRanges( + column: string, + requests: readonly BlobRangeRequest[], + options?: BlobReadOptions, + ): Promise<(Buffer | null)[]>; + /** * Opens lazy blob handles for `column` at the given row IDs using the * table's current checkout. @@ -1510,6 +1540,23 @@ export class LocalTable extends Table { return values.map((value) => value ?? null); } + async fetchBlobRanges( + column: string, + requests: readonly BlobRangeRequest[], + options?: BlobReadOptions, + ): Promise<(Buffer | null)[]> { + const rowIds = rowIdsToBigInts(requests.map((request) => request.rowId)); + const ranges = requests.map((request, i) => ({ + rowId: rowIds[i], + offset: request.offset, + length: request.length, + })); + const values = await runWithSignal(options?.signal, (signal) => + this.inner.fetchBlobRanges(column, ranges, signal), + ); + return values.map((value) => value ?? null); + } + async fetchBlobFiles( column: string, rowIds: readonly (bigint | number)[], diff --git a/nodejs/src/blob.rs b/nodejs/src/blob.rs index 59cf61af6..b55154019 100644 --- a/nodejs/src/blob.rs +++ b/nodejs/src/blob.rs @@ -21,6 +21,25 @@ pub struct BlobRange { pub end: BigInt, } +#[napi(object)] +pub struct BlobRangeRequest { + pub row_id: BigInt, + pub offset: BigInt, + pub length: BigInt, +} + +impl BlobRangeRequest { + pub(crate) fn parse(self) -> napi::Result { + let row_id = parse_u64(self.row_id, "row id")?; + let offset = parse_u64(self.offset, "offset")?; + let length = parse_u64(self.length, "length")?; + offset + .checked_add(length) + .ok_or_else(|| napi::Error::from_reason("blob range offset + length overflows u64"))?; + Ok(lancedb::blob::BlobRangeRequest::new(row_id, offset, length)) + } +} + #[napi] pub struct BlobFile { inner: Arc, @@ -205,14 +224,25 @@ pub fn parse_row_ids(row_ids: Vec) -> napi::Result> { .collect() } -pub fn copy_blob_buffers(array: LargeBinaryArray) -> Vec> { - (0..array.len()) - .map(|i| { - if array.is_null(i) { - None - } else { - Some(Buffer::from(array.value(i).to_vec())) - } - }) - .collect() +pub fn copy_blob_buffers(array: LargeBinaryArray) -> napi::Result>> { + let allocation_error = + |error| napi::Error::from_reason(format!("cannot allocate blob result: {error}")); + let mut result = Vec::new(); + result + .try_reserve_exact(array.len()) + .map_err(allocation_error)?; + for value in array.iter() { + let buffer = value + .map(|value| { + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(value.len()) + .map_err(allocation_error)?; + bytes.extend_from_slice(value); + Ok::<_, napi::Error>(Buffer::from(bytes)) + }) + .transpose()?; + result.push(buffer); + } + Ok(result) } diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 7076bbec1..f6abeb553 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -17,7 +17,7 @@ use napi::bindgen_prelude::*; use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode}; use napi_derive::napi; -use crate::blob::{BlobFile, copy_blob_buffers, parse_row_ids, spawn_abortable}; +use crate::blob::{BlobFile, BlobRangeRequest, copy_blob_buffers, parse_row_ids, spawn_abortable}; use crate::error::NapiErrorExt; use crate::index::Index; use crate::job::JobFailureInfo; @@ -420,7 +420,32 @@ impl Table { .fetch_blobs(column.as_str(), &row_ids) .await .default_error()?; - Ok(copy_blob_buffers(array)) + copy_blob_buffers(array) + }) + } + + #[napi(catch_unwind)] + pub fn fetch_blob_ranges<'env>( + &self, + env: &'env Env, + column: String, + requests: Vec, + signal: Option, + ) -> napi::Result>>> { + let mut ranges = Vec::new(); + ranges.try_reserve_exact(requests.len()).map_err(|error| { + napi::Error::from_reason(format!("cannot allocate blob range requests: {error}")) + })?; + for request in requests { + ranges.push(request.parse()?); + } + let table = self.inner_ref()?.clone(); + spawn_abortable(env, signal, async move { + let array = table + .fetch_blob_ranges(column.as_str(), ranges) + .await + .default_error()?; + copy_blob_buffers(array) }) } diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 3161eda30..5f0af7a66 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2661,8 +2661,12 @@ class Table(ABC): results are returned. Row IDs can be obtained from a query with ``with_row_id(True)``. - LanceDB Enterprise reads the requested ranges through version-pinned blob - handles using HTTP byte-range requests. + LanceDB Enterprise uses the batched endpoint when the server advertises + support, partitioning at 1024 ranges and 64 MiB, including individual + larger ranges. Older servers use descriptor queries and BlobFile range + reads. Both protocols read one exact branch and version without changing + the table's checkout. Errors from an advertised endpoint propagate without + switching protocols. The final output occupies memory. """ @abstractmethod @@ -8652,6 +8656,10 @@ class AsyncTable: column: str, requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: + """Materialize row-specific byte ranges. + + See [Table.fetch_blob_ranges][lancedb.table.Table.fetch_blob_ranges]. + """ return await self._inner.fetch_blob_ranges(column, list(requests)) async def fetch_blob_files( diff --git a/python/python/tests/test_blob_ranges.py b/python/python/tests/test_blob_ranges.py new file mode 100644 index 000000000..c65dfb8ae --- /dev/null +++ b/python/python/tests/test_blob_ranges.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Exercise the public sync/async range APIs over bounded loopback HTTP.""" + +import asyncio +import hashlib +import json +from pathlib import Path +import socket +import threading + +import lancedb._lancedb as native +import pyarrow as pa +import pytest + +from lancedb.remote.errors import HttpError, TableNotFoundError +from test_remote_db import mock_lancedb_connection + +MIB = 1024 * 1024 +IDENTITY = "v1:" + "a" * 64 +ROUTE = "/v1/table/test/fetch_blob_ranges/" + + +def payload(row, offset, length, version): + pattern = bytes((row * 17 + i + version) % 251 for i in range(251)) + start = offset % 251 + return (pattern * ((start + length + 250) // 251))[start : start + length] + + +class RangeServer: + def __init__(self, fault=None): + self.fault = fault + self.calls = [] + self.fetches = [] + self.identities = [] + self.stalled = threading.Event() + self.release = threading.Event() + + def __call__(self, request): + request.connection.settimeout(15) + size = int(request.headers.get("Content-Length", 0)) + assert size <= 256 * 1024 + body = json.loads(request.rfile.read(size) or b"{}") + self.calls.append((request.command, request.path, body)) + if request.path in ( + "/v1/table/test/describe/", + "/v1/table/test/branches/list/", + ): + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.send_header("phalanx-version", "1.13.0-beta.0") + if request.path.endswith("/describe/"): + request.send_header("x-lancedb-blob-ranges", "1") + request.end_headers() + response = ( + { + "branches": { + "experiment": { + "parentBranch": None, + "parentVersion": 1, + "createAt": 1, + "manifestSize": 1, + } + } + } + if request.path.endswith("branches/list/") + else {"version": 8, "schema": {"fields": []}} + ) + request.wfile.write(json.dumps(response).encode()) + return + assert (request.command, request.path) == ("POST", ROUTE) + self.fetches.append(body) + self.identities.append(request.headers.get("x-lancedb-query-snapshot")) + assert body["column"] == "image" + assert body["branch"] == "experiment" + ranges = body["ranges"] + assert len(ranges) <= 1024 + assert sum(r["length"] for r in ranges) <= 64 * MIB + version = body["version"] or (7 if len(self.fetches) == 1 else 8) + values = [ + None + if r["row_id"] == 20 + else payload(r["row_id"], r["offset"], r["length"], version) + for r in ranges + ] + fault = ( + self.fault + if len(self.fetches) > 1 or str(self.fault).startswith("first-") + else None + ) + if fault in ("recreated", "gone"): + request.send_response(400 if fault == "recreated" else 404) + request.end_headers() + request.wfile.write( + b"query snapshot identity mismatch: source gone or replaced" + ) + return + if fault == "stall" and len(self.fetches) == 2: + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.stream") + request.send_header("x-lancedb-version", str(version)) + request.send_header("x-lancedb-query-snapshot", IDENTITY) + request.end_headers() + request.wfile.flush() + self.stalled.set() + assert self.release.wait(15), "test did not release stalled response" + return + if fault == "structured400": + request.send_response(400) + request.send_header("x-lancedb-error-code", "invalid_input") + request.end_headers() + request.wfile.write(b'{"error":"range exceeds blob size"}') + return + if isinstance(fault, int): + request.send_response(fault) + request.end_headers() + request.wfile.write(b"ordinary server failure") + return + if fault == "disconnect": + request.connection.shutdown(socket.SHUT_RDWR) + return + column = "wrong" if fault == "column" else "value" + data_type = pa.binary() if fault == "type" else pa.large_binary() + if fault == "count": + values = [] + elif fault == "length": + values = [b""] + batch = pa.table({column: pa.array(values, type=data_type)}) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, batch.schema) as writer: + writer.write_table(batch, max_chunksize=257) + ipc = sink.getvalue().to_pybytes() + if fault == "eos": + ipc = ipc[:-8] + elif fault == "body": + ipc = ipc[: len(ipc) // 2] + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.stream") + if fault != "version": + request.send_header("x-lancedb-version", str(version)) + if fault not in ("missing-identity", "first-missing-identity"): + identity = ( + "broken" + if fault == "first-invalid-identity" + else "v1:" + "b" * 64 + if fault == "wrong-identity" + else IDENTITY + ) + request.send_header("x-lancedb-query-snapshot", identity) + request.send_header( + "Content-Length", str(len(ipc) + (10 if fault == "transport" else 0)) + ) + request.end_headers() + for i in range(0, len(ipc), 8192): + request.wfile.write(ipc[i : i + 8192]) + + +async def fetch(table, requests, use_async): + result = table.fetch_blob_ranges("image", requests) + return await result if use_async else result + + +def api(db, use_async, version=None): + table = db.open_table("test", branch="experiment", version=version) + return table._table if use_async else table + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("version", [None, 3]) +async def test_range_sdk_order_snapshot_and_overwrite(use_async, version): + handler = RangeServer() + with mock_lancedb_connection(handler) as db: + table = api(db, use_async, version) + requests = [(20 if i % 9 == 0 else 10 + i % 7, i, i % 4) for i in range(1025)] + requests += [(10, 0, 0), (10, 2, 3), (10, 2, 3)] + result = await fetch(table, requests, use_async) + assert isinstance(result, pa.LargeBinaryArray) + assert result.to_pylist() == [ + None if row == 20 else payload(row, offset, length, version or 7) + for row, offset, length in requests + ] + assert [b["version"] for b in handler.fetches] == [version, version or 7] + result = await fetch(table, [(10, 0, 1)], use_async) + assert result.to_pylist() == [payload(10, 0, 1, version or 8)] + assert handler.fetches[-1]["version"] == version + assert len(handler.fetches) == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_range_sdk_cached_reopen_negotiates_capability(use_async): + handler = RangeServer() + with mock_lancedb_connection(handler) as db: + for version in (7, 8): + table = api(db, use_async) + result = await fetch(table, [(10, 0, 1)], use_async) + assert result.to_pylist() == [payload(10, 0, 1, version)] + assert len(handler.fetches) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_range_sdk_large_single_and_aggregate_budget(use_async): + handler = RangeServer() + with mock_lancedb_connection(handler) as db: + table = api(db, use_async) + requests = [ + (10, 3, 33 * MIB), + (30, 7, 33 * MIB), + (40, 11, 70 * MIB), + (20, 0, 2**64 - 1), + ] + # The loopback server advances latest after the first partition. + result = await fetch(table, requests, use_async) + assert len(result) == 4 + for i, (row, offset, length) in enumerate(requests[:3]): + expected = payload(row, offset, length, 7) + assert ( + hashlib.sha256(result[i].as_buffer()).digest() + == hashlib.sha256(expected).digest() + ) + assert len(result[i].as_buffer()) == length + assert result[3].as_py() is None + assert len(handler.fetches) == 3 + assert [b["version"] for b in handler.fetches] == [None, 7, 7] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + "fault", + [ + "type", + "column", + "count", + "length", + "eos", + "body", + "transport", + "version", + "disconnect", + 401, + 403, + 404, + 429, + 500, + ], +) +async def test_range_sdk_later_failure_rejects_entire_call(use_async, fault): + handler = RangeServer(fault=fault) + with mock_lancedb_connection( + handler, retry_config={"retries": 2, "connect_retries": 2, "read_retries": 2} + ) as db: + table = api(db, use_async) + with pytest.raises(Exception) as exc: + await fetch(table, [(10, 0, 1)] * 1025, use_async) + if fault == 404: + assert isinstance(exc.value, TableNotFoundError) + assert isinstance(exc.value.__cause__, HttpError) + assert exc.value.__cause__.status_code == 404 + else: + assert not isinstance(exc.value, ValueError) + if isinstance(fault, int): + assert exc.value.status_code == fault + assert len(handler.fetches) == 2 + assert all( + path in ("/v1/table/test/describe/", "/v1/table/test/branches/list/", ROUTE) + for _, path, _ in handler.calls + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_range_sdk_structured_400_preserves_value_error(use_async): + handler = RangeServer(fault="structured400") + with mock_lancedb_connection(handler) as db: + table = api(db, use_async) + with pytest.raises(ValueError, match="range exceeds blob size"): + await fetch(table, [(10, 0, 1)] * 1025, use_async) + assert len(handler.fetches) == 2 + + +@pytest.mark.asyncio +async def test_range_sdk_async_cancel_releases_only_request_permit(): + handler = RangeServer(fault="stall") + with mock_lancedb_connection(handler, blob_request_concurrency=1) as db: + table = api(db, True) + pending = asyncio.create_task(fetch(table, [(10, 0, 1)] * 2049, True)) + try: + assert await asyncio.to_thread(handler.stalled.wait, 10) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert len(handler.fetches) == 2 + result = await asyncio.wait_for(fetch(table, [(10, 0, 1)], True), 5) + assert result.to_pylist() == [payload(10, 0, 1, 8)] + assert len(handler.fetches) == 3 + finally: + pending.cancel() + handler.release.set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_range_sdk_integer_validation_and_empty_zero_io(use_async): + handler = RangeServer() + with mock_lancedb_connection(handler) as db: + table = api(db, use_async) + before = len(handler.calls) + assert (await fetch(table, [], use_async)).to_pylist() == [] + for request in [ + (-1, 0, 1), + (2**64, 0, 1), + (10, -1, 1), + (10, 0, -1), + (10, 0, 2**64), + (10, 0.5, 1), + (20, 2**64 - 1, 1), + (10, 0), + ]: + with pytest.raises((ValueError, OverflowError, TypeError)): + await fetch(table, [request], use_async) + assert len(handler.calls) == before + + +def test_range_sdk_loaded_artifact_identity(): + artifact = Path(native.__file__).resolve() + digest = hashlib.sha256() + with artifact.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + print(f"loaded native artifact={artifact} sha256={digest.hexdigest()}") + assert artifact.is_file() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + "fault", ["missing-identity", "wrong-identity", "recreated", "gone"] +) +@pytest.mark.parametrize("huge", [False, True]) +async def test_range_sdk_original_identity_survives_every_partition( + use_async, fault, huge +): + handler = RangeServer(fault=fault) + with mock_lancedb_connection(handler) as db: + table = api(db, use_async) + requests = [(10, 0, 64 * MIB + 1)] if huge else [(10, 0, 1)] * 1025 + error = ( + ValueError + if fault == "recreated" + else TableNotFoundError + if fault == "gone" + else NotImplementedError + ) + with pytest.raises( + error, match="snapshot identity" if fault != "gone" else None + ): + await fetch(table, requests, use_async) + assert handler.identities == [None, IDENTITY] + assert len(handler.fetches) == 2 + assert [body["version"] for body in handler.fetches] == [None, 7] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("fault", ["first-missing-identity", "first-invalid-identity"]) +async def test_range_sdk_bulk_requires_identity_before_payload_acceptance( + use_async, fault +): + handler = RangeServer(fault=fault) + with mock_lancedb_connection(handler) as db: + table = api(db, use_async) + before = len(handler.calls) + with pytest.raises(NotImplementedError, match="snapshot identity"): + await fetch(table, [(10, 0, 1)] * 1025, use_async) + assert len(handler.fetches) == 1 + assert len(handler.calls) == before + 1 diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 561a0f900..4f707f041 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -5466,6 +5466,7 @@ def blob_remote_table( *, server_version=Version("0.5.0"), location=None, + range_capability=None, version_header="1", framing="file", requests=None, @@ -5487,6 +5488,8 @@ def blob_remote_table( request.send_response(200) request.send_header("Content-Type", "application/json") request.send_header("phalanx-version", str(server_version)) + if range_capability is not None: + request.send_header("x-lancedb-blob-ranges", range_capability) request.end_headers() response = dict(BLOB_DESCRIBE_RESPONSE) response["version"] = body.get("version") or state["version"] @@ -5594,6 +5597,42 @@ def blob_remote_table( writer_factory = pa.ipc.new_file if framing == "file" else pa.ipc.new_stream with writer_factory(request.wfile, response_table.schema) as writer: writer.write_table(response_table, max_chunksize=1) + elif ( + range_capability == "1" + and request.path == "/v1/table/test/fetch_blob_ranges/" + ): + size = int(request.headers["Content-Length"]) + body = json.loads(request.rfile.read(size)) + payloads = {10: b"alpha", 20: None, 30: b"gamma", 40: b""} + values = [] + for item in body["ranges"]: + row, offset, length = ( + item[key] for key in ("row_id", "offset", "length") + ) + error = None + if row not in payloads: + error = "only 1 exist" + elif payloads[row] is not None and offset + length > len(payloads[row]): + error = "range exceeds blob size" + if error: + request.send_response(400) + request.end_headers() + request.wfile.write(error.encode()) + return + payload = payloads[row] + values.append( + None if payload is None else payload[offset : offset + length] + ) + response_table = pa.table( + {"value": pa.array(values, type=pa.large_binary())} + ) + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.stream") + request.send_header("x-lancedb-version", "1") + request.send_header("x-lancedb-query-snapshot", "v1:" + "a" * 64) + request.end_headers() + with pa.ipc.new_stream(request.wfile, response_table.schema) as writer: + writer.write_table(response_table) elif request.path == "/v1/table/test/fetch_blobs/": content_len = int(request.headers.get("Content-Length", 0)) body = json.loads(request.rfile.read(content_len)) @@ -5901,8 +5940,11 @@ def test_sync_remote_blob_pandas_avoids_storage_access( @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True]) -async def test_remote_fetch_blob_ranges(use_async): - with blob_remote_table() as remote_table: +@pytest.mark.parametrize("range_capability", [None, "0", "2", "1"]) +async def test_remote_fetch_blob_ranges(use_async, range_capability): + with blob_remote_table( + server_version="1.13.0-beta.0", range_capability=range_capability + ) as remote_table: table = remote_table._table if use_async else remote_table requests = [ (30, 2, 3), diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index ead1f2410..82f502250 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -356,7 +356,7 @@ impl BlobFile { /// /// `row_id` is obtained from a query with row ids enabled. /// `offset` and `length` are relative to the beginning of the logical blob. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] pub struct BlobRangeRequest { /// Row id of the blob value to read. pub row_id: u64, diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index b562f41df..578e3bbab 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1655,9 +1655,9 @@ impl Database for RemoteDatabase { let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?; let rsp = RemoteTable::::handle_table_not_found(&request.name, rsp, &request_id).await?; - self.client.check_response(&request_id, rsp).await?; + let rsp = self.client.check_response(&request_id, rsp).await?; - Ok(Arc::new(RemoteTable::new_with_sql_client( + let table = Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), self.db_name.clone(), request.name, @@ -1665,7 +1665,9 @@ impl Database for RemoteDatabase { identifier, version, self.sql_client.clone(), - ))) + )); + table.seed_blob_ranges_capability(rsp.headers()); + Ok(table) } else { // Describe the table to confirm it exists before moving on. let req = self @@ -1676,6 +1678,7 @@ impl Database for RemoteDatabase { RemoteTable::::handle_table_not_found(&request.name, rsp, &request_id).await?; let rsp = self.client.check_response(&request_id, rsp).await?; let version = parse_server_version(&request_id, &rsp)?; + let describe_headers = rsp.headers().clone(); let describe_body = rsp.text().await.ok(); let table = Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), @@ -1686,6 +1689,7 @@ impl Database for RemoteDatabase { version.clone(), self.sql_client.clone(), )); + table.seed_blob_ranges_capability(&describe_headers); // This describe already carries the schema, so hand it to the table // instead of making the first schema read fetch it again. A version or // branch pin applied after this invalidates the cache. diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index d560a4802..fdb547afa 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -86,6 +86,7 @@ use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; use std::io::Cursor; use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; use tokio::sync::RwLock; @@ -634,6 +635,8 @@ pub struct RemoteTable { version: Arc>>, location: RwLock>, schema_cache: BackgroundCache, + // Only a successful describe can opt into the batched range protocol. + blob_ranges_supported: Arc, // WAL-PK-FUSION: delete, with its initializers below. wal_pk_fusion: PkFusionMemory, freshness: Arc>, @@ -881,6 +884,7 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + blob_ranges_supported: Arc::new(AtomicBool::new(false)), wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -888,6 +892,15 @@ impl RemoteTable { } } + pub(super) fn seed_blob_ranges_capability(&self, headers: &http::HeaderMap) { + self.blob_ranges_supported.store( + headers + .get("x-lancedb-blob-ranges") + .is_some_and(|value| value == "1"), + AtomicOrdering::Relaxed, + ); + } + /// Seed the schema cache from a `describe` body the caller already fetched. /// /// Best effort. `open_table` succeeds today without reading this body, so a @@ -918,6 +931,7 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + blob_ranges_supported: self.blob_ranges_supported.clone(), wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch, @@ -1067,6 +1081,7 @@ impl RemoteTable { }; let response = self.check_table_response(&request_id, response).await?; + self.seed_blob_ranges_capability(response.headers()); match response.text().await { Ok(body) => { @@ -1943,6 +1958,7 @@ async fn fetch_schema( read_snapshot: ReadSnapshot, branch: Option, freshness: Arc>, + blob_ranges_supported: Arc, ) -> Result { let mut body = serde_json::json!({ "version": read_snapshot.version }); if let Some(branch) = &branch { @@ -1969,6 +1985,13 @@ async fn fetch_schema( let response = client.check_response(&request_id, response).await?; freshness_headers.observe_headers(&freshness, response.headers()); + blob_ranges_supported.store( + response + .headers() + .get("x-lancedb-blob-ranges") + .is_some_and(|value| value == "1"), + AtomicOrdering::Relaxed, + ); let body = response.text().await.map_err(|e| { let status_code = e.status(); Error::Http { @@ -2026,6 +2049,7 @@ mod test_utils { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + blob_ranges_supported: Arc::new(AtomicBool::new(false)), wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -2054,6 +2078,7 @@ mod test_utils { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + blob_ranges_supported: Arc::new(AtomicBool::new(false)), wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -2091,6 +2116,7 @@ mod test_utils { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + blob_ranges_supported: Arc::new(AtomicBool::new(false)), wal_pk_fusion: PkFusionMemory::default(), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -2885,6 +2911,7 @@ impl BaseTable for RemoteTable { let table_name = self.name.clone(); let branch = self.branch.clone(); let freshness = self.freshness.clone(); + let blob_ranges_supported = self.blob_ranges_supported.clone(); match self .schema_cache @@ -2896,6 +2923,7 @@ impl BaseTable for RemoteTable { read_snapshot, branch, freshness, + blob_ranges_supported, ) .await }) diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index a0051c955..3367d2218 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -35,11 +35,13 @@ use crate::table::{AnyQuery, BaseTable}; use super::{ FreshnessHeaders, FreshnessState, QUERY_SNAPSHOT_HEADER, ReadSnapshot, RemoteTable, VERSION_HEADER, check_query_snapshot_response, freshness_headers_snapshot, + valid_query_snapshot_identity, }; // The Cloud route rejects larger row-id lists. Its separate 64 MiB byte limit // is handled by splitting a rejected request below. const MAX_FETCH_BLOBS_ROW_IDS: usize = 1024; +const MAX_FETCH_BLOB_RANGE_BYTES: u64 = 64 * 1024 * 1024; fn is_fetch_blobs_byte_limit_error(error: &Error) -> bool { matches!(error, Error::Http { @@ -599,7 +601,7 @@ impl RemoteTable { read_snapshot.version = Some(self.describe_read_snapshot(read_snapshot).await?.version); } let mut pending: Vec<&[u64]> = row_ids.chunks(MAX_FETCH_BLOBS_ROW_IDS).rev().collect(); - let mut blobs = BlobAccumulator::new(row_ids.len()); + let mut blobs = BlobAccumulator::new(row_ids.len())?; while let Some(ids) = pending.pop() { match self .fetch_blobs_chunk(column, ids, read_snapshot, &mut blobs) @@ -642,7 +644,7 @@ impl RemoteTable { } } - Ok(blobs.finish()) + Ok(blobs.finish()?) } async fn fetch_blobs_chunk( @@ -725,11 +727,88 @@ impl RemoteTable { Ok(()) } - /// Materialize ranges through version-pinned handles, preserving nulls and order. + /// Materialize ranges using the protocol advertised by a successful describe. pub(super) async fn fetch_blob_ranges_impl( &self, column: &str, requests: &[BlobRangeRequest], + ) -> Result { + let mut blobs = BlobAccumulator::new(requests.len())?; + if requests.is_empty() { + return Ok(blobs.finish()?); + } + for request in requests { + request + .offset + .checked_add(request.length) + .ok_or_else(|| Error::InvalidInput { + message: "blob range offset + length overflows u64".into(), + })?; + } + self.reject_list_blob_path_if_known(column).await?; + if !self.blob_ranges_supported.load(Ordering::Relaxed) { + return self.fetch_blob_ranges_legacy(column, requests).await; + } + // Select the protocol before issuing data requests. An advertised endpoint + // failing in a mixed deployment must fail truthfully, without fallback. + let mut snapshot = self.snapshot_read_state().await; + let mut identity = self.expected_query_snapshot.clone(); + let mut index = 0; + let mut consumed = 0; + let mut ranges = Vec::new(); + ranges + .try_reserve_exact(MAX_FETCH_BLOBS_ROW_IDS) + .map_err(|error| Error::InvalidInput { + message: format!("cannot allocate blob range requests: {error}"), + })?; + while index < requests.len() { + ranges.clear(); + let continuation = consumed != 0; + let mut budget = MAX_FETCH_BLOB_RANGE_BYTES; + while index < requests.len() && ranges.len() < MAX_FETCH_BLOBS_ROW_IDS { + let request = requests[index]; + let remaining = request.length - consumed; + if budget == 0 && remaining != 0 { + break; + } + let length = remaining.min(budget); + // The full range was checked above; all subranges fit within it. + ranges.push(BlobRangeRequest::new( + request.row_id, + request.offset + consumed, + length, + )); + consumed += length; + budget -= length; + if consumed == request.length { + index += 1; + consumed = 0; + } else { + break; + } + } + let array = self + .fetch_blob_ranges_chunk(column, &ranges, &mut snapshot, &mut identity) + .await?; + // A null has no bytes to partition. Avoid sending arbitrarily many + // empty fragments for a null with a very large requested length. + if consumed != 0 && array.is_null(array.len() - 1) { + index += 1; + consumed = 0; + } + if continuation { + blobs.append_continuation(&array)?; + } else { + blobs.append(&array)?; + } + } + Ok(blobs.finish()?) + } + + async fn fetch_blob_ranges_legacy( + &self, + column: &str, + requests: &[BlobRangeRequest], ) -> Result { let row_ids = requests .iter() @@ -738,16 +817,12 @@ impl RemoteTable { let files = self.fetch_blob_files_impl(column, &row_ids).await?; let mut builder = LargeBinaryBuilder::with_capacity(requests.len(), 0); for (file, request) in files.into_iter().zip(requests) { - let end = - request - .offset - .checked_add(request.length) - .ok_or_else(|| Error::InvalidInput { - message: "blob range offset + length overflows u64".into(), - })?; match file { Some(file) => { - let bytes = file.read_range(request.offset..end).await?; + // Overflow was checked before selecting either protocol. + let bytes = file + .read_range(request.offset..request.offset + request.length) + .await?; builder.append_value(bytes); } None => builder.append_null(), @@ -756,6 +831,162 @@ impl RemoteTable { Ok(builder.finish()) } + async fn fetch_blob_ranges_chunk( + &self, + column: &str, + ranges: &[BlobRangeRequest], + snapshot: &mut ReadSnapshot, + identity: &mut Option, + ) -> Result { + let _permit = self + .client + .blob_request_semaphore + .clone() + .acquire_owned() + .await + .expect("blob request semaphore is never closed"); + let body = serde_json::json!({ + "column": column, + "version": snapshot.version, + "branch": self.branch, + "ranges": ranges, + }); + let request = self + .client + .post(&format!("/v1/table/{}/fetch_blob_ranges/", self.identifier)) + .json(&body); + let request = if let Some(expected) = identity.as_deref() { + request.header(super::QUERY_SNAPSHOT_HEADER, expected) + } else { + request + }; + let (request_id, response) = self + .send_with_freshness(request, false, snapshot.freshness) + .await?; + // The range operation defines 400 as invalid input. Let the established + // checker decode structured/plain server errors before normalizing it. + let response = self + .check_table_response(&request_id, response) + .await + .map_err(|error| { + if let Error::Http { + source, + request_id, + status_code: Some(StatusCode::BAD_REQUEST), + } = error + { + Error::InvalidInput { + message: format!("fetch_blob_ranges (request {request_id}): {source}"), + } + } else { + error + } + })?; + let invalid = |message: String| Error::Http { + source: format!("fetch_blob_ranges: {message}").into(), + request_id: request_id.clone(), + status_code: None, + }; + let version = response.headers().get(&VERSION_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .filter(|version| *version != 0) + .ok_or_else(|| invalid("missing or invalid x-lancedb-version; a server supporting the batched range endpoint is required".into()))?; + if snapshot.version.is_some_and(|expected| expected != version) { + return Err(invalid(format!( + "response version {version} differs from requested version {:?}", + snapshot.version + ))); + } + // Do not decode or accept any payload before establishing/verifying the + // identity. A later partition can never replace the original expectation. + check_query_snapshot_response(identity.as_deref(), &response)?; + if identity.is_none() { + let actual = response.headers().get(super::QUERY_SNAPSHOT_HEADER) + .and_then(|value| value.to_str().ok()) + .filter(|value| valid_query_snapshot_identity(value)) + .ok_or_else(|| Error::NotSupported { + message: "Batched Blob ranges require a snapshot identity; upgrade every serving node".into(), + })?; + *identity = Some(actual.to_owned()); + } + snapshot.version = Some(version); + // This endpoint's contract is stream framing. Keep the bounded, fallible + // IPC path from #4979 instead of buffering an Arrow file response. + let content_type = response + .headers() + .get(header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or(""); + if !content_type + .split(';') + .next() + .unwrap_or("") + .trim() + .eq_ignore_ascii_case(crate::remote::ARROW_STREAM_CONTENT_TYPE) + { + return Err(invalid(format!( + "expected Arrow IPC stream, got '{content_type}'" + ))); + } + let mut stream = super::ipc_stream::read_ipc_stream( + response.bytes_stream(), + request_id.clone(), + Some(version), + ) + .await?; + let schema = stream.schema(); + if !matches!( + schema + .field_with_name("value") + .map(|field| field.data_type()), + Ok(DataType::LargeBinary) + ) { + return Err(invalid( + "response requires a LargeBinary column named 'value'".into(), + )); + } + let mut blobs = BlobAccumulator::new(ranges.len())?; + let mut rows = 0usize; + let mut bytes = 0u64; + while let Some(batch) = stream.try_next().await? { + let array = batch + .column_by_name("value") + .and_then(|column| column.as_any().downcast_ref::()) + .ok_or_else(|| { + invalid("response requires a LargeBinary column named 'value'".into()) + })?; + let end = rows + .checked_add(array.len()) + .filter(|end| *end <= ranges.len()) + .ok_or_else(|| invalid("response has more rows than requested".into()))?; + let offsets = array.value_offsets(); + bytes = bytes + .checked_add((offsets[array.len()] - offsets[0]) as u64) + .filter(|bytes| *bytes <= MAX_FETCH_BLOB_RANGE_BYTES) + .ok_or_else(|| invalid("response exceeds 64 MiB of range bytes".into()))?; + for (i, range) in ranges[rows..end].iter().enumerate() { + if array.is_valid(i) && array.value_length(i) as u64 != range.length { + return Err(invalid(format!( + "row {} returned {} bytes for requested length {}", + rows + i, + array.value_length(i), + range.length + ))); + } + } + blobs.append(array)?; + rows = end; + } + if rows != ranges.len() { + return Err(invalid(format!( + "returned {rows} rows for {} ranges", + ranges.len() + ))); + } + Ok(blobs.finish()?) + } + /// Open seekable handles for `row_ids`. /// /// Descriptors for the requested rows come from one `_rowid` take query per @@ -1666,10 +1897,155 @@ mod tests { let error = table.fetch_blob_ranges("image", &ranges).await.unwrap_err(); assert!(error.to_string().contains(message), "{error}"); // Missing rows and invalid ranges fail without reading any bytes. - assert_eq!(requests.lock().unwrap().len(), 1); + assert!(matches!(error, Error::InvalidInput { .. })); + assert_eq!( + requests.lock().unwrap().len(), + usize::from(!message.contains("overflows")) + ); } } + #[tokio::test] + async fn remote_blob_ranges_keep_producing_query_identity_from_first_request() { + const IDENTITY: &str = + "v1:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + let calls = Arc::new(StdMutex::new(Vec::new())); + let captured = calls.clone(); + let table = RemoteTable::new_mock( + "my_table".into(), + move |request| { + assert_eq!( + request.url().path(), + "/v1/table/my_table/fetch_blob_ranges/" + ); + assert_eq!( + request.headers()[super::super::QUERY_SNAPSHOT_HEADER], + IDENTITY + ); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body["version"], 42); + let rows = body["ranges"].as_array().unwrap().len(); + captured.lock().unwrap().push(rows); + let batch = RecordBatch::try_from_iter([( + "value", + Arc::new(LargeBinaryArray::from(vec![Some(b"A".as_slice()); rows])) + as Arc, + )]) + .unwrap(); + let mut bytes = Vec::new(); + let mut writer = + arrow_ipc::writer::StreamWriter::try_new(&mut bytes, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + http::Response::builder() + .status(200) + .header("content-type", "application/vnd.apache.arrow.stream") + .header(VERSION_HEADER, "42") + .header(super::super::QUERY_SNAPSHOT_HEADER, IDENTITY) + .body(bytes) + .unwrap() + }, + Some(Version::new(1, 13, 0)), + ); + let mut headers = http::HeaderMap::new(); + headers.insert("x-lancedb-blob-ranges", "1".parse().unwrap()); + table.seed_blob_ranges_capability(&headers); + let snapshot = table + .query_snapshot_with_identity(42, IDENTITY) + .await + .unwrap(); + let ranges = snapshot + .fetch_blob_ranges("image", &[BlobRangeRequest::new(10, 0, 1); 1025]) + .await + .unwrap(); + assert_eq!( + ranges, + LargeBinaryArray::from(vec![Some(b"A".as_slice()); 1025]) + ); + assert_eq!(*calls.lock().unwrap(), [1024, 1]); + assert!(table.snapshot_read_state().await.version.is_none()); + } + + #[tokio::test] + async fn remote_blob_ranges_cancel_releases_the_only_request_permit() { + use std::sync::atomic::AtomicUsize; + let entered = Arc::new(tokio::sync::Notify::new()); + let started = entered.clone(); + let calls = Arc::new(AtomicUsize::new(0)); + let captured = calls.clone(); + let table = Arc::new(RemoteTable::new_mock_with_version_and_config( + "my_table".into(), + move |request| { + assert!(request.url().path().ends_with("/fetch_blob_ranges/")); + if captured.fetch_add(1, Ordering::SeqCst) == 0 { + started.notify_one(); + http::Response::builder() + .status(200) + .header("content-type", "application/vnd.apache.arrow.stream") + .header(VERSION_HEADER, "7") + .header( + super::super::QUERY_SNAPSHOT_HEADER, + format!("v1:{}", "a".repeat(64)), + ) + .body(reqwest::Body::wrap_stream(futures::stream::pending::< + std::result::Result, + >())) + .unwrap() + } else { + http::Response::builder() + .status(400) + .header("x-lancedb-error-code", "invalid_input") + .body(reqwest::Body::from( + r#"{"error":"range exceeds blob size"}"#, + )) + .unwrap() + } + }, + Some(Version::new(1, 13, 0)), + crate::remote::ClientConfig { + blob_request_concurrency: Some(1), + ..Default::default() + }, + )); + let mut headers = http::HeaderMap::new(); + headers.insert("x-lancedb-blob-ranges", "1".parse().unwrap()); + table.seed_blob_ranges_capability(&headers); + let reader_table = table.clone(); + let reader = tokio::spawn(async move { + reader_table + .fetch_blob_ranges("image", &[BlobRangeRequest::new(10, 0, 1); 1025]) + .await + }); + tokio::time::timeout(Duration::from_secs(5), entered.notified()) + .await + .unwrap(); + reader.abort(); + assert!(reader.await.unwrap_err().is_cancelled()); + let error = tokio::time::timeout( + Duration::from_secs(5), + table.fetch_blob_ranges("image", &[BlobRangeRequest::new(10, 0, 1)]), + ) + .await + .unwrap() + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert!(error.to_string().contains("range exceeds blob size")); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn remote_blob_ranges_reject_overflow_before_http() { + let (table, requests) = mock_take_blob_table(vec![], Some("7")); + let error = table + .fetch_blob_ranges("image", &[BlobRangeRequest::new(20, u64::MAX, 1)]) + .await + .unwrap_err(); + assert!(error.to_string().contains("overflows u64"), "{error}"); + assert!(requests.lock().unwrap().is_empty()); + } + #[tokio::test] async fn remote_blob_ranges_accept_empty_requests_without_http() { let (table, requests) = mock_take_blob_table(vec![], Some("7")); diff --git a/rust/lancedb/src/remote/table/blobs/aggregate.rs b/rust/lancedb/src/remote/table/blobs/aggregate.rs index e8e794fa0..68bc20c88 100644 --- a/rust/lancedb/src/remote/table/blobs/aggregate.rs +++ b/rust/lancedb/src/remote/table/blobs/aggregate.rs @@ -6,75 +6,143 @@ use std::collections::VecDeque; use arrow_array::{Array, LargeBinaryArray}; -use arrow_buffer::{Buffer, MutableBuffer, NullBufferBuilder, OffsetBuffer}; +use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer}; use arrow_schema::ArrowError; +fn allocation_error(error: impl std::fmt::Display) -> ArrowError { + ArrowError::ComputeError(format!("cannot allocate blob result: {error}")) +} + pub(super) struct BlobAccumulator { offsets: Vec, - nulls: NullBufferBuilder, + validity: Vec, values: VecDeque, bytes: i64, + has_nulls: bool, } impl BlobAccumulator { - pub(super) fn new(rows: usize) -> Self { - let mut offsets = Vec::with_capacity(rows + 1); + pub(super) fn new(rows: usize) -> Result { + let capacity = rows + .checked_add(1) + .ok_or_else(|| allocation_error("row count overflow"))?; + let mut offsets = Vec::new(); + offsets + .try_reserve_exact(capacity) + .map_err(allocation_error)?; offsets.push(0); - Self { + let mut validity = Vec::new(); + let validity_bytes = rows / 8 + usize::from(!rows.is_multiple_of(8)); + validity + .try_reserve_exact(validity_bytes) + .map_err(allocation_error)?; + Ok(Self { offsets, - nulls: NullBufferBuilder::new(rows), + validity, values: VecDeque::new(), bytes: 0, - } + has_nulls: false, + }) } pub(super) fn append(&mut self, array: &LargeBinaryArray) -> Result<(), ArrowError> { + self.append_inner(array, false) + } + + /// The first value continues the preceding response's final logical range. + pub(super) fn append_continuation( + &mut self, + array: &LargeBinaryArray, + ) -> Result<(), ArrowError> { + self.append_inner(array, true) + } + + fn append_inner( + &mut self, + array: &LargeBinaryArray, + continuation: bool, + ) -> Result<(), ArrowError> { + let rows = self.offsets.len() - 1; + if continuation + && (rows == 0 + || array.is_empty() + || (self.validity[(rows - 1) / 8] & (1 << ((rows - 1) % 8)) != 0) + != array.is_valid(0)) + { + return Err(ArrowError::ComputeError( + "inconsistent nulls in split blob range".into(), + )); + } let offsets = array.value_offsets(); let start = offsets[0]; let length = offsets[array.len()] - start; let total = self.bytes.checked_add(length).ok_or_else(|| { - ArrowError::ComputeError("fetch_blobs result exceeds i64 byte offsets".into()) + ArrowError::ComputeError("blob result exceeds i64 byte offsets".into()) })?; - self.offsets.extend( - offsets[1..] - .iter() - .map(|offset| self.bytes + (*offset - start)), - ); - if let Some(nulls) = array.nulls() { - self.nulls.append_buffer(nulls); - } else { - self.nulls.append_n_non_nulls(array.len()); - } + let skip = usize::from(continuation); + self.offsets + .try_reserve(array.len() - skip) + .map_err(allocation_error)?; + let new_rows = rows + .checked_add(array.len() - skip) + .ok_or_else(|| allocation_error("row count overflow"))?; + let validity_bytes = new_rows / 8 + usize::from(new_rows % 8 != 0); + self.validity + .try_reserve(validity_bytes - self.validity.len()) + .map_err(allocation_error)?; + self.validity.resize(validity_bytes, 0); if length != 0 { + self.values.try_reserve(1).map_err(allocation_error)?; self.values.push_back( array .values() .slice_with_length(start as usize, length as usize), ); } + if continuation { + *self.offsets.last_mut().unwrap() = self.bytes + (offsets[1] - start); + } + self.offsets.extend( + offsets[1 + skip..] + .iter() + .map(|offset| self.bytes + (*offset - start)), + ); + for i in skip..array.len() { + let row = rows + i - skip; + if array.is_valid(i) { + self.validity[row / 8] |= 1 << (row % 8); + } else { + self.has_nulls = true; + } + } self.bytes = total; Ok(()) } - pub(super) fn finish(mut self) -> LargeBinaryArray { + pub(super) fn finish(mut self) -> Result { let values = match self.values.len() { 0 => Buffer::default(), 1 => self.values.pop_front().unwrap(), _ => { - let mut values = MutableBuffer::new(self.bytes as usize); - // Do not retain all decoded batches alongside the final allocation. + let length = usize::try_from(self.bytes).map_err(allocation_error)?; + let mut values = Vec::new(); + values.try_reserve_exact(length).map_err(allocation_error)?; // Drop each source allocation immediately after copying its bytes. while let Some(chunk) = self.values.pop_front() { values.extend_from_slice(&chunk); } - values.into() + Buffer::from_vec(values) } }; - LargeBinaryArray::new( + let rows = self.offsets.len() - 1; + let nulls = self + .has_nulls + .then(|| NullBuffer::new(BooleanBuffer::new(Buffer::from_vec(self.validity), 0, rows))); + Ok(LargeBinaryArray::new( OffsetBuffer::new(self.offsets.into()), values, - self.nulls.finish(), - ) + nulls, + )) } } @@ -85,9 +153,9 @@ mod tests { #[test] fn single_buffer_is_not_copied() { let array = LargeBinaryArray::from(vec![Some(b"alpha".as_slice()), None, Some(b"")]); - let mut result = BlobAccumulator::new(3); + let mut result = BlobAccumulator::new(3).unwrap(); result.append(&array).unwrap(); - let output = result.finish(); + let output = result.finish().unwrap(); assert_eq!(output.values().as_ptr(), array.values().as_ptr()); assert_eq!(output, array); } @@ -101,21 +169,27 @@ mod tests { Some(b""), Some(b"beta"), ]); - let mut result = BlobAccumulator::new(6); + let mut result = BlobAccumulator::new(6).unwrap(); result.append(&array.slice(1, 3)).unwrap(); result.append(&array.slice(4, 1)).unwrap(); result.append(&array.slice(2, 2)).unwrap(); assert_eq!( - result.finish(), + result.finish().unwrap(), LargeBinaryArray::from(vec![ Some(b"alpha".as_slice()), None, Some(b""), Some(b"beta"), None, - Some(b""), + Some(b"") ]) ); - assert_eq!(BlobAccumulator::new(0).finish().len(), 0); + assert_eq!(BlobAccumulator::new(0).unwrap().finish().unwrap().len(), 0); + } + + #[test] + fn impossible_capacity_is_an_error() { + assert!(BlobAccumulator::new(usize::MAX).is_err()); + assert!(BlobAccumulator::new(usize::MAX / 2).is_err()); } } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index db68ab290..4959b2842 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1452,8 +1452,13 @@ impl Table { /// ranges for the same blob. The output has the same length and order as /// the requests. Null blobs produce null output slots; empty ranges on /// non-null blobs produce empty byte strings. - /// LanceDB Enterprise tables read through version-pinned [`BlobFile`] handles using - /// HTTP byte-range requests. + /// LanceDB Enterprise uses the batched endpoint when the server advertises + /// support, splitting calls at 1024 ranges and 64 MiB of requested bytes. + /// Individual larger ranges are partitioned and reassembled into one result. + /// Older servers use descriptor queries and [`Self::fetch_blob_files`] range + /// reads. Both protocols read one exact branch and version without changing + /// this handle's checkout. The final materialized output still occupies memory. + /// Errors from an advertised endpoint propagate without switching protocols. /// /// `_rowid` values stay valid after compaction when the table has stable /// row ids. One missing or deleted row id rejects the whole request; no @@ -1486,10 +1491,14 @@ impl Table { column: impl AsRef, requests: impl IntoIterator, ) -> Result { - let requests = requests.into_iter().collect::>(); - self.inner - .fetch_blob_ranges(column.as_ref(), &requests) - .await + let mut ranges = Vec::new(); + for request in requests { + ranges.try_reserve(1).map_err(|error| Error::InvalidInput { + message: format!("cannot allocate blob range requests: {error}"), + })?; + ranges.push(request); + } + self.inner.fetch_blob_ranges(column.as_ref(), &ranges).await } /// Open lazy [`BlobFile`] handles for the given row ids. diff --git a/rust/lancedb/tests/remote_blob_ranges.rs b/rust/lancedb/tests/remote_blob_ranges.rs new file mode 100644 index 000000000..51c49fe09 --- /dev/null +++ b/rust/lancedb/tests/remote_blob_ranges.rs @@ -0,0 +1,634 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +#![cfg(feature = "remote")] + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use arrow_array::{Array, ArrayRef, LargeBinaryArray, RecordBatch, StringArray}; +use arrow_ipc::writer::StreamWriter; +use lancedb::Table; +use lancedb::blob::BlobRangeRequest; +use lancedb::remote::{ClientConfig, RetryConfig, TimeoutConfig}; +use serde_json::{Value, json}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; +use tokio::net::TcpListener; + +const MIB: u64 = 1024 * 1024; +const IDENTITY: &str = "v1:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; +const ROUTE: &str = "/v1/table/test/fetch_blob_ranges/"; + +struct Reply { + status: u16, + version: Option, + body: Vec, + declared_extra: usize, + identity: Option<&'static str>, +} + +fn ipc(column: &str, values: ArrayRef) -> Vec { + let batch = RecordBatch::try_from_iter([(column, values)]).unwrap(); + let mut body = Vec::new(); + let mut writer = StreamWriter::try_new(&mut body, &batch.schema()).unwrap(); + // Multiple batches exercise response assembly as well as HTTP partitions. + let mid = batch.num_rows() / 2; + writer.write(&batch.slice(0, mid)).unwrap(); + writer + .write(&batch.slice(mid, batch.num_rows() - mid)) + .unwrap(); + writer.finish().unwrap(); + drop(writer); + body +} + +fn legacy_reply(path: &str, body: &Value, takes: usize) -> Reply { + use arrow_array::{StructArray, UInt8Array, UInt32Array, UInt64Array}; + let version = body["version"] + .as_u64() + .unwrap_or(if takes == 1 { 7 } else { 8 }); + if path.ends_with("/query/") { + assert_eq!(body["branch"], "experiment"); + assert_eq!(body["columns"], json!(["image"])); + let filter = body["filter"].as_str().unwrap(); + let ids = filter + .split_once('(') + .unwrap() + .1 + .trim_end_matches(')') + .split(',') + .map(|id| id.trim().parse::().unwrap()) + .filter(|id| [10, 20, 30, 40].contains(id)) + .rev() + .collect::>(); + let fields = lance_core::datatypes::BLOB_V2_DESC_FIELDS.clone(); + let descriptors = StructArray::new( + fields, + vec![ + Arc::new(UInt8Array::from(vec![0; ids.len()])), + Arc::new(UInt64Array::from(vec![0; ids.len()])), + Arc::new(UInt64Array::from( + ids.iter() + .map(|id| if *id == 40 { 0 } else { 5 }) + .collect::>(), + )), + Arc::new(UInt32Array::from(vec![0; ids.len()])), + Arc::new(StringArray::from(vec![""; ids.len()])), + ], + Some(ids.iter().map(|id| *id != 20).collect()), + ); + let batch = RecordBatch::try_from_iter([ + ("image", Arc::new(descriptors) as ArrayRef), + ("_rowid", Arc::new(UInt64Array::from(ids)) as ArrayRef), + ]) + .unwrap(); + let mut bytes = Vec::new(); + let mut writer = StreamWriter::try_new(&mut bytes, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + Reply { + status: 200, + version: Some(version), + body: bytes, + declared_extra: 0, + identity: Some(IDENTITY), + } + } else { + assert!(path.contains("branch=experiment"), "{path}"); + let version = path + .split("version=") + .nth(1) + .unwrap() + .split('&') + .next() + .unwrap() + .parse::() + .unwrap(); + let range = body["range"] + .as_str() + .unwrap() + .strip_prefix("bytes=") + .unwrap(); + let (start, end) = range.split_once('-').unwrap(); + let start = start.parse::().unwrap(); + let end = end.parse::().unwrap(); + let payload = if version == 8 { b"ALPHA" } else { b"alpha" }; + Reply { + status: 206, + version: Some(version), + body: payload[start..=end].to_vec(), + declared_extra: 0, + identity: Some(IDENTITY), + } + } +} + +fn reply(values: Vec>>, version: u64) -> Reply { + Reply { + status: 200, + version: Some(version), + body: ipc( + "value", + Arc::new(LargeBinaryArray::from_iter( + values.iter().map(|v| v.as_deref()), + )), + ), + declared_extra: 0, + identity: Some(IDENTITY), + } +} + +fn byte(row: u64, offset: u64, version: u64) -> u8 { + ((row * 17 + offset % 251 + version) % 251) as u8 +} + +fn serve_ranges(body: &Value, version: u64) -> Reply { + assert_eq!(body["column"], "image"); + let ranges = body["ranges"].as_array().unwrap(); + assert!(ranges.len() <= 1024); + assert!( + ranges + .iter() + .map(|r| r["length"].as_u64().unwrap()) + .sum::() + <= 64 * MIB + ); + let values = ranges + .iter() + .map(|r| { + let row = r["row_id"].as_u64().unwrap(); + let offset = r["offset"].as_u64().unwrap(); + let length = r["length"].as_u64().unwrap(); + (row != 20).then(|| { + (offset..offset + length) + .map(|i| byte(row, i, version)) + .collect() + }) + }) + .collect(); + reply(values, version) +} + +struct Server { + url: String, + calls: Arc>>, + task: tokio::task::JoinHandle<()>, + identities: Arc>>>, +} + +impl Drop for Server { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl Server { + async fn new(handler: impl Fn(&Value, usize) -> Reply + Send + Sync + 'static) -> Self { + Self::with_capability(handler, Some("1")).await + } + + async fn with_capability( + handler: impl Fn(&Value, usize) -> Reply + Send + Sync + 'static, + capability: Option<&'static str>, + ) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let calls = Arc::new(Mutex::new(Vec::new())); + let captured = calls.clone(); + let identities = Arc::new(Mutex::new(Vec::new())); + let captured_identities = identities.clone(); + let task = tokio::spawn(async move { + let mut fetches = 0; + let mut takes = 0; + loop { + let (socket, _) = listener.accept().await.unwrap(); + let captured = captured.clone(); + let result = tokio::time::timeout(Duration::from_secs(20), async { + let mut socket = BufReader::new(socket); + let mut line = String::new(); + socket.read_line(&mut line).await?; + let path = line.split_whitespace().nth(1).unwrap().to_string(); + let method = line.split_whitespace().next().unwrap().to_string(); + let mut range = String::new(); + let mut length = 0; + let mut identity = None; + loop { + line.clear(); + socket.read_line(&mut line).await?; + assert!(line.len() <= 16384); + if line == "\r\n" { break; } + if let Some((name, value)) = line.split_once(':') { + if name.eq_ignore_ascii_case("content-length") { + length = value.trim().parse::().unwrap(); + } + if name.eq_ignore_ascii_case("x-lancedb-query-snapshot") { identity = Some(value.trim().to_owned()); } + if name.eq_ignore_ascii_case("range") { range = value.trim().to_string(); } + } + } + assert!(length <= 256 * 1024); + let mut bytes = vec![0; length]; + socket.read_exact(&mut bytes).await?; + let body = if method == "GET" { json!({"range": range}) } else if bytes.is_empty() { json!({}) } else { serde_json::from_slice(&bytes).unwrap() }; + captured.lock().unwrap().push((path.clone(), body.clone())); + let response = if path == "/v1/table/test/describe/" { + Reply { status: 200, version: None, body: serde_json::to_vec(&json!({"version": 8, "schema": {"fields": []}})).unwrap(), declared_extra: 0, identity: None } + } else if path == "/v1/table/test/branches/list/" { + Reply { status: 200, version: None, body: br#"{"branches":{"experiment":{"parentBranch":null,"parentVersion":1,"createAt":1,"manifestSize":1}}}"#.to_vec(), declared_extra: 0, identity: None } + } else if capability != Some("1") { + assert_ne!(path, ROUTE, "unknown capability must select old routes first"); + takes += usize::from(path.ends_with("/query/")); + legacy_reply(&path, &body, takes) + } else { + assert_eq!(path, ROUTE, "unexpected legacy or fallback request"); + fetches += 1; + captured_identities.lock().unwrap().push(identity); + handler(&body, fetches) + }; + if response.status == 0 { return Ok::<_, std::io::Error>(()); } + let version = response.version.map(|v| format!("x-lancedb-version: {v}\r\n")).unwrap_or_default(); + let capability = if path.ends_with("/describe/") { capability.map(|value| format!("x-lancedb-blob-ranges: {value}\r\n")).unwrap_or_default() } else { String::new() }; + let content_type = if path == ROUTE || path.ends_with("/query/") { "application/vnd.apache.arrow.stream" } else { "application/json" }; + let content_range = if method == "GET" { format!("Content-Range: {}/5\r\n", body["range"].as_str().unwrap().replace("bytes=", "bytes ")) } else { String::new() }; + let identity = if path == ROUTE { response.identity.map(|id| format!("x-lancedb-query-snapshot: {id}\r\n")).unwrap_or_default() } else { String::new() }; + let header = format!("HTTP/1.1 {} Response\r\nContent-Type: {content_type}\r\nphalanx-version: 1.13.0-beta.0\r\n{capability}{version}{identity}{content_range}Content-Length: {}\r\nConnection: close\r\n\r\n", response.status, response.body.len() + response.declared_extra); + socket.write_all(header.as_bytes()).await?; + for chunk in response.body.chunks(8192) { socket.write_all(chunk).await?; } + socket.shutdown().await?; + Ok(()) + }).await; + // A client rejecting a malformed schema may close the body early. + assert!(result.is_ok(), "loopback request timed out"); + } + }); + Self { + url, + calls, + task, + identities, + } + } + + async fn table(&self, version: Option) -> Table { + let db = lancedb::connect("db://dev") + .api_key("fake") + .region("us-east-1") + .host_override(&self.url) + .client_config(ClientConfig { + retry_config: RetryConfig { + retries: Some(2), + connect_retries: Some(2), + read_retries: Some(2), + ..Default::default() + }, + timeout_config: TimeoutConfig { + timeout: Some(Duration::from_secs(30)), + ..Default::default() + }, + ..Default::default() + }) + .execute() + .await + .unwrap(); + let mut open = db.open_table("test").branch("experiment"); + if let Some(version) = version { + open = open.version(version); + } + open.execute().await.unwrap() + } +} + +#[tokio::test] +async fn public_ranges_work_on_old_routes_only_servers() { + for capability in [None, Some("0"), Some("2")] { + for checkout in [None, Some(3)] { + let server = Server::with_capability( + |_, _| panic!("new endpoint must not be requested"), + capability, + ) + .await; + let table = server.table(checkout).await; + let cloned = table.clone(); + let ranges = [ + BlobRangeRequest::new(10, 1, 3), + BlobRangeRequest::new(20, 0, 1), + BlobRangeRequest::new(40, 0, 0), + BlobRangeRequest::new(10, 5, 0), + BlobRangeRequest::new(10, 1, 3), + ]; + let result = cloned.fetch_blob_ranges("image", ranges).await.unwrap(); + assert_eq!( + result, + LargeBinaryArray::from(vec![ + Some(b"lph".as_slice()), + None, + Some(b"".as_slice()), + Some(b"".as_slice()), + Some(b"lph".as_slice()) + ]) + ); + let result = table + .fetch_blob_ranges("image", [BlobRangeRequest::new(10, 0, 1)]) + .await + .unwrap(); + assert_eq!( + result.value(0), + if checkout.is_none() { b"A" } else { b"a" } + ); + for ranges in [ + vec![BlobRangeRequest::new(10, 4, 2)], + vec![BlobRangeRequest::new(99, 0, 1)], + ] { + assert!(matches!( + table.fetch_blob_ranges("image", ranges).await.unwrap_err(), + lancedb::Error::InvalidInput { .. } + )); + } + assert!( + !server + .calls + .lock() + .unwrap() + .iter() + .any(|(path, _)| path == ROUTE) + ); + } + } +} + +#[tokio::test] +async fn public_ranges_keep_order_nulls_and_snapshot_without_pinning_live_handle() { + for checkout in [None, Some(3)] { + let server = Server::new(move |body, call| { + assert_eq!(body["branch"], "experiment"); + let expected = checkout.or((call > 1 && call < 3).then_some(7)); + assert_eq!(body["version"].as_u64(), expected); + serve_ranges( + body, + body["version"] + .as_u64() + .unwrap_or(if call == 1 { 7 } else { 8 }), + ) + }) + .await; + let table = server.table(checkout).await; + let mut requests = (0..1025) + .map(|i| BlobRangeRequest::new(if i % 9 == 0 { 20 } else { 10 + i % 7 }, i, i % 4)) + .collect::>(); + requests.extend([ + BlobRangeRequest::new(10, 0, 0), + BlobRangeRequest::new(10, 2, 3), + BlobRangeRequest::new(10, 2, 3), + ]); + let result = table + .fetch_blob_ranges("image", requests.clone()) + .await + .unwrap(); + assert_eq!(result.len(), requests.len()); + for (i, request) in requests.iter().enumerate() { + assert_eq!(result.is_null(i), request.row_id == 20); + if request.row_id != 20 { + let expected = (request.offset..request.offset + request.length) + .map(|j| byte(request.row_id, j, checkout.unwrap_or(7))) + .collect::>(); + assert_eq!(result.value(i), expected); + } + } + let next = table + .fetch_blob_ranges("image", [BlobRangeRequest::new(10, 0, 1)]) + .await + .unwrap(); + assert_eq!(next.value(0), [byte(10, 0, checkout.unwrap_or(8))]); + assert_eq!( + server + .calls + .lock() + .unwrap() + .iter() + .filter(|(path, _)| path == ROUTE) + .count(), + 3 + ); + } +} + +#[tokio::test] +async fn public_ranges_partition_bytes_and_reassemble_large_single_range() { + let server = Server::new(|body, call| { + assert_eq!(body["version"].as_u64(), (call > 1).then_some(7)); + serve_ranges(body, 7) + }) + .await; + let table = server.table(None).await; + let requests = [ + BlobRangeRequest::new(10, 3, 33 * MIB), + BlobRangeRequest::new(30, 7, 33 * MIB), + BlobRangeRequest::new(40, 11, 70 * MIB), + BlobRangeRequest::new(20, 0, u64::MAX), + ]; + let result = table.fetch_blob_ranges("image", requests).await.unwrap(); + assert_eq!(result.len(), 4); + for (i, request) in requests[..3].iter().enumerate() { + assert_eq!(result.value_length(i) as u64, request.length); + assert!( + result + .value(i) + .iter() + .enumerate() + .all(|(j, value)| *value == byte(request.row_id, request.offset + j as u64, 7)) + ); + } + assert!(result.is_null(3)); + assert_eq!( + server + .calls + .lock() + .unwrap() + .iter() + .filter(|(path, _)| path == ROUTE) + .count(), + 3 + ); +} + +#[tokio::test] +async fn public_ranges_reject_malformed_responses_and_http_failures_without_fallback() { + for case in 0..16 { + let server = Server::new(move |body, call| { + let mut response = serve_ranges(body, 7); + // Fail after one successful partition to prove no partial success. + if call == 1 { + return response; + } + match case { + 0 => response.body = ipc("value", Arc::new(StringArray::from(vec!["bad"]))), + 1 => { + response.body = ipc( + "wrong", + Arc::new(LargeBinaryArray::from(vec![Some(b"a".as_slice())])), + ) + } + 2 => response = reply(vec![], 7), + 3 => response = reply(vec![Some(vec![1]), Some(vec![2])], 7), + 4 => response.body.truncate(response.body.len() - 8), + 5 => response.body.truncate(response.body.len() / 2), + 6 => response.declared_extra = 10, + 7 => response = reply(vec![Some(vec![])], 7), + 8 => response.version = None, + 9 => response.version = Some(8), + 10..=14 => { + response.status = [401, 403, 404, 429, 500][case - 10]; + response.body = b"ordinary failure".to_vec(); + } + 15 => response.status = 0, + _ => unreachable!(), + } + response + }) + .await; + let table = server.table(None).await; + let requests = vec![BlobRangeRequest::new(10, 0, 1); 1025]; + assert!( + table.fetch_blob_ranges("image", requests).await.is_err(), + "case {case}" + ); + assert_eq!( + server + .calls + .lock() + .unwrap() + .iter() + .filter(|(path, _)| path == ROUTE) + .count(), + 2, + "case {case}" + ); + } +} + +#[tokio::test] +async fn public_ranges_validate_locally_and_propagate_server_range_errors() { + let server = Server::new(|_, _| Reply { + status: 400, + version: None, + body: b"range exceeds blob size or missing row".to_vec(), + declared_extra: 0, + identity: Some(IDENTITY), + }) + .await; + let table = server.table(None).await; + let before = server.calls.lock().unwrap().len(); + assert_eq!(table.fetch_blob_ranges("image", []).await.unwrap().len(), 0); + assert!( + table + .fetch_blob_ranges("image", [BlobRangeRequest::new(20, u64::MAX, 1)]) + .await + .unwrap_err() + .to_string() + .contains("overflows u64") + ); + assert_eq!(server.calls.lock().unwrap().len(), before); + let error = table + .fetch_blob_ranges("image", [BlobRangeRequest::new(10, 9, 2)]) + .await + .unwrap_err(); + assert!(matches!(error, lancedb::Error::InvalidInput { .. })); + assert!(error.to_string().contains("range exceeds")); +} + +#[tokio::test] +async fn public_ranges_require_original_identity_on_every_partition() { + for huge in [false, true] { + for case in 0..5 { + let server = Server::new(move |body, call| { + let mut response = serve_ranges(body, 7); + if call == 2 { + response.identity = match case { + 0 => None, // Old serving node ignores the precondition. + 1 => Some("invalid"), + _ => Some( + "v1:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + ), + }; + if case == 3 || case == 4 { + response.status = if case == 3 { 400 } else { 404 }; + response.body = + b"query snapshot identity mismatch: gone or replaced".to_vec(); + } else { + // Identity verification must win over malformed payload. + response.body = b"not Arrow".to_vec(); + } + } + response + }) + .await; + let table = server.table(None).await; + let requests = if huge { + vec![BlobRangeRequest::new(10, 0, 64 * MIB + 1)] + } else { + vec![BlobRangeRequest::new(10, 0, 1); 1025] + }; + let error = table + .fetch_blob_ranges("image", requests) + .await + .unwrap_err(); + if case < 3 { + assert!( + matches!(error, lancedb::Error::NotSupported { .. }), + "{error}" + ); + assert!(error.to_string().contains("snapshot identity"), "{error}"); + } else if case == 3 { + assert!( + matches!(error, lancedb::Error::InvalidInput { .. }), + "{error}" + ); + } else { + assert!( + matches!(error, lancedb::Error::TableNotFound { .. }), + "{error}" + ); + } + assert_eq!( + *server.identities.lock().unwrap(), + [None, Some(IDENTITY.into())] + ); + assert_eq!( + server + .calls + .lock() + .unwrap() + .iter() + .filter(|(path, _)| path == ROUTE) + .count(), + 2 + ); + } + } +} + +#[tokio::test] +async fn public_ranges_cannot_establish_identity_from_numeric_version_or_describe() { + for identity in [None, Some("invalid"), Some("v2:unsupported")] { + let server = Server::new(move |_, call| { + assert_eq!(call, 1, "no failure fallback, replay or second partition"); + let mut response = reply(vec![Some(b"A".to_vec()); 1024], 7); + response.identity = identity; + response.body = b"not Arrow".to_vec(); + response + }) + .await; + let table = server.table(None).await; + let before = server.calls.lock().unwrap().len(); + let error = table + .fetch_blob_ranges("image", vec![BlobRangeRequest::new(10, 0, 1); 1025]) + .await + .unwrap_err(); + assert!( + matches!(error, lancedb::Error::NotSupported { .. }), + "{error}" + ); + assert!(error.to_string().contains("snapshot identity"), "{error}"); + assert_eq!(server.calls.lock().unwrap().len(), before + 1); + } +}