From 6d24446a3043fcb22a73c2c9aed21b62a5648938 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 13 Nov 2026 04:03:04 +0800 Subject: [PATCH] feat: add negotiated batched blob range reads across SDKs (#4994) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Existing Rust and Python blob range calls must keep working when only the SDK is upgraded. Add negotiated batched range reads across Rust, Python sync/async and TypeScript, preserving ordered NULLs, duplicates and branch/version reads with splitting and reassembly at 1,024 ranges or 64 MiB per request. Only a successful describe response advertising `x-lancedb-blob-ranges: 1` selects the batched endpoint. Missing or unknown capability values select the existing descriptor-query plus BlobFile Range GET route before the operation. Endpoint failures never switch protocols. HTTP 400 preserves the server message and request ID as Rust `InvalidInput` / Python `ValueError`; other errors retain their existing semantics. The operation binds both numeric version and immutable snapshot identity from its producing query or first range response. Every later partition and oversized-range piece sends that original identity and requires an equal response identity before decoding. Same-name table recreation therefore fails explicitly instead of returning mixed old/new payloads; no partial array is returned on later failure. Targets `main`. The read-buffering prerequisite #4979, safe-write prerequisite #4977 and deferred-query identity prerequisite #4973 are merged into main; #5000 also supplies the merged SQL timeout-test prerequisite. Server support is [Sophon #10219](https://github.com/lancedb/sophon/pull/10219), including #10211. Exact SDK `1497b508324662180086a4d1e0a6eb9e0b3970db` and server `f39712c2c150b220eb831e00f6663c9eef32dce5` passed a real local SDK → REST → PE experiment: one Rust harness test, 14 assertion groups, 38 REST range requests and 34 PE RPCs. All four sync/async × warm/cold table-recreation cases failed explicitly before the second descriptor lookup or PE dispatch, with no mixed payload or partial result. The other groups retained 1,025-range success, 70 MiB reassembly, branch/version reads, overwrite snapshot consistency, error behavior and TypeScript buffer ownership. Bindings were rebuilt and their loaded hashes verified. This is local interoperability evidence, not Cloud, mixed-fleet or deployed header-forwarding acceptance. The server dependency upgrade, External/cache/URI-policy acceptance and validated rollout remain required. Enforcing PE support must precede QN and SDK enablement; identity does not retain versions or freeze External objects. This PR does not claim deployed support or complete ENT-3100/GA acceptance. Published documentation retains the Early access badge and released Enterprise instructions, including explicit fetches for copying. The new Enterprise `blob_mode="bytes"` guidance, snapshot-bound copy example, future streaming restrictions and GA wording remain source-only MDX drafts until release. The correction distinguishing released Enterprise SQL bytes from SDK descriptors remains visible. Documentation validation checked comment boundaries, visible instructions, Python example syntax and whitespace; the site was not built. The main integration keeps these release gates unchanged. Main integration at `0491d9bdcc65ce5d52ea93e1083c89e951ea4f9d` normally merges main `9b13e8f72530c8b1950c09685f6920711154a539`, including its Lance beta.18 dependency. Focused current-source validation passed: 493 remote-table/query/Blob/IPC/write-progress Rust tests, 23 shared create/write-body tests, 59 native Blob integration tests and 7 public HTTP range tests. Rebuilt Python passed 468 selected tests (one optional PyAV test skipped), and rebuilt TypeScript passed 194 selected tests across four suites; lint/format, TypeScript compilation and API documentation generation passed. Shared tests include actual HTTP multipart failures, acknowledgment boundaries, dictionary compaction and native Blob compaction. The 598 unrelated Jest cases excluded by the focused shared-suite selector were not run. The retained 14-group joint REST/PE experiment used the earlier exact candidates stated above. Operation-identity/range/IPC code is unchanged, but dependencies and shared write interfaces changed with main; the new local checks validate the SDK combination and do not constitute a fresh final-head SDK/server joint, Cloud or rollout run. Early access and every remaining server/External/cache/policy/dependency/rollout condition remain in force. CI snapshot on this head (2026-11-12 18:30 UTC): `build-no-lock` failed in the standalone `ci/tests/test_build_no_lock.py` fixture at the retained-rlib assertion, before the actual SDK no-lock build. The fixture, helper and Rust workflow are byte-identical to main; the fixture does not load SDK source. The underlying CI/environment cause is not diagnosed here. Other jobs remained pending/running at the snapshot, so CI is not reported as passing. Failure: https://github.com/lancedb/lancedb/actions/runs/43280277756/job/131399418284. --- docs/src/js/classes/Table.md | 42 ++ docs/src/js/globals.md | 1 + docs/src/js/type-aliases/BlobRangeRequest.md | 39 ++ docs/web/management/table.mdx | 27 + nodejs/__test__/blob_ranges.test.ts | 589 ++++++++++++++++ nodejs/__test__/table.test.ts | 37 + nodejs/lancedb/blob.ts | 10 + nodejs/lancedb/index.ts | 1 + nodejs/lancedb/table.ts | 49 +- nodejs/src/blob.rs | 50 +- nodejs/src/table.rs | 29 +- python/python/lancedb/table.py | 12 +- python/python/tests/test_blob_ranges.py | 380 +++++++++++ python/python/tests/test_remote_db.py | 46 +- rust/lancedb/src/blob.rs | 2 +- rust/lancedb/src/remote/db.rs | 10 +- rust/lancedb/src/remote/table.rs | 28 + rust/lancedb/src/remote/table/blobs.rs | 400 ++++++++++- .../src/remote/table/blobs/aggregate.rs | 136 +++- rust/lancedb/src/table.rs | 21 +- rust/lancedb/tests/remote_blob_ranges.rs | 634 ++++++++++++++++++ 21 files changed, 2473 insertions(+), 70 deletions(-) create mode 100644 docs/src/js/type-aliases/BlobRangeRequest.md create mode 100644 nodejs/__test__/blob_ranges.test.ts create mode 100644 python/python/tests/test_blob_ranges.py create mode 100644 rust/lancedb/tests/remote_blob_ranges.rs 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); + } +}