From f5f4463b15bf1e6b08725da26f05cf8dd5188888 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:42:22 +0800 Subject: [PATCH 01/36] fix(node): materialize nested struct blobs in bytes mode (#4611) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `blobMode: "bytes"` left blobs inside structs as descriptors because projection matching ignored parent structs and byte replacement only rebuilt top-level fields. Expand direct struct projections and aliases to their blob leaves, then recursively replace the corresponding Arrow fields and data with `LargeBinary`. Track literal output names separately from descendant paths so scalar aliases stay intact and every overlapping blob projection is converted. Preserve null structs, null and empty blobs, sibling fields, schema metadata, and the existing row ID and snapshot behavior. Blobs inside lists retain their descriptors. Update the API docs to describe struct support. Regression coverage includes full and parent-struct projections, leaf and struct aliases, dotted output names, overlapping projections in both orders, nested structs, nulls, empty blobs and results, multiple batches, list siblings, query reuse, row IDs, and concurrent writes. Regression tests reproduce the nested-blob and projection-collision failures before the corresponding changes. Validation: - `pnpm build` - `pnpm lint` - `pnpm run docs` - `pnpm test __test__/nested_blob.test.ts __test__/blob.test.ts __test__/query.test.ts --runInBand --verbose=false` — 142 passed - `pnpm test __test__/table.test.ts --runInBand --verbose=false -t 'when dealing with blob columns'` — 60 passed Fixes #4609 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- .../js/interfaces/QueryExecutionOptions.md | 13 +- nodejs/__test__/nested_blob.test.ts | 327 ++++++++++++++++++ nodejs/__test__/table.test.ts | 125 ++++--- nodejs/lancedb/query.ts | 259 +++++++------- 4 files changed, 528 insertions(+), 196 deletions(-) create mode 100644 nodejs/__test__/nested_blob.test.ts diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index c0b33dcf0..f18ada7aa 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -20,13 +20,14 @@ How to return blob v2 columns. Applies to `toArrow()` and `toArray()`. - `"descriptions"` (default): each blob is a descriptor. Read the bytes with [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles). -- `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - including blobs nested in structs. The query also reads `_rowid` to - fetch the bytes, and leaves it out of the result unless - [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was called. All the bytes are held in - memory. +- `"bytes"`: each blob column, including blobs nested in structs, holds + the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + leaves it out of the result unless [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was + called. All the bytes are held in memory. -Blobs nested in a list keep their descriptors. +Selecting a parent struct or directly aliasing a blob or struct also returns +bytes. Blobs inside lists keep their descriptors. Arbitrary projection +expressions are not materialized as bytes. The rows and the bytes come from the same table version. If the table changes while the query runs, the query runs again, up to three times. diff --git a/nodejs/__test__/nested_blob.test.ts b/nodejs/__test__/nested_blob.test.ts new file mode 100644 index 000000000..508b57c77 --- /dev/null +++ b/nodejs/__test__/nested_blob.test.ts @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import * as tmp from "tmp"; +import { type Table, blob, connect, isBlobField } from "../lancedb"; +import { + DataType, + Field, + Int64, + List, + Schema, + Struct, + Utf8, +} from "../lancedb/arrow"; + +describe('nested blobs with blobMode "bytes"', () => { + let tmpDir: tmp.DirResult; + let table: Table; + + const caption = new Field( + "caption", + new Utf8(), + false, + new Map([["label", "caption"]]), + ); + const info = new Field( + "info", + new Struct([ + blob("image"), + caption, + new Field("preview", new Struct([blob("image")]), true), + ]), + true, + new Map([["label", "info"]]), + ); + const schema = new Schema( + [new Field("id", new Int64(), false), info, blob("cover")], + new Map([["source", "nested-blobs"]]), + ); + + beforeEach(async () => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + table = await db.createTable( + "nested_blobs", + [ + { + id: 1n, + info: { + image: Buffer.from("hello"), + caption: "first", + preview: { image: Buffer.from("preview") }, + }, + cover: Buffer.from("cover"), + }, + { id: 2n, info: null, cover: null }, + { + id: 3n, + info: { image: null, caption: "third", preview: null }, + cover: null, + }, + { + id: 4n, + info: { + image: Buffer.alloc(0), + caption: "fourth", + preview: { image: null }, + }, + cover: Buffer.alloc(0), + }, + ], + { schema }, + ); + }); + + afterEach(() => tmpDir.removeCallback()); + + it.each([false, true])( + "reads struct leaves across batches (select parent: %s)", + async (selectParent) => { + const query = table + .query() + .orderBy({ columnName: "id", ascending: true }); + if (selectParent) { + query.select(["info"]); + } + const result = await query.toArrow({ + blobMode: "bytes", + maxBatchLength: 1, + }); + expect(result.batches.length).toBeGreaterThan(1); + expect(result.schema.fields.map((field) => field.name)).toEqual( + selectParent ? ["info"] : ["id", "info", "cover"], + ); + expect(result.schema.metadata).toEqual(schema.metadata); + const infoField = result.schema.fields.find( + (field) => field.name === "info", + )!; + expect(infoField.nullable).toBe(true); + expect(infoField.metadata).toEqual(info.metadata); + const fields = (infoField.type as Struct).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[0].metadata.size).toBe(0); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + + const rows = result.toArray(); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(Buffer.from(rows[0].info.preview.image).toString()).toBe( + "preview", + ); + expect(rows[1].info).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(rows[2].info.caption).toBe("third"); + expect(rows[2].info.preview).toBeNull(); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + expect(rows[3].info.preview.image).toBeNull(); + if (!selectParent) { + expect(Buffer.from(rows[0].cover).toString()).toBe("cover"); + expect(rows[1].cover).toBeNull(); + } + }, + ); + + it("reads direct leaf and struct aliases", async () => { + const rows = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select({ + image: "info.image", + details: "info", + preview: "info.preview", + caption: "info.caption", + }) + .toArray({ blobMode: "bytes", maxBatchLength: 1 }); + expect(Object.keys(rows[0])).toEqual([ + "image", + "details", + "preview", + "caption", + ]); + expect(Buffer.from(rows[0].image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.preview.image).toString()).toBe( + "preview", + ); + expect(Buffer.from(rows[0].preview.image).toString()).toBe("preview"); + expect(rows[0].caption).toBe("first"); + expect(rows[1].details).toBeNull(); + expect(rows[1].image).toBeNull(); + expect(rows[2].image).toBeNull(); + expect(Buffer.from(rows[3].image)).toEqual(Buffer.alloc(0)); + }); + + it.each([false, true])( + "preserves a literal scalar alias matching a nested blob path (reverse: %s)", + async (reverse) => { + const projection: [string, string][] = [ + ["info", "info"], + ["info.image", "id"], + ]; + if (reverse) { + projection.reverse(); + } + const [row] = await table + .query() + .where("id = 1") + .select(new Map(projection)) + .toArray({ blobMode: "bytes" }); + expect(row["info.image"]).toBe(1n); + expect(Buffer.from(row.info.image).toString()).toBe("hello"); + }, + ); + + it.each([false, true])( + "materializes every overlapping blob projection (reverse: %s)", + async (reverse) => { + const columns = ["info.image", "info"]; + if (reverse) { + columns.reverse(); + } + const result = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select(columns) + .toArrow({ blobMode: "bytes", maxBatchLength: 1 }); + expect(result.schema.fields.map((field) => field.name)).toEqual(columns); + expect(DataType.isLargeBinary(result.getChild("info.image")!.type)).toBe( + true, + ); + expect( + DataType.isLargeBinary( + result.getChild("info")!.getChild("image")!.type, + ), + ).toBe(true); + const rows = result.toArray(); + expect(Buffer.from(rows[0]["info.image"]).toString()).toBe("hello"); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(rows[1]["info.image"]).toBeNull(); + expect(rows[1].info).toBeNull(); + expect(rows[2]["info.image"]).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(Buffer.from(rows[3]["info.image"])).toEqual(Buffer.alloc(0)); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + }, + ); + + it("materializes dotted aliases independently of nested fields", async () => { + const [row] = await table + .query() + .where("id = 1") + .select({ + details: "info", + "details.image": "info.image", + "details.preview": "info.preview", + "details.preview.image": "info.preview.image", + }) + .toArray({ blobMode: "bytes" }); + expect(Buffer.from(row.details.image).toString()).toBe("hello"); + expect(Buffer.from(row["details.image"]).toString()).toBe("hello"); + expect(Buffer.from(row.details.preview.image).toString()).toBe("preview"); + expect(Buffer.from(row["details.preview"].image).toString()).toBe( + "preview", + ); + expect(Buffer.from(row["details.preview.image"]).toString()).toBe( + "preview", + ); + }); + + it.each([false, true])( + "converts the empty result schema (alias: %s)", + async (alias) => { + const query = table.query().where("id < 0"); + if (alias) { + query.select({ image: "info.image", details: "info" }); + } + const result = await query.toArrow({ blobMode: "bytes" }); + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual( + alias ? ["image", "details"] : ["id", "info", "cover"], + ); + const fields = ( + result.schema.fields.find( + (field) => field.name === (alias ? "details" : "info"), + )!.type as Struct + ).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + if (alias) { + expect(DataType.isLargeBinary(result.schema.fields[0].type)).toBe(true); + } + }, + ); + + it("preserves requested row IDs and leaves subsequent descriptors intact", async () => { + const query = table.query().where("id = 1").select(["info"]); + const [bytes] = await query.toArray({ blobMode: "bytes" }); + expect(Object.keys(bytes)).toEqual(["info"]); + const [description] = await query.toArray(); + expect(Object.keys(description)).toEqual(["info"]); + expect(Number(description.info.image.size)).toBe(5); + const [withRowId] = await query.withRowId().toArray({ blobMode: "bytes" }); + expect(typeof withRowId._rowid).toBe("bigint"); + expect(await table.fetchBlobs("info.image", [withRowId._rowid])).toEqual([ + Buffer.from(withRowId.info.image), + ]); + }); + + it("keeps descriptors in description mode", async () => { + const [description] = await table + .query() + .where("id = 1") + .toArray({ blobMode: "descriptions" }); + expect(Number(description.info.image.size)).toBe(5); + }); + + it("keeps list elements as descriptors alongside struct blob bytes", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable( + "list_blobs", + [ + { + info: { image: Buffer.from("hello"), images: [Buffer.from("list")] }, + }, + ], + { + schema: new Schema([ + new Field( + "info", + new Struct([ + blob("image"), + new Field("images", new List(blob("item")), true), + ]), + true, + ), + ]), + }, + ); + const [bytes] = await table.query().toArray({ blobMode: "bytes" }); + expect(Buffer.from(bytes.info.image).toString()).toBe("hello"); + expect(Number(bytes.info.images.get(0).size)).toBe(4); + const [row] = await table + .query() + .select(["info.images"]) + .toArray({ blobMode: "bytes" }); + expect(Number(row["info.images"].get(0).size)).toBe(4); + expect(Object.keys(row)).toEqual(["info.images"]); + const result = await table.query().toArrow({ blobMode: "bytes" }); + const fields = (result.schema.fields[0].type as Struct).children; + const descriptions = await table.query().toArrow(); + expect(fields[1]).toEqual( + (descriptions.schema.fields[0].type as Struct).children[1], + ); + expect(DataType.isStruct((fields[1].type as List).children[0].type)).toBe( + true, + ); + }); +}); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 92e3df456..d848676cd 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3314,51 +3314,64 @@ describe("when dealing with blob columns", () => { ).rejects.toThrow(/blobMode must be/); }); - it("reruns the query when the table changes before the version check", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - await update(2n, "new"); - return liveVersion(); - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(2n); - expect(Buffer.from(row.image).toString()).toBe("new"); - }); + it.each([false, true])( + "reruns the query when the table changes before the version check (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + await update(2n, "new"); + return liveVersion(); + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(2n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "new", + ); + }, + ); - it("reads bytes from the query's version when a write lands before the fetch", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - const version = await liveVersion(); - await update(2n, "new"); - return version; - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(1n); - expect(Buffer.from(row.image).toString()).toBe("old"); - }); + it.each([false, true])( + "reads bytes from the query's version when a write lands before the fetch (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + const version = await liveVersion(); + await update(2n, "new"); + return version; + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(1n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "old", + ); + }, + ); - it("fails when the table changes during every attempt", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - let generation = 1n; - native.version = async () => { - generation += 1n; - await update(generation, `generation ${generation}`); - return liveVersion(); - }; - try { - await expect( - table.query().toArray({ blobMode: "bytes" }), - ).rejects.toThrow(/changed during each of 3 attempts/); - expect(generation).toBe(4n); - } finally { - native.version = liveVersion; - } - }); + it.each([false, true])( + "fails when the table changes during every attempt (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + let generation = 1n; + native.version = async () => { + generation += 1n; + await update(generation, `generation ${generation}`); + return liveVersion(); + }; + try { + await expect( + table.query().toArray({ blobMode: "bytes" }), + ).rejects.toThrow(/changed during each of 3 attempts/); + expect(generation).toBe(4n); + } finally { + native.version = liveVersion; + } + }, + ); it("rejects fetchBlobs on a non-blob column", async () => { const { table, rowIds } = await openBlobTable(); @@ -3642,18 +3655,28 @@ describe("when dealing with blob columns", () => { ).map((value) => (value == null ? null : Number(value.size))); } - async function openVersionedBlobTable() { + async function openVersionedBlobTable(nested = false) { const db = await connect(tmpDir.name, { storageOptions: { newTableEnableStableRowIds: "true" }, }); const schema = new Schema([ new Field("id", new Int64(), false), new Field("generation", new Int64(), true), - blob("image"), + nested + ? new Field("info", new Struct([blob("image")]), true) + : blob("image"), ]); const table = await db.createTable( "versions", - [{ id: 1n, generation: 1n, image: Buffer.from("old") }], + [ + { + id: 1n, + generation: 1n, + ...(nested + ? { info: { image: Buffer.from("old") } } + : { image: Buffer.from("old") }), + }, + ], { schema }, ); const native = ( @@ -3663,7 +3686,15 @@ describe("when dealing with blob columns", () => { await table .mergeInsert("id") .whenMatchedUpdateAll() - .execute([{ id: 1n, generation, image: Buffer.from(image) }]); + .execute([ + { + id: 1n, + generation, + ...(nested + ? { info: { image: Buffer.from(image) } } + : { image: Buffer.from(image) }), + }, + ]); }; return { table, native, update }; } diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 823239ef9..b96e4fda6 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +import type { Data as ArrowData } from "apache-arrow"; import { - type ArrowData, Table as ArrowTable, DataType, Field, @@ -14,7 +14,6 @@ import { extractVectorBuffer, fromBufferToRecordBatch, fromRecordBatchToBuffer, - makeData, tableFromIPC, vectorFromArray, } from "./arrow"; @@ -68,84 +67,42 @@ export interface QueryBlobContext { const BLOB_BYTES_ATTEMPTS = 3; +interface BlobSource { + // Output names are literal: an alias containing dots is still one field. + output: string; + path: string[]; + source: string; +} + function blobSources( blobColumns: string[], projection: [string, string][] | undefined, -): Map { - const holds = (source: string) => - blobColumns.some( - (column) => column === source || column.startsWith(`${source}.`), - ); - const pairs = - projection ?? - Array.from(new Set(blobColumns.map((column) => column.split(".")[0]))).map( - (name): [string, string] => [name, name], - ); - return new Map(pairs.filter(([, source]) => holds(source))); -} - -// Fetched bytes for the blobs under one output column. A blob column is a -// leaf; a struct holds the blobs nested in it, keyed by field name. A Map -// keeps struct field names from reaching shared JavaScript prototypes. -type BlobBytes = (Buffer | null)[] | Map; - -function isLeaf(bytes: BlobBytes): bytes is (Buffer | null)[] { - return Array.isArray(bytes); -} - -function nestedFieldType(type: Struct, bytes: Map): Struct { - return new Struct( - type.children.map((field) => { - const inner = bytes.get(field.name); - if (inner === undefined) { - return field; +): BlobSource[] { + if (projection === undefined) { + return blobColumns.map((source) => { + const [output, ...path] = source.split("."); + return { output, path, source }; + }); + } + const sources: BlobSource[] = []; + for (const [output, source] of projection) { + for (const column of blobColumns) { + if (column === source || column.startsWith(`${source}.`)) { + sources.push({ + output, + path: + column === source ? [] : column.slice(source.length + 1).split("."), + source: column, + }); } - if (isLeaf(inner)) { - return new Field(field.name, new LargeBinary(), true); - } - return new Field( - field.name, - nestedFieldType(field.type as Struct, inner), - field.nullable, - ); - }), - ); -} - -// Rebuilds a struct with the blobs under it as bytes. The struct keeps its own -// validity, so a null struct stays null. -function nestedData( - data: ArrowData, - bytes: Map, - offset: number, -): ArrowData { - const type = data.type as Struct; - const children = data.children.map((child, index) => { - const inner = bytes.get(type.children[index].name); - if (inner === undefined) { - return child; } - if (isLeaf(inner)) { - return vectorFromArray( - inner.slice(offset, offset + data.length), - new LargeBinary(), - ).data[0]; - } - return nestedData(child, inner, offset); - }); - return makeData({ - type: nestedFieldType(type, bytes), - length: data.length, - nullCount: data.nullCount, - nullBitmap: data.nullBitmap, - children, - }); + } + return sources; } async function readBlobBytes( table: ArrowTable, - sources: Map, - blobColumns: string[], + sources: BlobSource[], snapshot: NativeTable, dropRowId: boolean, ): Promise { @@ -157,42 +114,27 @@ async function readBlobBytes( rowIdColumn === null ? [] : Array.from(rowIdColumn.toArray() as BigUint64Array); - const fetch = async (column: string) => - (rowIds.length === 0 ? [] : await snapshot.fetchBlobs(column, rowIds)).map( - (value) => value ?? null, + const outputFields = new Map( + table.schema.fields.map((field) => [field.name, field]), + ); + const replacements = new Map(); + for (const { output, path, source } of sources) { + let field = outputFields.get(output); + for (const name of path) { + field = + field !== undefined && DataType.isStruct(field.type) + ? field.type.children.find((child) => child.name === name) + : undefined; + } + if (field === undefined) { + continue; + } + const values = + rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds); + replacements.set( + field, + values.map((value) => value ?? null), ); - const replacements = new Map(); - for (const [output, source] of sources) { - const index = table.schema.fields.findIndex( - (field) => field.name === output, - ); - if (index < 0) { - continue; - } - if (blobColumns.includes(source)) { - replacements.set(index, await fetch(source)); - continue; - } - if (!DataType.isStruct(table.schema.fields[index].type)) { - continue; - } - const nested = new Map(); - for (const column of blobColumns.filter((c) => - c.startsWith(`${source}.`), - )) { - const path = column.slice(source.length + 1).split("."); - let level = nested; - for (const name of path.slice(0, -1)) { - let next = level.get(name); - if (next === undefined) { - next = new Map(); - level.set(name, next); - } - level = next as Map; - } - level.set(path[path.length - 1], await fetch(column)); - } - replacements.set(index, nested); } const result = withBinaryColumns(table, replacements); if (!dropRowId) { @@ -205,48 +147,76 @@ async function readBlobBytes( ); } -// A fresh field, not a clone: the descriptor field's metadata must not follow -// the bytes. function withBinaryColumns( table: ArrowTable, - replacements: Map, + replacements: Map, ): ArrowTable { if (replacements.size === 0) { return table; } - const fields = table.schema.fields.map((field, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return field; + const replaceField = (field: Field): Field => { + if (replacements.has(field)) { + // The descriptor's extension and encoding metadata must not follow bytes. + return new Field(field.name, new LargeBinary(), true); } - return isLeaf(bytes) - ? new Field(field.name, new LargeBinary(), true) - : new Field( - field.name, - nestedFieldType(field.type as Struct, bytes), - field.nullable, - ); - }); + if (DataType.isStruct(field.type)) { + const children = field.type.children.map(replaceField); + if (children.some((child, i) => child !== field.type.children[i])) { + return field.clone({ type: new Struct(children) }); + } + } + return field; + }; + const fields = table.schema.fields.map(replaceField); const schema = new Schema(fields, table.schema.metadata); let offset = 0; const batches = table.batches.map((batch) => { - const children = batch.data.children.map((child, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return child; + const replaceData = ( + data: ArrowData, + field: Field, + updated: Field, + ): ArrowData => { + if (field === updated) { + return data; } - if (!isLeaf(bytes)) { - return nestedData(child, bytes, offset); + const values = replacements.get(field); + if (values !== undefined) { + return vectorFromArray( + values.slice(offset, offset + batch.numRows), + new LargeBinary(), + ).data[0]; } - return vectorFromArray( - bytes.slice(offset, offset + batch.numRows), - new LargeBinary(), - ).data[0]; - }); + const children = data.children.map((child, index) => + replaceData( + child, + field.type.children[index], + updated.type.children[index], + ), + ); + // Keep the struct's validity bitmap and offset, including null parents. + return data.clone( + updated.type, + undefined, + undefined, + undefined, + undefined, + children, + ); + }; + const children = batch.data.children.map((child, index) => + replaceData(child, table.schema.fields[index], fields[index]), + ); offset += batch.numRows; return new RecordBatch( schema, - makeData({ type: new Struct(fields), length: batch.numRows, children }), + batch.data.clone( + new Struct(fields), + undefined, + undefined, + undefined, + undefined, + children, + ), ); }); return new ArrowTable(schema, batches); @@ -274,13 +244,14 @@ export interface QueryExecutionOptions { * * - `"descriptions"` (default): each blob is a descriptor. Read the bytes * with {@link Table.fetchBlobs} or {@link Table.fetchBlobFiles}. - * - `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - * including blobs nested in structs. The query also reads `_rowid` to - * fetch the bytes, and leaves it out of the result unless - * {@link QueryBase.withRowId} was called. All the bytes are held in - * memory. + * - `"bytes"`: each blob column, including blobs nested in structs, holds + * the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + * leaves it out of the result unless {@link QueryBase.withRowId} was + * called. All the bytes are held in memory. * - * Blobs nested in a list keep their descriptors. + * Selecting a parent struct or directly aliasing a blob or struct also returns + * bytes. Blobs inside lists keep their descriptors. Arbitrary projection + * expressions are not materialized as bytes. * * The rows and the bytes come from the same table version. If the table * changes while the query runs, the query runs again, up to three times. @@ -529,15 +500,17 @@ export class QueryBase< // unchanged version afterwards proves it read the snapshot's version. const snapshot = await context.table.querySnapshot(); const version = await snapshot.version(); - const blobColumns = await snapshot.blobColumns(); - const sources = blobSources(blobColumns, context.projection); - if (sources.size === 0) { + const sources = blobSources( + await snapshot.blobColumns(), + context.projection, + ); + if (sources.length === 0) { return this.collect(options, false); } const addRowId = !context.withRowId; const table = await this.collect(options, addRowId); if ((await context.table.version()) === version) { - return readBlobBytes(table, sources, blobColumns, snapshot, addRowId); + return readBlobBytes(table, sources, snapshot, addRowId); } if (attempt === BLOB_BYTES_ATTEMPTS) { throw new Error( From 88203d25156a0dd7fc7d2c51aee66410d546ae18 Mon Sep 17 00:00:00 2001 From: Christopher Krah Date: Thu, 22 Oct 2026 21:42:36 +0200 Subject: [PATCH 02/36] fix(python): keep a renamed Permutation column in its position (#4626) The `Permutation.rename_column` docstring says it "only changes the name of the column that is visible from this permutation". But `permutation.py:868-870` adds the new name at the end of the selection and then deletes the old key, so the renamed column also moves to the last position. Hugging Face `datasets`, which `Permutation` mirrors (#2725), keeps it in place. Positional formats feel it most: `with_format("numpy")` swaps the feature order of every row. Before, on `d97df0c`: ``` $ python repro.py /python/python # lancedb 0.41.0b12 extension, Python sources from the tree ['id', 'value'] ['value', 'row_id'] [{'value': 0, 'row_id': 0}] [[10 1]] ``` After: ``` $ python repro.py /python/python ['id', 'value'] ['row_id', 'value'] [{'row_id': 0, 'value': 0}] [[ 1 10]] ```
repro.py (its first output line, lancedb.__file__, is cut above) ```python import sys; sys.path.insert(0, sys.argv[1]) import tempfile import lancedb from lancedb.permutation import Permutation print(lancedb.__file__) db = lancedb.connect(tempfile.mkdtemp()) tbl = db.create_table("t", [{"id": i, "value": i * 10} for i in range(3)]) perm = Permutation.identity(tbl) print(perm.column_names) renamed = perm.rename_column("id", "row_id") print(renamed.column_names) print(renamed.__getitems__([0])) print(renamed.with_format("numpy").__getitems__([1])) ```
For comparison, `datasets` 5.1.0 gives `['new_id', 'value']` for `Dataset.from_dict({'id': [1, 2], 'value': [3, 4]}).rename_column('id', 'new_id').column_names`. No issue exists for this. The fix rebuilds the selection dict in order. The cost: code that renamed a column and then read it by position in numpy batches (say `[:, -1]`) gets a different column after this, with no error. `rename_columns` calls `rename_column` once per name, so it gets the fix too. The tests that check schema order only rename `value`, the last column, and `test_rename_rowid` checks membership but not order. The new test renames `id`. ### Verification ``` # Python 3.12.6; lancedb 0.41.0b12 wheel from pypi.fury.io, its _lancedb.abi3.so copied into the tree; pylance 14.0.0-rc.1 $ PYTHONPATH=python python -m pytest python/tests/test_permutation.py -q -p no:cacheprovider d97df0c: 72 passed in 2.53s this branch: 73 passed in 2.36s permutation.py reverted, test kept: FAILED python/tests/test_permutation.py::test_rename_column_keeps_position assert value: int64\nnew_id: int64 == new_id: int64\nvalue: int64 1 failed, 72 passed in 2.41s $ ruff format --check . && ruff check . # ruff 0.9.9, as python.yml pins it 152 files already formatted All checks passed! # not run: pyright (its include list covers neither permutation.py nor tests/), the Rust build ``` @westonpace, this is your class from #2725. Do you want the docstring to say outright that the column keeps its position? If so, I'll add that line to both rename methods in this PR. Signed-off-by: Christopher Krah <48338417+chrikrah@users.noreply.github.com> --- python/python/lancedb/permutation.py | 8 +++++--- python/python/tests/test_permutation.py | 7 +++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 355ebf798..e95db27f6 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -909,9 +909,11 @@ class Permutation: f"Cannot rename column `{old_name}` to `{new_name}` because a column " "with that name already exists" ) - new_selection = self.selection.copy() - new_selection[new_name] = new_selection[old_name] - del new_selection[old_name] + # Rebuild the mapping so the renamed column keeps its position + new_selection = { + (new_name if name == old_name else name): value + for name, value in self.selection.items() + } return self._with_selection(new_selection) def rename_columns(self, column_map: dict[str, str]) -> "Permutation": diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index ed1894340..50ca53061 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -811,6 +811,13 @@ def test_rename_column(some_permutation: Permutation): some_permutation.rename_column("non_existent", "new_value") +def test_rename_column_keeps_position(some_permutation: Permutation): + renamed = some_permutation.rename_column("id", "new_id") + assert renamed.schema == pa.schema([("new_id", pa.int64()), ("value", pa.int64())]) + batch = renamed.with_format("arrow").__getitems__([0]) + assert batch.schema.names == ["new_id", "value"] + + def test_rename_columns(some_permutation: Permutation): assert some_permutation.rename_columns({"value": "new_value"}).schema == pa.schema( [("id", pa.int64()), ("new_value", pa.int64())] From 631ea41a635890218bf7ee1ab80df6caccc7606c Mon Sep 17 00:00:00 2001 From: Colin Patrick McCabe Date: Thu, 22 Oct 2026 15:41:30 -0700 Subject: [PATCH 03/36] feat!: return lazy iterators for table and namespace listings (#4599) Make table and namespace listings consistent with other resource listings across Rust, Python, and TypeScript. Callers can iterate over individual names across all pages, with configurable page sizes and continuation tokens. In order to avoid breaking existing Python users, this change adds some compatibility methods to the returned Python iterators. Code such as: ``` for t in db.list_tables().tables: print(t) ``` will continue to work as expected because the iterator has a `tables` field and a `page_token` field, which function as expected to reveal the first page of results. In other words, it is never necessary to use the iterator __as an iterator__. It may be used as if it were an old-style response object. (As the iterator progresses, the `tables` list may switch to the second, third, etc. page of results.) Similar compatibility code allows the async python API to keep working as well. --- docs/src/js/classes/Connection.md | 80 ++------ docs/src/js/globals.md | 4 - .../js/interfaces/ListNamespacesOptions.md | 27 --- .../js/interfaces/ListNamespacesResponse.md | 23 --- docs/src/js/interfaces/ListTablesOptions.md | 35 ---- docs/src/js/interfaces/ListTablesResponse.md | 23 --- docs/src/js/interfaces/TableNamesOptions.md | 2 +- docs/src/python/python.md | 23 ++- docs/web-tests/ts/connection.test.ts | 16 +- docs/web-tests/ts/quickstart.test.ts | 7 +- nodejs/__test__/connection.test.ts | 156 ++++++++-------- nodejs/lancedb/connection.ts | 133 ++++---------- nodejs/lancedb/index.ts | 4 - nodejs/src/connection.rs | 101 ++--------- python/python/lancedb/_lancedb.pyi | 11 +- python/python/lancedb/db.py | 70 ++++--- python/python/lancedb/listing.py | 171 ++++++++++++++++++ python/python/lancedb/namespace.py | 39 ++-- python/python/lancedb/remote/db.py | 16 +- python/python/tests/test_embeddings.py | 4 + python/python/tests/test_listing.py | 137 ++++++++++++++ python/python/tests/test_ollama_embeddings.py | 2 +- python/src/connection.rs | 56 ++---- python/src/listing.rs | 30 ++- rust/lancedb/src/connection.rs | 161 +++++++++-------- rust/lancedb/src/listing.rs | 130 +++++++++---- rust/lancedb/src/remote/db.rs | 66 ++----- 27 files changed, 814 insertions(+), 713 deletions(-) delete mode 100644 docs/src/js/interfaces/ListNamespacesOptions.md delete mode 100644 docs/src/js/interfaces/ListNamespacesResponse.md delete mode 100644 docs/src/js/interfaces/ListTablesOptions.md delete mode 100644 docs/src/js/interfaces/ListTablesResponse.md diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index d5ab8c3d1..c110eb38c 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -653,94 +653,48 @@ See [Listing](Listing.md) for pagination state and error behavior. ### listNamespaces() ```ts -abstract listNamespaces(namespacePath?, options?): Promise +abstract listNamespaces(namespacePath?, options?): Listing ``` -List the immediate child namespaces under the given parent. - -Results may be paginated. To retrieve subsequent pages, pass the -`pageToken` returned by a previous call. +Lazily list immediate child namespaces, 100 per page by default. #### Parameters * **namespacePath?**: `string`[] - The parent namespace path. Defaults - to the root namespace if omitted. -* **options?**: `Partial`<[`ListNamespacesOptions`](../interfaces/ListNamespacesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) #### Returns -`Promise`<[`ListNamespacesResponse`](../interfaces/ListNamespacesResponse.md)> - -Child namespace names and - an optional token for fetching the next page. +[`Listing`](Listing.md)<`string`> *** ### listTables() -#### listTables(options) - ```ts -abstract listTables(options?): Promise +abstract listTables(namespacePath?, options?): Listing ``` -List a page of the tables in this database. +Lazily list tables in a namespace, requesting 100 per page by default. -To retrieve the tables after the page, pass the `pageToken` the response -carries back in. A page can be shorter than `limit` without being the last -one, so walk until a response carries no page token: - -```ts -const names = []; -let pageToken = undefined; -do { - const page = await conn.listTables({ pageToken, limit: 100 }); - names.push(...page.tables); - pageToken = page.pageToken; -} while (pageToken); -``` - -##### Parameters - -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). - -##### Returns - -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> - -A page of table names and an - optional token for the tables after it. - -#### listTables(namespacePath, options) - -```ts -abstract listTables(namespacePath?, options?): Promise -``` - -List a page of the tables in this database. - -##### Parameters +#### Parameters * **namespacePath?**: `string`[] - The namespace path to list tables from - (defaults to root namespace) -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) -##### Returns +#### Returns -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> +[`Listing`](Listing.md)<`string`> -A page of table names and an - optional token for the tables after it. +#### Example + +```ts +for await (const name of conn.listTables([], { pageLimit: 20 })) { + console.log(name); +} +``` *** diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 646341190..e3830e9ad 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -108,10 +108,6 @@ - [JobEventsOptions](interfaces/JobEventsOptions.md) - [JobFailureInfo](interfaces/JobFailureInfo.md) - [JobInfo](interfaces/JobInfo.md) -- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) -- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) -- [ListTablesOptions](interfaces/ListTablesOptions.md) -- [ListTablesResponse](interfaces/ListTablesResponse.md) - [ListingOptions](interfaces/ListingOptions.md) - [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md) - [MemtableStats](interfaces/MemtableStats.md) diff --git a/docs/src/js/interfaces/ListNamespacesOptions.md b/docs/src/js/interfaces/ListNamespacesOptions.md deleted file mode 100644 index 79680417f..000000000 --- a/docs/src/js/interfaces/ListNamespacesOptions.md +++ /dev/null @@ -1,27 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesOptions - -# Interface: ListNamespacesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An optional limit to the number of results to return. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response for pagination. diff --git a/docs/src/js/interfaces/ListNamespacesResponse.md b/docs/src/js/interfaces/ListNamespacesResponse.md deleted file mode 100644 index a99ab480f..000000000 --- a/docs/src/js/interfaces/ListNamespacesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesResponse - -# Interface: ListNamespacesResponse - -## Properties - -### namespaces - -```ts -namespaces: string[]; -``` - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` diff --git a/docs/src/js/interfaces/ListTablesOptions.md b/docs/src/js/interfaces/ListTablesOptions.md deleted file mode 100644 index b8dd8b5ef..000000000 --- a/docs/src/js/interfaces/ListTablesOptions.md +++ /dev/null @@ -1,35 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesOptions - -# Interface: ListTablesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An upper bound on how many tables to return. -Defaults to 100. Zero returns an empty page without a continuation token. - -A page may hold fewer than this and still not be the last one, so keep -going while the response carries a page token rather than while pages are -full. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response, to resume listing where it left off. - -The token is opaque: it carries whatever the database needs to resume, and -callers should not construct or interpret one. diff --git a/docs/src/js/interfaces/ListTablesResponse.md b/docs/src/js/interfaces/ListTablesResponse.md deleted file mode 100644 index 76cac2b23..000000000 --- a/docs/src/js/interfaces/ListTablesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesResponse - -# Interface: ListTablesResponse - -## Properties - -### pageToken? - -```ts -optional pageToken: string; -``` - -*** - -### tables - -```ts -tables: string[]; -``` diff --git a/docs/src/js/interfaces/TableNamesOptions.md b/docs/src/js/interfaces/TableNamesOptions.md index 123aa6c9f..e1fd8c382 100644 --- a/docs/src/js/interfaces/TableNamesOptions.md +++ b/docs/src/js/interfaces/TableNamesOptions.md @@ -8,7 +8,7 @@ ## Deprecated -Use [ListTablesOptions](ListTablesOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) +Use [ListingOptions](ListingOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) instead. ## Properties diff --git a/docs/src/python/python.md b/docs/src/python/python.md index f4c697497..c465d0ac4 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -26,14 +26,17 @@ is also an [asynchronous API client](#connections-asynchronous). ::: lancedb.db.DBConnection -Resource listings (`list_secrets`, `list_views`, `list_jobs`, `list_functions`, and -`list_materialized_views`) return lazy iterators. Use `list(db.list_views())` to -collect synchronous results, or `[name async for name in db.list_views()]` with -an asynchronous connection. `page_limit` controls each request, and `page_token` -resumes a listing. Drain the iterator's cached results before saving its token. +Resource listings return iterators. Use `list(db.list_tables())` to collect all +table names. `db.list_tables().tables` is a list of the first page; `limit` +controls its page size and `page_token` resumes at a saved token. The same +patterns apply to `list_namespaces()`. Other resource listings use `page_limit`. ::: lancedb.listing.Listing +::: lancedb.listing.TableListing + +::: lancedb.listing.NamespaceListing + ::: lancedb.Session ## Catalogs (Synchronous) @@ -631,6 +634,16 @@ can be used to create, list, or open tables. ::: lancedb.listing.AsyncListing +`async for name in db.list_tables()` loads pages as needed. To inspect the +first page, use `page = await db.list_tables()` and then `page.tables`. Reading +the page list does not consume the iterator. `page.page_token` is the token for +the next request; drain cached results before using it to resume in a new +listing. The same behavior applies to namespaces. + +::: lancedb.listing.AsyncTableListing + +::: lancedb.listing.AsyncNamespaceListing + ## Namespaces (Asynchronous) ::: lancedb.connect_namespace_async diff --git a/docs/web-tests/ts/connection.test.ts b/docs/web-tests/ts/connection.test.ts index 891501d4c..651be7d36 100644 --- a/docs/web-tests/ts/connection.test.ts +++ b/docs/web-tests/ts/connection.test.ts @@ -4,6 +4,12 @@ import { afterEach, describe, expect, jest, test } from "@jest/globals"; import * as lancedb from "@lancedb/lancedb"; import { withTempDirectory } from "./util.ts"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + test("connect to a local database", async () => { await withTempDirectory(async (tempDir) => { const cwd = process.cwd(); @@ -13,12 +19,12 @@ test("connect to a local database", async () => { const db = await lancedb.connect("ex_lancedb"); // --8<-- [end:connect] await db.createTable("probe", [{ id: 1 }]); - expect((await db.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(db.listTables())).toEqual(["probe"]); } finally { process.chdir(cwd); } const reopened = await lancedb.connect(`${tempDir}/ex_lancedb`); - expect((await reopened.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(reopened.listTables())).toEqual(["probe"]); }); }); @@ -109,8 +115,8 @@ async function namespaceTableOpsExample() { ); // Verify - console.log((await db.listNamespaces()).namespaces); // ["prod"] - console.log((await db.listNamespaces(["prod"])).namespaces); // ["recommendations", "search"] + console.log(await collectNames(db.listNamespaces())); // ["prod"] + console.log(await collectNames(db.listNamespaces(["prod"]))); // ["recommendations", "search"] console.log(await db.tableNames(["prod", "search"])); // ["user"] console.log(await db.tableNames(["prod", "recommendations"])); // ["user"] // --8<-- [end:namespace_table_ops] @@ -124,7 +130,7 @@ async function namespaceAdminOpsExample() { await db.createNamespace(["prod"]); await db.createNamespace(["prod", "search"]); - const childNamespaces = (await db.listNamespaces(["prod"])).namespaces; + const childNamespaces = await collectNames(db.listNamespaces(["prod"])); console.log(`Child namespaces under ${JSON.stringify(namespace)}:`, childNamespaces); // Child namespaces under ["prod","search"]: [ 'search' ] diff --git a/docs/web-tests/ts/quickstart.test.ts b/docs/web-tests/ts/quickstart.test.ts index e32a14f86..2c3ee8252 100644 --- a/docs/web-tests/ts/quickstart.test.ts +++ b/docs/web-tests/ts/quickstart.test.ts @@ -20,12 +20,7 @@ const connections = async function tableNames(db: lancedb.Connection) { const names: string[] = []; - let pageToken: string | undefined; - do { - const page = await db.listTables({ pageToken }); - names.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); + for await (const name of db.listTables()) names.push(name); return names; } diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index bbd0c3c52..73f99803a 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -5,15 +5,15 @@ import { mkdirSync, readdirSync } from "fs"; import { join } from "path"; import { Field, Float64, Schema } from "apache-arrow"; import * as tmp from "tmp"; -import { - Connection, - ListTablesResponse, - Table, - connect, - connectNamespace, -} from "../lancedb"; +import { Connection, Table, connect, connectNamespace } from "../lancedb"; import { LocalTable } from "../lancedb/table"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + describe("when connecting", () => { let tmpDir: tmp.DirResult; beforeEach(() => { @@ -135,7 +135,7 @@ describe("given a connection", () => { await db.close(); expect(db.isOpen()).toBe(false); await expect(db.tableNames()).rejects.toThrow("Connection is closed"); - await expect(db.listTables()).rejects.toThrow("Connection is closed"); + expect(() => db.listTables()).toThrow("Connection is closed"); await expect(db.renameTable("a", "b")).rejects.toThrow( "Connection is closed", ); @@ -247,15 +247,15 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const first = await db.listTables(namespacePath, { limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables(namespacePath, { pageLimit: 1 }); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables(namespacePath, { - limit: 1, - pageToken: first.pageToken, + const second = db.listTables(namespacePath, { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); }); it("should keep the existing table when mode is exist_ok", async () => { @@ -299,19 +299,18 @@ describe("given a connection", () => { await db.createTable("a", [{ id: 1 }]); await db.createTable("c", [{ id: 1 }]); - const all = await db.listTables(); - expect(all.tables).toEqual(["a", "b", "c"]); - expect(all.pageToken).toBeUndefined(); + expect(await collectNames(db.listTables())).toEqual(["a", "b", "c"]); - const first = await db.listTables({ limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables({ - limit: 1, - pageToken: first.pageToken, + const second = db.listTables([], { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); it("should visit every table exactly once when walking pages", async () => { @@ -322,18 +321,9 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const seen: string[] = []; - let pageToken: string | undefined = undefined; - do { - const page: ListTablesResponse = await db.listTables({ - limit: 2, - pageToken, - }); - seen.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); - - expect(seen).toEqual(created); + expect(await collectNames(db.listTables([], { pageLimit: 2 }))).toEqual( + created, + ); }); it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])( @@ -347,11 +337,10 @@ describe("given a connection", () => { ]) { await expect(names).resolves.toEqual(["a", "b"]); } - for (const page of [ - db.listTables({ limit }), - db.listTables([], { limit }), - ]) { - await expect(page).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + expect(() => + db.listTables(namespacePath, { pageLimit: limit }), + ).toThrow(/limit/); } }, ); @@ -366,8 +355,11 @@ describe("given a connection", () => { ])("should reject an invalid table listing limit %s", async (limit) => { await expect(db.tableNames({ limit })).rejects.toThrow(/limit/); await expect(db.tableNames([], { limit })).rejects.toThrow(/limit/); - await expect(db.listTables({ limit })).rejects.toThrow(/limit/); - await expect(db.listTables([], { limit })).rejects.toThrow(/limit/); + for (const namespacePath of [undefined, []]) { + expect(() => db.listTables(namespacePath, { pageLimit: limit })).toThrow( + /limit/, + ); + } }); it("should preserve zero, small and unsigned-32-bit table listing limits", async () => { @@ -379,11 +371,21 @@ describe("given a connection", () => { [2 ** 32 - 1, ["a", "b"]], ] as const) { await expect(db.tableNames({ limit })).resolves.toEqual(names); - const page = await db.listTables({ limit }); - expect(page.tables).toEqual(names); } await expect(db.tableNames()).resolves.toEqual(["a", "b"]); - await expect(db.listTables()).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + for (const pageLimit of [1, 2 ** 31 - 1]) { + expect( + await collectNames(db.listTables(namespacePath, { pageLimit })), + ).toEqual(["a", "b"]); + } + for (const pageLimit of [0, 2 ** 31, 2 ** 32 - 1]) { + expect(() => db.listTables(namespacePath, { pageLimit })).toThrow( + /limit/, + ); + } + } + expect(await collectNames(db.listTables())).toEqual(["a", "b"]); }); it("should list all names and use a default page size of 100", async () => { @@ -396,20 +398,17 @@ describe("given a connection", () => { } await expect(db.tableNames()).resolves.toEqual(names); await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]); - const first = await db.listTables(); - expect(first.tables).toEqual(names.slice(0, 100)); - expect(first.pageToken).toBeDefined(); - const second = await db.listTables({ pageToken: first.pageToken }); - expect(second.tables).toEqual(names.slice(100)); - expect(second.pageToken).toBeUndefined(); - await expect(db.listTables({ limit: 200 })).resolves.toEqual({ - tables: names, - pageToken: undefined, - }); - await expect(db.listTables({ limit: 0 })).resolves.toEqual({ - tables: [], - pageToken: undefined, - }); + const first = db.listTables(); + const firstPage: string[] = []; + for (let i = 0; i < 100; i++) firstPage.push((await first.next()).value!); + expect(firstPage).toEqual(names.slice(0, 100)); + expect(first.pageToken()).toBeDefined(); + expect(first.numPageResults()).toBe(0); + expect(await collectNames(first)).toEqual(names.slice(100)); + expect(await collectNames(db.listTables([], { pageLimit: 200 }))).toEqual( + names, + ); + expect(() => db.listTables([], { pageLimit: 0 })).toThrow(/limit/); }); it("should list tables in a namespace", async () => { @@ -420,12 +419,8 @@ describe("given a connection", () => { await db.createNamespace(["child"]); await db.createTable("nested", [{ id: 1 }], ["child"]); - await expect(db.listTables(["child"])).resolves.toEqual( - expect.objectContaining({ tables: ["nested"] }), - ); - await expect(db.listTables()).resolves.toEqual( - expect.objectContaining({ tables: [] }), - ); + expect(await collectNames(db.listTables(["child"]))).toEqual(["nested"]); + expect(await collectNames(db.listTables())).toEqual([]); }); it("should create tables in v2 mode", async () => { @@ -719,22 +714,35 @@ describe("namespaces", () => { it("should list namespaces created at the root", async () => { await db.createNamespace(["alpha"]); await db.createNamespace(["beta"]); - const list = await db.listNamespaces(); - expect(list.namespaces).toEqual(expect.arrayContaining(["alpha", "beta"])); + const list = await collectNames(db.listNamespaces()); + expect(list).toEqual(expect.arrayContaining(["alpha", "beta"])); }); it("should list child namespaces under a parent", async () => { await db.createNamespace(["parent"]); await db.createNamespace(["parent", "child"]); - const list = await db.listNamespaces(["parent"]); - expect(list.namespaces).toContain("child"); + const list = await collectNames(db.listNamespaces(["parent"])); + expect(list).toContain("child"); + }); + + it("should resume a namespace listing from its page token", async () => { + await db.createNamespace(["alpha"]); + await db.createNamespace(["beta"]); + const first = db.listNamespaces([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect((await first.next()).value).toBe("alpha"); + const second = db.listNamespaces([], { + pageLimit: 1, + pageToken: first.pageToken(), + }); + expect(await collectNames(second)).toEqual(["beta"]); }); it("should drop a namespace", async () => { await db.createNamespace(["ephemeral"]); await db.dropNamespace(["ephemeral"]); - const list = await db.listNamespaces(); - expect(list.namespaces).not.toContain("ephemeral"); + const list = await collectNames(db.listNamespaces()); + expect(list).not.toContain("ephemeral"); }); it("should raise an error on any namespace op after close", async () => { @@ -742,7 +750,7 @@ describe("namespaces", () => { await expect(db.describeNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); - await expect(db.listNamespaces()).rejects.toThrow("Connection is closed"); + expect(() => db.listNamespaces()).toThrow("Connection is closed"); await expect(db.createNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index c93414fab..5f067eeef 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -30,16 +30,12 @@ import type { DescribeNamespaceResponse, DropNamespaceResponse, JobInfo, - ListNamespacesResponse, - ListTablesResponse, } from "./native"; import { ViewDescription, viewDescriptionFromNative } from "./view"; export type { CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, }; import { sanitizeTable } from "./sanitize"; import { LocalTable, Table } from "./table"; @@ -141,7 +137,7 @@ export interface OpenTableOptions { } /** - * @deprecated Use {@link ListTablesOptions} with {@link Connection.listTables} + * @deprecated Use {@link ListingOptions} with {@link Connection.listTables} * instead. */ export interface TableNamesOptions { @@ -157,32 +153,6 @@ export interface TableNamesOptions { limit?: number; } -export interface ListTablesOptions { - /** - * Token from a previous response, to resume listing where it left off. - * - * The token is opaque: it carries whatever the database needs to resume, and - * callers should not construct or interpret one. - */ - pageToken?: string; - /** - * An upper bound on how many tables to return. - * Defaults to 100. Zero returns an empty page without a continuation token. - * - * A page may hold fewer than this and still not be the last one, so keep - * going while the response carries a page token rather than while pages are - * full. - */ - limit?: number; -} - -export interface ListNamespacesOptions { - /** Token from a previous response for pagination. */ - pageToken?: string; - /** An optional limit to the number of results to return. */ - limit?: number; -} - export interface CreateNamespaceOptions { /** Creation mode. */ mode?: "create" | "exist_ok" | "overwrite"; @@ -278,45 +248,18 @@ export abstract class Connection { options?: Partial, ): Promise; - /** - * List a page of the tables in this database. - * - * To retrieve the tables after the page, pass the `pageToken` the response - * carries back in. A page can be shorter than `limit` without being the last - * one, so walk until a response carries no page token: - * + /** Lazily list tables in a namespace, requesting 100 per page by default. + * @example * ```ts - * const names = []; - * let pageToken = undefined; - * do { - * const page = await conn.listTables({ pageToken, limit: 100 }); - * names.push(...page.tables); - * pageToken = page.pageToken; - * } while (pageToken); + * for await (const name of conn.listTables([], { pageLimit: 20 })) { + * console.log(name); + * } * ``` - * - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. - */ - abstract listTables( - options?: Partial, - ): Promise; - /** - * List a page of the tables in this database. - * - * @param {string[]} namespacePath - The namespace path to list tables from - * (defaults to root namespace) - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. */ abstract listTables( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Open a table in the database. @@ -542,23 +485,11 @@ export abstract class Connection { namespacePath: string[], ): Promise; - /** - * List the immediate child namespaces under the given parent. - * - * Results may be paginated. To retrieve subsequent pages, pass the - * `pageToken` returned by a previous call. - * - * @param {string[]} namespacePath - The parent namespace path. Defaults - * to the root namespace if omitted. - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} Child namespace names and - * an optional token for fetching the next page. - */ + /** Lazily list immediate child namespaces, 100 per page by default. */ abstract listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Create a new namespace at the given path. @@ -821,22 +752,17 @@ export class LocalConnection extends Connection { ); } - async listTables( - namespacePathOrOptions?: string[] | Partial, - options?: Partial, - ): Promise { - // Detect if first argument is namespacePath array or options object - const namespacePath = Array.isArray(namespacePathOrOptions) - ? namespacePathOrOptions - : undefined; - const listTablesOptions = Array.isArray(namespacePathOrOptions) - ? options - : (namespacePathOrOptions ?? options); - - return this.inner.listTables( - namespacePath ?? [], - listTablesOptions?.pageToken, - normalizeTableListingLimit(listTablesOptions?.limit), + listTables( + namespacePath?: string[], + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listTables( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } @@ -1039,12 +965,15 @@ export class LocalConnection extends Connection { listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise { - return this.inner.listNamespaces( - namespacePath ?? [], - options?.pageToken, - options?.limit, + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listNamespaces( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 4ad8b2a93..0f1b5f333 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -97,13 +97,9 @@ export { Connection, CreateTableOptions, TableNamesOptions, - ListTablesOptions, OpenTableOptions, - ListNamespacesOptions, CreateNamespaceOptions, DropNamespaceOptions, - ListNamespacesResponse, - ListTablesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 67c4706e4..914f9e4c9 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -16,8 +16,7 @@ use crate::table::Table; use lancedb::connection::{ConnectBuilder, Connection as LanceDBConnection, connect_namespace}; use lance_namespace::models::{ - CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest, - ListTablesRequest, + CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, }; use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema}; @@ -31,18 +30,6 @@ pub struct DescribeNamespaceResponse { pub properties: Option>, } -#[napi(object)] -pub struct ListNamespacesResponse { - pub namespaces: Vec, - pub page_token: Option, -} - -#[napi(object)] -pub struct ListTablesResponse { - pub tables: Vec, - pub page_token: Option, -} - #[napi(object)] pub struct CreateNamespaceResponse { pub properties: Option>, @@ -239,31 +226,20 @@ impl Connection { op.execute().await.default_error() } - /// List a page of tables in the database. + /// Lazily list tables in the database. #[napi(catch_unwind)] - pub async fn list_tables( + pub fn list_tables( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let request = ListTablesRequest { - // The root namespace is an empty path, not an absent one: a namespace-backed - // database rejects a request that names no namespace. - id: Some(namespace_path.unwrap_or_default()), - page_token, - limit: limit.map(|limit| i32::try_from(limit).unwrap_or(i32::MAX)), - ..Default::default() - }; - let response = self - .get_inner()? - .list_tables(request) - .await - .default_error()?; - Ok(ListTablesResponse { - tables: response.tables, - page_token: response.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Create table from a Apache Arrow IPC (file) buffer. @@ -663,28 +639,19 @@ impl Connection { } #[napi(catch_unwind)] - /// List child namespaces under the given namespace path - pub async fn list_namespaces( + /// Lazily list child namespaces under the given namespace path + pub fn list_namespaces( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let req = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: checked_limit(limit)?, - ..Default::default() - }; - let resp = self - .get_inner()? - .list_namespaces(req) - .await - .default_error()?; - Ok(ListNamespacesResponse { - namespaces: resp.namespaces, - page_token: resp.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[napi(catch_unwind)] @@ -787,33 +754,3 @@ impl Connection { .default_error() } } - -fn checked_limit(limit: Option) -> napi::Result> { - limit - .map(i32::try_from) - .transpose() - .map_err(|_| napi::Error::from_reason("limit must not exceed 2147483647")) -} - -#[cfg(test)] -mod tests { - use super::checked_limit; - - #[test] - fn checked_limit_rejects_values_above_i32_max() { - // u32 values above i32::MAX wrapped negative through `as i32`, - // sending a negative page size to the server. - assert!(checked_limit(Some(3_000_000_000)).is_err()); - assert!(checked_limit(Some(u32::MAX)).is_err()); - } - - #[test] - fn checked_limit_passes_boundary_and_missing_values() { - assert_eq!( - checked_limit(Some(i32::MAX as u32)).unwrap(), - Some(i32::MAX) - ); - assert_eq!(checked_limit(Some(0)).unwrap(), Some(0)); - assert_eq!(checked_limit(None).unwrap(), None); - } -} diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 96eef064c..6ef23abc7 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -19,11 +19,9 @@ from .index import ( FTS, ) from lance_namespace import ( - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, ) from .remote import ClientConfig from .types import BaseTokenizerType @@ -114,12 +112,12 @@ class Connection(object): uri: str async def is_open(self): ... async def close(self): ... - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: ... + ) -> NameListing: ... async def create_namespace( self, namespace_path: List[str], @@ -136,12 +134,12 @@ class Connection(object): self, namespace_path: List[str], ) -> DescribeNamespaceResponse: ... - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: ... + ) -> NameListing: ... async def table_names( self, namespace_path: Optional[List[str]], @@ -1010,6 +1008,7 @@ async def connect_catalog( class NameListing: def num_page_results(self) -> int: ... def page_token(self) -> Optional[str]: ... + async def next_page(self) -> Optional[dict[str, Any]]: ... def __aiter__(self) -> NameListing: ... async def __anext__(self) -> str: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index a74763afb..90dc1cccf 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -32,18 +32,24 @@ if sys.version_info >= (3, 12): else: from overrides import EnforceOverrides, override # type: ignore -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, + validate_page_limit, +) from lancedb.embeddings.registry import EmbeddingFunctionRegistry from lancedb.common import data_to_reader, sanitize_uri, validate_schema from lancedb.background_loop import LOOP from lance_namespace import ( LanceNamespace, - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, connect as namespace_connect, ) from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError @@ -171,7 +177,7 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -187,12 +193,12 @@ class DBConnection(EnforceOverrides): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return ListNamespacesResponse(namespaces=[], page_token=None) + return NamespaceListing(AsyncNamespaceListing(None, empty=True)) def create_namespace( self, @@ -317,8 +323,8 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -334,7 +340,7 @@ class DBConnection(EnforceOverrides): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ raise NotImplementedError( @@ -1436,7 +1442,7 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -1452,16 +1458,14 @@ class LanceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( - namespace_path=namespace_path, - page_token=page_token, - limit=limit, + namespace_path=namespace_path, page_token=page_token, limit=limit ) ) @@ -1512,8 +1516,8 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -1529,12 +1533,12 @@ class LanceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -2286,12 +2290,12 @@ class AsyncConnection(object): else: return None - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -2306,15 +2310,18 @@ class AsyncConnection(object): Returns ------- - ListNamespacesResponse + AsyncNamespaceListing Response containing namespace names and optional pagination token """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_namespaces( + validate_page_limit(limit) + inner = self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListNamespacesResponse(**result) + if limit == 0: + return AsyncNamespaceListing(None, empty=True) + return AsyncNamespaceListing(inner, page_token) async def create_namespace( self, @@ -2393,13 +2400,13 @@ class AsyncConnection(object): result = await self._inner.describe_namespace(namespace_path) return DescribeNamespaceResponse(**result) - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> AsyncTableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -2415,15 +2422,18 @@ class AsyncConnection(object): Returns ------- - ListTablesResponse + AsyncTableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_tables( + validate_page_limit(limit) + inner = self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListTablesResponse(**result) + if limit == 0: + return AsyncTableListing(None, empty=True) + return AsyncTableListing(inner, page_token) async def table_names( self, diff --git a/python/python/lancedb/listing.py b/python/python/lancedb/listing.py index 0620aa7d8..8abc9d85b 100644 --- a/python/python/lancedb/listing.py +++ b/python/python/lancedb/listing.py @@ -5,6 +5,7 @@ from __future__ import annotations +import asyncio from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar from .background_loop import LOOP @@ -83,3 +84,173 @@ class Listing(Iterator[T], Generic[T]): for initial, cached, final-page, and error behavior. """ return self._inner.page_token() + + +def validate_page_limit(limit: Optional[int]) -> None: + """Validate the legacy limit argument before passing it to the bindings.""" + if limit is not None and ( + not isinstance(limit, int) or limit < 0 or limit > 2**31 - 1 + ): + raise ValueError("limit must be between 0 and 2147483647") + + +class _AsyncPagedNameListing(AsyncIterator[str]): + """An async listing with an inspectable copy of its current backend page.""" + + def __init__( + self, inner: Any, page_token: Optional[str] = None, *, empty: bool = False + ): + self._inner = inner + self._token = page_token + self._page: Optional[list[str]] = [] if empty else None + self._position = 0 + self._done = empty + self._failure: Optional[BaseException] = None + self._lock = asyncio.Lock() + + @property + def page_token(self) -> Optional[str]: + """The continuation token for the next backend request.""" + return self._token + + def num_page_results(self) -> int: + """Count the remaining names in the current page.""" + return 0 if self._page is None else len(self._page) - self._position + + def _page_copy(self) -> list[str]: + if self._page is None: + raise RuntimeError("Await or iterate the listing before reading its page") + return self._page.copy() + + async def _fetch_page(self) -> None: + try: + page = await self._inner.next_page() + except asyncio.CancelledError: + raise + except BaseException as exc: + self._failure = exc + self._done = True + raise + if page is None: + self._page = [] + self._token = None + self._done = True + else: + self._page = page["items"] + self._token = page["page_token"] + self._done = self._token is None + self._position = 0 + + async def _load_initial(self): + async with self._lock: + if self._failure is not None: + raise self._failure + if self._page is None: + await self._fetch_page() + return self + + def __await__(self): + return self._load_initial().__await__() + + def __aiter__(self): + return self + + async def __anext__(self) -> str: + async with self._lock: + if self._failure is not None: + raise StopAsyncIteration + if self._page is None: + await self._fetch_page() + assert self._page is not None + while self._position >= len(self._page): + if self._done: + raise StopAsyncIteration + await self._fetch_page() + assert self._page is not None + item = self._page[self._position] + self._position += 1 + return item + + +class AsyncTableListing(_AsyncPagedNameListing): + """Awaitable table iterator; awaiting loads one page. + + For example:: + + page = await db.list_tables(limit=20) + print(page.tables, page.page_token) + async for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + """A copy of every table name in the most recently fetched page.""" + return self._page_copy() + + +class AsyncNamespaceListing(_AsyncPagedNameListing): + """Awaitable namespace iterator; awaiting loads one page. + + For example:: + + page = await db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + """A copy of every name in the most recently fetched page.""" + return self._page_copy() + + +class _PagedNameListing(Iterator[str]): + def __init__(self, inner: _AsyncPagedNameListing): + self._inner = inner + LOOP.run(inner._load_initial()) + + def __iter__(self): + return self + + def __next__(self) -> str: + try: + return LOOP.run(self._inner.__anext__()) + except StopAsyncIteration: + raise StopIteration from None + + @property + def page_token(self) -> Optional[str]: + return self._inner.page_token + + def num_page_results(self) -> int: + return self._inner.num_page_results() + + +class TableListing(_PagedNameListing): + """Table iterator retaining the current page as a list. + + For example:: + + page = db.list_tables(limit=20) + print(page.tables, page.page_token) + for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + return self._inner.tables + + +class NamespaceListing(_PagedNameListing): + """Namespace iterator retaining the current page as a list. + + For example:: + + page = db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + return self._inner.namespaces diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index a5d88eb1f..55baa39a3 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -52,7 +52,14 @@ from lancedb.background_loop import LOOP from lancedb.arrow import AsyncRecordBatchReader from lancedb.db import AsyncConnection, DBConnection from lancedb.job import AsyncJob, Job -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, +) from lancedb.sql import AsyncQuery as AsyncSqlQuery from lancedb.sql import QueryDescription, QueryParameters from lance_namespace import ( @@ -61,8 +68,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, NamespaceExistsRequest, TableExistsRequest, ) @@ -854,7 +859,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """ List child namespaces under the given namespace. @@ -871,12 +876,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1023,7 +1028,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """ List all tables in this database with pagination support. @@ -1040,12 +1045,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1459,12 +1464,12 @@ class AsyncLanceNamespaceDBConnection: namespace_path = [] await self._inner.drop_all_tables(namespace_path=namespace_path) - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """ List child namespaces under the given namespace. @@ -1481,12 +1486,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_namespaces( + return self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1616,12 +1621,12 @@ class AsyncLanceNamespaceDBConnection: except TableNotFoundError: return False - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> AsyncTableListing: """ List all tables in this database with pagination support. @@ -1638,12 +1643,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_tables( + return self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 89c7f8a78..08125cc42 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -36,7 +36,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig import pyarrow as pa from ..common import DATA -from lancedb.listing import Listing +from lancedb.listing import Listing, NamespaceListing, TableListing from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job @@ -54,8 +54,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, ) from ..pydantic import LanceModel from ..table import Table @@ -268,7 +266,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -284,12 +282,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -379,7 +377,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """List a page of tables in this database. Parameters @@ -396,12 +394,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index e319d3048..207a3d253 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -896,6 +896,8 @@ def test_transient_embedding_error_is_retried_in_table_add(tmp_path): assert len(calls) == 3 assert mock_sleep.call_count == 2 assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list) + + @patch("time.sleep") @pytest.mark.parametrize( "error", @@ -941,6 +943,8 @@ def test_retry_on_wrapped_provider_error(mock_sleep): assert wrapped_func() == [0.1] assert mock_func.call_count == 2 assert mock_sleep.call_count == 1 + + def test_openai_bad_request_partial_batch_recovery(): # Regression test for issue #1677 diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py index 4a2745eb8..b0b973bbf 100644 --- a/python/python/tests/test_listing.py +++ b/python/python/tests/test_listing.py @@ -88,6 +88,143 @@ async def advance(items, asynchronous): KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"] +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_compatible_name_listing_pages(listing_server, kind, asynchronous): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a", "b"], "next")), + (200, response_page(kind, ["c"], None)), + ] + ) + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)( + namespace_path=["team"], page_token="start", limit=2 + ) + field = kind + assert len(requests) == (0 if asynchronous else 1) + if asynchronous: + with pytest.raises(RuntimeError, match="Await or iterate"): + getattr(listing, field) + assert listing.page_token == "start" + assert await listing is listing + assert len(requests) == 1 + page_names = getattr(listing, field) + assert type(page_names) is list + assert page_names == ["a", "b"] + page_names.append("changed") + assert getattr(listing, field) == ["a", "b"] + assert listing.page_token == "next" + assert listing.num_page_results() == 2 + assert await advance(listing, asynchronous) == "a" + assert getattr(listing, field) == ["a", "b"] + assert await advance(listing, asynchronous) == "b" + assert len(requests) == 1 + assert await advance(listing, asynchronous) == "c" + assert getattr(listing, field) == ["c"] + assert listing.page_token is None + assert len(requests) == 2 + assert ([item async for item in listing] if asynchronous else list(listing)) == [] + if asynchronous: + assert await listing is listing + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_skips_empty_page_only_during_iteration( + listing_server, kind +): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, [], "next")), + (200, response_page(kind, ["a"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + await listing + assert getattr(listing, kind) == [] + assert listing.page_token == "next" + assert len(requests) == 1 + assert [name async for name in listing] == ["a"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_name_listing_zero_limit_uses_no_request( + listing_server, kind, asynchronous +): + endpoint, requests, _ = listing_server + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)(limit=0) + if asynchronous: + assert await listing is listing + assert [name async for name in listing] == [] + else: + assert list(listing) == [] + assert getattr(listing, kind) == [] + assert listing.page_token is None + assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_failure_retains_token(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((401, {})) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(page_token="retry") + with pytest.raises(HttpError): + await listing + assert listing.page_token == "retry" + with pytest.raises(HttpError): + await listing + assert [name async for name in listing] == [] + assert len(requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_direct_async_iteration(listing_server, kind): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a"], "next")), + (200, response_page(kind, ["b"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(limit=1) + assert requests == [] + assert [name async for name in listing] == ["a", "b"] + assert getattr(listing, kind) == ["b"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_serializes_concurrent_advances(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((200, response_page(kind, ["a", "b"], None))) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + results = await asyncio.gather( + listing.__anext__(), + listing.__anext__(), + listing.__anext__(), + return_exceptions=True, + ) + assert results[:2] == ["a", "b"] + assert isinstance(results[2], StopAsyncIteration) + assert len(requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", KINDS) @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/python/python/tests/test_ollama_embeddings.py b/python/python/tests/test_ollama_embeddings.py index 9d8568ca7..de36dc24b 100644 --- a/python/python/tests/test_ollama_embeddings.py +++ b/python/python/tests/test_ollama_embeddings.py @@ -3,7 +3,6 @@ import json -import httpx import numpy as np import pyarrow as pa import pytest @@ -14,6 +13,7 @@ from lancedb.embeddings import get_registry @pytest.fixture def ollama_embeddings(): + httpx = pytest.importorskip("httpx") pytest.importorskip("ollama") requests = [] diff --git a/python/src/connection.rs b/python/src/connection.rs index baeb29780..2cdf71f4a 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -590,29 +590,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_namespaces( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListNamespacesRequest; - let request = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_namespaces(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("namespaces", response.namespaces)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[pyo3(signature = (namespace_path, mode=None, properties=None))] @@ -725,29 +713,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_tables( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListTablesRequest; - let request = ListTablesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_tables(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("tables", response.tables)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Get the configuration for constructing an equivalent namespace client. diff --git a/python/src/listing.rs b/python/src/listing.rs index 2d1a3916e..79a7effdd 100644 --- a/python/src/listing.rs +++ b/python/src/listing.rs @@ -3,10 +3,11 @@ use crate::error::PythonErrorExt; use crate::runtime::future_into_py; -use futures::{StreamExt, future::poll_fn}; +use futures::{Future, StreamExt, future::poll_fn}; use lancedb::listing::{Listing, ListingOptions}; use pyo3::exceptions::PyStopAsyncIteration; -use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyDict, PyDictMethods}; +use pyo3::{Bound, Py, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; use std::sync::{Arc, Mutex}; pub fn options(page_token: Option, page_limit: Option) -> ListingOptions { @@ -52,6 +53,31 @@ impl NameListing { Ok(value) }) } + + fn next_page(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + let next_lock = self_.next_lock.clone(); + future_into_py(self_.py(), async move { + let _guard = next_lock.lock().await; + let page = futures::future::poll_fn(|cx| { + let mut inner = inner.lock().unwrap(); + let future = inner.next_page(); + futures::pin_mut!(future); + future.poll(cx) + }) + .await + .infer_error()?; + Python::attach(|py| -> PyResult>> { + page.map(|page| { + let dict = PyDict::new(py); + dict.set_item("items", page.items)?; + dict.set_item("page_token", page.page_token)?; + Ok(dict.unbind()) + }) + .transpose() + }) + }) + } } #[pyclass] diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 8a0e489c2..124f6a54a 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -13,7 +13,7 @@ use lance::dataset::refs::MAIN_BRANCH; use lance_namespace::models::{ CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest, DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest, - ListNamespacesResponse, ListTablesRequest, ListTablesResponse, + ListTablesRequest, }; #[cfg(feature = "aws")] use object_store::aws::AwsCredential; @@ -28,7 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; -use crate::listing::{Listing, ListingOptions}; +use crate::listing::{Listing, ListingOptions, ListingPage}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -1126,12 +1126,43 @@ impl Connection { self.internal.drop_all_tables(namespace_path).await } - /// List immediate child namespace names in the given namespace - pub async fn list_namespaces( + /// Lazily list immediate child namespace names. Each request asks for at + /// most 100 names unless `options.page_limit` is set. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_namespaces(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_namespaces( &self, - request: ListNamespacesRequest, - ) -> Result { - self.internal.list_namespaces(request).await + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_namespaces(ListNamespacesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.namespaces, + page_token: response.page_token, + }) + } + }) } /// Create a new namespace @@ -1182,23 +1213,43 @@ impl Connection { self.internal.namespace_client_config().await } - /// List tables with pagination support + /// Lazily list tables. Each request asks for at most 100 names unless + /// `options.page_limit` is set. /// - /// The default limit is 100 tables per page for both local and remote connections, - /// including namespaces. Zero returns an empty page without a continuation token. - /// Follow the response's opaque `page_token` until it is absent to retrieve every table; - /// a page can contain fewer than the limit even when more tables remain. - pub async fn list_tables(&self, mut request: ListTablesRequest) -> Result { - // Apply the SDK's page size here so every backend receives the same request. - let limit = request.limit.get_or_insert(100); - if *limit <= 0 { - return Ok(ListTablesResponse { - context: None, - tables: Vec::new(), - page_token: None, - }); - } - self.internal.list_tables(request).await + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_tables(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_tables( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_tables(ListTablesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.tables, + page_token: response.page_token, + }) + } + }) } /// Get the in-memory embedding registry. @@ -2260,26 +2311,11 @@ mod tests { } names.sort(); - // Walking in pages has to reach every table exactly once, with nothing lost at a - // page boundary. - let mut seen = Vec::with_capacity(names.len()); - let mut page_token = None; - loop { - let page = db - .list_tables(ListTablesRequest { - id: Some(Vec::new()), - limit: Some(2), - page_token, - ..Default::default() - }) - .await - .unwrap(); - seen.extend(page.tables); - page_token = page.page_token.filter(|token| !token.is_empty()); - if page_token.is_none() { - break; - } - } + let seen = futures::TryStreamExt::try_collect::>( + db.list_tables(&[], ListingOptions::default().page_limit(2)), + ) + .await + .unwrap(); assert_eq!(seen, names); } @@ -2305,37 +2341,20 @@ mod tests { .is_empty() ); - let first = db.list_tables(ListTablesRequest::default()).await.unwrap(); - assert_eq!(first.tables, names[..100]); + let mut listing = db.list_tables(&[], ListingOptions::default()); + assert_eq!(listing.num_page_results(), 0); + let first = listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items, names[..100]); assert!(first.page_token.is_some()); - let second = db - .list_tables(ListTablesRequest { - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(second.tables, names[100..]); + let second = listing.next_page().await.unwrap().unwrap(); + assert_eq!(second.items, names[100..]); assert!(second.page_token.is_none()); + assert!(listing.next_page().await.unwrap().is_none()); - let all = db - .list_tables(ListTablesRequest { - limit: Some(200), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(all.tables, names); - assert!(all.page_token.is_none()); - let empty = db - .list_tables(ListTablesRequest { - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(empty.tables.is_empty()); - assert!(empty.page_token.is_none()); + let mut all = db.list_tables(&[], ListingOptions::default().page_limit(200)); + assert_eq!(all.next_page().await.unwrap().unwrap().items, names); + let mut invalid = db.list_tables(&[], ListingOptions::default().page_limit(0)); + assert!(invalid.next_page().await.is_err()); } #[tokio::test] diff --git a/rust/lancedb/src/listing.rs b/rust/lancedb/src/listing.rs index 710202d9d..a0f6e5d5e 100644 --- a/rust/lancedb/src/listing.rs +++ b/rust/lancedb/src/listing.rs @@ -51,8 +51,9 @@ type FetchPage = /// A lazy stream of resources with inspectable pagination state. /// -/// No requests are made until the stream is polled. Empty pages with continuation -/// tokens are skipped. Errors terminate the stream and retain the failed request's +/// No requests are made until the stream is polled or [`Self::next_page`] is called. +/// Item iteration skips empty pages with continuation tokens; `next_page` returns +/// them. Errors terminate the listing and retain the failed request's /// token, allowing a new stream to retry it. Repeated tokens are rejected. /// /// A token advances when its page is fetched. Drain [`Self::num_page_results`] @@ -100,6 +101,64 @@ impl Listing { pub fn page_token(&self) -> Option<&str> { self.options.page_token.as_deref() } + + /// Return exactly one page without consuming a following page. + /// + /// If item iteration has started, this returns the remaining items in its + /// cached page. An empty page with a continuation token is returned as is. + pub async fn next_page(&mut self) -> Result>> { + if self.items.len() == 0 { + if self.done { + return Ok(None); + } + futures::future::poll_fn(|cx| self.poll_fetch_page(cx)).await?; + } + let items = std::mem::replace(&mut self.items, Vec::new().into_iter()).collect(); + Ok(Some(ListingPage { + items, + page_token: self.options.page_token.clone(), + })) + } + + fn poll_fetch_page(&mut self, cx: &mut Context<'_>) -> Poll> { + if self + .options + .page_limit + .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) + { + self.done = true; + return Poll::Ready(Err(Error::InvalidInput { + message: "Listing page limit must be between 1 and 2147483647".into(), + })); + } + let pending = self + .pending + .get_or_insert_with(|| (self.fetch)(self.options.clone())); + let response = futures::ready!(pending.as_mut().poll(cx)); + self.pending = None; + match response { + Ok(response) => { + let token = response.page_token.filter(|token| !token.is_empty()); + if token + .as_ref() + .is_some_and(|token| !self.seen_tokens.insert(token.clone())) + { + self.done = true; + return Poll::Ready(Err(Error::Runtime { + message: "Listing response repeated a page_token".into(), + })); + } + self.done = token.is_none(); + self.options.page_token = token; + self.items = response.items.into_iter(); + Poll::Ready(Ok(())) + } + Err(error) => { + self.done = true; + Poll::Ready(Err(error)) + } + } + } } impl Stream for Listing { @@ -114,41 +173,8 @@ impl Stream for Listing { if this.done { return Poll::Ready(None); } - if this - .options - .page_limit - .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) - { - this.done = true; - return Poll::Ready(Some(Err(Error::InvalidInput { - message: "Listing page limit must be between 1 and 2147483647".into(), - }))); - } - let pending = this - .pending - .get_or_insert_with(|| (this.fetch)(this.options.clone())); - let response = futures::ready!(pending.as_mut().poll(cx)); - this.pending = None; - match response { - Ok(response) => { - let token = response.page_token.filter(|token| !token.is_empty()); - if token - .as_ref() - .is_some_and(|token| !this.seen_tokens.insert(token.clone())) - { - this.done = true; - return Poll::Ready(Some(Err(Error::Runtime { - message: "Listing response repeated a page_token".into(), - }))); - } - this.done = token.is_none(); - this.options.page_token = token; - this.items = response.items.into_iter(); - } - Err(error) => { - this.done = true; - return Poll::Ready(Some(Err(error))); - } + if let Err(error) = futures::ready!(this.poll_fetch_page(cx)) { + return Poll::Ready(Some(Err(error))); } } } @@ -224,6 +250,36 @@ mod tests { assert!(listing.next().await.is_none()); } + #[tokio::test] + async fn next_page_preserves_empty_pages_and_remaining_items() { + let requests = Arc::new(Mutex::new(0)); + let count = requests.clone(); + let mut listing = Listing::new(ListingOptions::default(), move |options| { + *count.lock().unwrap() += 1; + async move { + let (items, page_token) = match options.page_token.as_deref() { + None => (vec![], Some("second".to_string())), + Some("second") => (vec!["a", "b"], Some("last".to_string())), + Some("last") => (vec!["c"], None), + _ => unreachable!(), + }; + Ok(ListingPage { items, page_token }) + } + }); + assert_eq!(*requests.lock().unwrap(), 0); + assert!(listing.next_page().await.unwrap().unwrap().items.is_empty()); + assert_eq!(*requests.lock().unwrap(), 1); + assert_eq!(listing.try_next().await.unwrap(), Some("a")); + assert_eq!(*requests.lock().unwrap(), 2); + let remaining = listing.next_page().await.unwrap().unwrap(); + assert_eq!(remaining.items, vec!["b"]); + assert_eq!(remaining.page_token.as_deref(), Some("last")); + assert_eq!(*requests.lock().unwrap(), 2); + assert_eq!(listing.next_page().await.unwrap().unwrap().items, vec!["c"]); + assert!(listing.next_page().await.unwrap().is_none()); + assert_eq!(*requests.lock().unwrap(), 3); + } + #[tokio::test] async fn invalid_limits_do_not_fetch() { for limit in [0, u32::MAX] { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 4a8cbf32e..8a456ab3b 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -2349,16 +2349,8 @@ mod tests { .unwrap() .is_empty() ); - let page = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(page.tables.is_empty()); - assert!(page.page_token.is_none()); + let mut listing = conn.list_tables(&namespace, ListingOptions::default().page_limit(0)); + assert!(listing.next_page().await.is_err()); } #[rstest::rstest] @@ -2393,25 +2385,17 @@ mod tests { ) .unwrap() }); - let first = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(first.tables.len(), 100); + let mut first_listing = conn.list_tables(&namespace, ListingOptions::default()); + let first = first_listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items.len(), 100); assert_eq!(first.page_token.as_deref(), Some("opaque-token")); - let second = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); + let mut second_listing = conn.list_tables( + &namespace, + ListingOptions::default().page_token(first.page_token.unwrap()), + ); + let second = second_listing.next_page().await.unwrap().unwrap(); let expected: Vec<_> = (100..130).map(|i| format!("t{i:03}")).collect(); - assert_eq!(second.tables, expected); + assert_eq!(second.items, expected); assert!(second.page_token.is_none()); } @@ -3680,14 +3664,9 @@ mod tests { assert!(table.is_ok(), "Failed to create table: {:?}", table.err()); // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .expect("Failed to list tables"); - assert_eq!(list_response.tables, vec!["test_table"]); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names, vec!["test_table"]); // Get namespace client and verify it can also list tables let namespace_client = conn.namespace_client().await.unwrap(); @@ -3754,17 +3733,12 @@ mod tests { } // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(list_response.tables.len(), 3); - assert!(list_response.tables.contains(&"table1".to_string())); - assert!(list_response.tables.contains(&"table2".to_string())); - assert!(list_response.tables.contains(&"table3".to_string())); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names.len(), 3); + assert!(names.contains(&"table1".to_string())); + assert!(names.contains(&"table2".to_string())); + assert!(names.contains(&"table3".to_string())); } } From 444f7f3bde8d18aa6ef39cba714dbe1bb448fb1d Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 23 Oct 2026 12:45:25 +0800 Subject: [PATCH 04/36] ci: refresh apt lists before installing protobuf in docs build (#4636) The docs deploy workflow installs protobuf-compiler without refreshing apt lists first. Ubuntu noble-updates replaced `protobuf 3.21.12-8.2ubuntu0.3` with `.0.4`, so stale package lists in older runner images make `apt install` fail with 404 and exit code 100. Run `apt update` before installing, matching the other workflows. --- .github/workflows/docs.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index d0ec583bf..fb3fb1d6a 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -38,6 +38,7 @@ jobs: uses: actions/checkout@v6 - name: Install dependencies needed for ubuntu run: | + sudo apt update sudo apt install -y protobuf-compiler libssl-dev rustup update && rustup default - name: Set up Python From 45c4b159cab990a749d592b42e21e030f7479192 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Thu, 22 Oct 2026 22:51:40 -0700 Subject: [PATCH 05/36] fix: cap the remote job poll interval at one second (#4629) job.wait() polled on a backoff that doubled up to 5 s, so a job's completion was seen 2.5 s late on average, as long as a small refresh itself takes. Cap the interval at 1 s; the first polls still back off from 200 ms, so a quick job is seen quickly and a long one costs one request a second. Linear: OSS-2704 https://linear.app/lancedb/issue/OSS-2704 --- rust/lancedb/src/remote/job.rs | 36 ++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 41d8971ac..7d2e70ac9 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -18,7 +18,9 @@ use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. const INITIAL_POLL_INTERVAL: Duration = Duration::from_millis(200); -const MAX_POLL_INTERVAL: Duration = Duration::from_secs(5); +/// A job's completion is seen at most this long after it happens. A 5 s cap +/// added 2.5 s on average to every short refresh, as much as the job itself. +const MAX_POLL_INTERVAL: Duration = Duration::from_secs(1); #[derive(Debug, Clone, PartialEq, Eq)] enum JobState { @@ -268,13 +270,18 @@ impl JobHandle for RemoteJob { #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + use async_trait::async_trait; use crate::Result; use crate::function::{FunctionVersion, RefreshColumnResult}; use crate::job::{Job, JobHandle, TerminalResult}; + use crate::remote::client::test_utils::client_with_handler; - use super::DescribeJobResponse; + use super::{DescribeJobResponse, RemoteJob}; const FUNCTION_JOB: &str = include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); @@ -333,6 +340,31 @@ mod tests { ); } + /// Five in-progress polls then done: 0.2 + 0.4 + 0.8 + 1 + 1 s of waiting. + /// Under the old 5 s cap the same job would have waited 10 s. + #[tokio::test(start_paused = true)] + async fn wait_polls_at_most_every_second() { + let describes = Arc::new(AtomicUsize::new(0)); + let seen = describes.clone(); + let client = client_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/jobs/describe"); + let n = seen.fetch_add(1, Ordering::SeqCst); + let body = if n < 5 { + r#"{"job_state":"IN_PROGRESS"}"#.to_string() + } else { + UNIT_JOB.to_string() + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let started = tokio::time::Instant::now(); + RemoteJob::new(client, "job-1".to_string()) + .wait() + .await + .expect("unit job completes"); + assert_eq!(describes.load(Ordering::SeqCst), 6); + assert_eq!(started.elapsed(), Duration::from_millis(3400)); + } + #[test] fn remote_wire_unknown_fields_are_forward_decodable() { let response: DescribeJobResponse = From a9535684cc1d9901c54c852f103dd0d492d61c3f Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 23 Oct 2026 20:17:38 +0800 Subject: [PATCH 06/36] chore(deps): bump ctutils to 0.4.4 for RUSTSEC-2026-0385 (#4638) Fixes RUSTSEC-2026-0385: `ctutils` < 0.4.4 generates masks with branches on secret values (`u128` mask generation / comparisons are not constant-time on some targets). - Advisory: https://rustsec.org/advisories/RUSTSEC-2026-0385 - Affected: `ctutils 0.4.2` pinned in `Cargo.lock`, reached transitively via `digest 0.11` (RustCrypto) - Change: `cargo update -p ctutils --precise 0.4.4`; single lockfile-only change (patch bump, no API impact) The `deny` CI job started failing for every PR head after the advisory was published at 06:45Z today; `main`'s last `deny` run (05:51Z) predates it and does not yet cover this advisory. --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 34ee05cb1..8bad6ca92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2230,9 +2230,9 @@ dependencies = [ [[package]] name = "ctutils" -version = "0.4.2" +version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +checksum = "12218662ccac37d9714f2191252553d67585b97df8c839cbe49e0bde278dc865" dependencies = [ "cmov", ] From b762ec221ba5d99ab368f2959f50381c589c762b Mon Sep 17 00:00:00 2001 From: LuQQiu Date: Fri, 23 Oct 2026 10:06:20 -0700 Subject: [PATCH 07/36] docs: add Enterprise distributed search reference page (#4625) Adds the Enterprise reference page at `docs/web/enterprise/distributed-search.mdx`, following the same arrangement as Enterprise architecture. The page describes distributed search and the embedded execution model, with links to shared search guides. Enterprise architecture links to it, and the official sidebar lists it immediately after Architecture in the existing Deployment group. Adds stable anchors for the RaBitQ and `approx_mode` sections linked by the Enterprise guide. Companion: https://github.com/lancedb/sophon/pull/8986 supplies the detailed Enterprise overlay at the same path. Merge this reference-page PR first so the overlay has an existing route to replace. Validation: - Combined-site assembly with both PR branches passed (309 files). - Mintlify build validation and link, anchor, and redirect checks passed. --- docs/web/docs.json | 1 + docs/web/enterprise/architecture.mdx | 3 ++ docs/web/enterprise/distributed-search.mdx | 33 ++++++++++++++++++++++ docs/web/indexing/vector-index.mdx | 2 +- docs/web/search/vector-search.mdx | 2 +- 5 files changed, 39 insertions(+), 2 deletions(-) create mode 100644 docs/web/enterprise/distributed-search.mdx diff --git a/docs/web/docs.json b/docs/web/docs.json index 9c9468a3b..54203f1a0 100644 --- a/docs/web/docs.json +++ b/docs/web/docs.json @@ -350,6 +350,7 @@ "pages": [ "enterprise/index", "enterprise/architecture", + "enterprise/distributed-search", "enterprise/deployment/index", "deployment/storage", "deployment/monitoring", diff --git a/docs/web/enterprise/architecture.mdx b/docs/web/enterprise/architecture.mdx index c47dcbd2e..8f887c142 100644 --- a/docs/web/enterprise/architecture.mdx +++ b/docs/web/enterprise/architecture.mdx @@ -18,3 +18,6 @@ run on machines that can be replaced without any of them owning the data. [Architecture](/architecture) covers this alongside the embedded form, which is the comparison most readers want first. + +See [Distributed search](/enterprise/distributed-search) for how Enterprise +distributes index execution and table-data reads. diff --git a/docs/web/enterprise/distributed-search.mdx b/docs/web/enterprise/distributed-search.mdx new file mode 100644 index 000000000..ebac960ba --- /dev/null +++ b/docs/web/enterprise/distributed-search.mdx @@ -0,0 +1,33 @@ +--- +title: "Distributed search" +sidebarTitle: "Distributed search" +description: "How LanceDB Enterprise distributes index execution and table-data reads." +icon: "network-wired" +availability: + oss: unavailable + enterprise: available + summary: >- + LanceDB Enterprise distributes index execution and table-data reads across + plan executors. Embedded LanceDB executes searches in the application process. +--- + +Distributed search spreads index execution, table-data reads, and caching across +multiple serving machines in LanceDB Enterprise. Query nodes plan queries and +combine results. Plan executors cache index structures and table data and execute +the remote work. See [Enterprise architecture](/enterprise/architecture). + +## Embedded search {#embedded-search} + +LanceDB OSS runs inside your application process. It searches indexes and reads +table data using that process's resources, including its index cache. Reading a +table from object storage does not by itself distribute query execution across +machines. + +## Search interfaces {#search-interfaces} + +The shared search guides describe the query APIs and index settings: + +- [Vector search](/search/vector-search) and [vector indexes](/indexing/vector-index). +- [Full-text search](/search/full-text-search) and [FTS indexes](/indexing/fts-index). +- [Filtering](/search/filtering) and [scalar indexes](/indexing/scalar-index). +- [Query plans](/search/explain) for inspecting execution and runtime measurements. diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index 2742a992f..b26e0c295 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -384,7 +384,7 @@ The inner workings of RaBitQ quantization are rather mathematically dense. It ge applying a uniformly random, approximately distance-preserving orthogonal transformation of the vertices of the $d$-dimensional hypercube, where $d$ is the dimensionality of the dataset. We defer the details, and an elegant theoretical error bound, to [the original paper.](https://arxiv.org/pdf/2405.12497) -##### Using RaBitQ +##### Using RaBitQ {#using-rabitq} Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, call `create_index` with `index_type="IVF_RQ"`. Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`. diff --git a/docs/web/search/vector-search.mdx b/docs/web/search/vector-search.mdx index 3519df553..cc6ae5421 100644 --- a/docs/web/search/vector-search.mdx +++ b/docs/web/search/vector-search.mdx @@ -235,7 +235,7 @@ The table below summarizes the behavior of `_distance` in search results based o See [vector indexes](/indexing/vector-index) for index configuration. The following sections cover query-time settings. -### Tuning `approx_mode` +### Tuning `approx_mode` {#tuning-approx-mode} Use `approx_mode` when you want to adjust the speed/recall trade-off for approximate vector search at query time. This setting currently applies only to RQ-quantized indexes, such as `IVF_RQ`; other index types ignore it. From 5a79c7d22da9f9729d27adc7e9d0b166767bbd4e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 24 Oct 2026 01:52:45 +0800 Subject: [PATCH 08/36] ci: run aarch64 Windows tests on main only (#4637) Run the aarch64 Windows build/test cycle only on pushes to main so PRs wait for a single Windows platform (x86_64). Main keeps the same native aarch64 coverage as before. GitHub evaluates job-level `if` before the matrix strategy is applied, so the aarch64 target is split into its own `windows-arm` job instead of being filtered inside the existing matrix. The job keeps the historical check name. --- .github/workflows/rust.yml | 43 ++++++++++++++++++++++++++++++++------ 1 file changed, 37 insertions(+), 6 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index cbfefc83d..930ae67e2 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -245,12 +245,6 @@ jobs: include: - target: x86_64-pc-windows-msvc runner: windows-2022 - # windows-11-arm is a standard runner, so it is free on public repos. - # Running natively lets the aarch64 tests actually execute -- this - # job used to cross-compile them and then skip the test step, paying - # full codegen and link cost for a compile check. - - target: aarch64-pc-windows-msvc - runner: windows-11-arm runs-on: ${{ matrix.runner }} defaults: run: @@ -279,6 +273,43 @@ jobs: # dependency graph a second time. cargo test --profile ci --features aws,remote --locked --target ${{ matrix.target }} + # Run the aarch64 Windows tests only on main so PRs wait for a single + # Windows platform. windows-11-arm is a standard runner, so it is free on + # public repos. Running natively lets the aarch64 tests actually execute -- + # this job used to cross-compile them and then skip the test step, paying + # full codegen and link cost for a compile check. + windows-arm: + # Keep the historical name so the check shows up as before on main. + name: windows (aarch64-pc-windows-msvc, windows-11-arm) + if: github.ref == 'refs/heads/main' + runs-on: windows-11-arm + defaults: + run: + working-directory: rust/lancedb + steps: + - uses: actions/checkout@v6 + - name: Set target + run: rustup target add aarch64-pc-windows-msvc + - uses: Swatinem/rust-cache@v2 + with: + # Restore everywhere, but only save from main. Per-PR saves are + # unreadable outside their own branch anyway, since GitHub scopes + # caches to the creating ref. + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Install Protoc v21.12 + run: choco install --no-progress protoc + - name: Build + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + cargo build --profile ci --features aws,remote --tests --locked --target aarch64-pc-windows-msvc + - name: Run tests + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + # `--target` has to match the build step above. Without it cargo uses + # target/ci/ rather than target//ci/ and rebuilds the entire + # dependency graph a second time. + cargo test --profile ci --features aws,remote --locked --target aarch64-pc-windows-msvc + msrv: # Check the minimum supported Rust version name: MSRV Check - Rust v${{ matrix.msrv }} From 702389654e11d1af5cd120d64866f028f41d6b90 Mon Sep 17 00:00:00 2001 From: KarthikS Date: Fri, 23 Oct 2026 15:27:32 -0400 Subject: [PATCH 09/36] fix(python): align pylance extra and preview install indexes (#4556) Fixes #4248. The preview index still contains pylance 1.0.0b6. uv's default first-index resolution can select it, or backtrack to an old LanceDB release when installing the extra. Meanwhile, the extra still permits pylance 5 even though current main bundles Lance 14. Raise the optional dependency to `pylance>=14.0.0b12`, the version already pinned in the test extra, and keep the lockfile metadata consistent. Document both Fury indexes and uv's `unsafe-best-match` strategy in the Python README and the source-owned quickstart (`docs/web/quickstart.mdx`). Explain the index trust requirement and recommend installing the extra to enforce the version floor. Add CI checks that reject older Lance majors and ensure the pylance and tests extras remain compatible. Run these checks when either dependency manifest changes, so a future Lance major bump cannot silently leave the Python minimum behind. ### Validation - Recreated the reported install on Ubuntu 26.04 / CPython 3.12.15 with uv 0.12.23: `lancedb==0.40.0b5` plus the old index selects `pylance==1.0.0b6`. Reproduced both the missing `blob_array` import and the unexpected `blob_handling` scanner argument using a real blob table. - Resolved the modified `python/pyproject.toml` with the pylance extra against both indexes: selects `pylance==14.0.0b18`, matching the current Rust pin. - Ran both documented pip and uv installations in isolated environments: `lancedb==0.41.0b7`, `pylance==14.0.0b18`; blob-array construction and a blob-bytes round trip through `to_pandas(blob_mode="bytes")` pass. Both package dependency checks pass. - New regression fails on the original dependency bound (27 rejected-version cases incorrectly accepted), then passes with the fix. Full CI-script suite: 3 tests pass. - `uv lock --check` with both indexes passes; focused Ruff checks pass; `git diff --check` passes. - Repository-wide Ruff reports 49 existing errors in untouched `docs/web-tests` files. Unrelated formatting changes were discarded. Runtime smoke checks used released native wheels; source dependency resolution was tested directly from the modified pyproject. No Rust or Python runtime code changes. --------- Co-authored-by: Xuanwo --- .../lancedb-update-lance-dependency/SKILL.md | 27 +++++++++++ .github/workflows/ci-scripts.yml | 6 +++ ci/tests/test_python_dependencies.py | 47 +++++++++++++++++++ docs/web/quickstart.mdx | 27 +++++++++++ python/README.md | 27 +++++++++++ python/pyproject.toml | 5 +- python/uv.lock | 18 +++---- 7 files changed, 146 insertions(+), 11 deletions(-) create mode 100644 ci/tests/test_python_dependencies.py diff --git a/.agents/skills/lancedb-update-lance-dependency/SKILL.md b/.agents/skills/lancedb-update-lance-dependency/SKILL.md index e97fae7ac..3cba2ae75 100644 --- a/.agents/skills/lancedb-update-lance-dependency/SKILL.md +++ b/.agents/skills/lancedb-update-lance-dependency/SKILL.md @@ -48,6 +48,23 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r This updates the Rust workspace Lance dependencies through `ci/set_lance_version.py`, updates `java/pom.xml`, refreshes Cargo metadata, and prints JSON metadata containing `branch_name`, `commit_message`, and `pr_title`. + Also synchronize the Python requirements in `python/pyproject.toml`: set the + `tests` extra's exact `pylance` pin to the matching published Python release, + and update the `pylance` extra's minimum when needed to exclude older bundled + Lance majors. The minimum must still accept the test pin. Python prerelease + spelling differs from Rust (for example, `15.0.0-beta.1` becomes `15.0.0b1`). + Confirm the matching Python release is published before changing the pin. + + Refresh and check `python/uv.lock` from the repository root: + + ```bash + uv lock --project python --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + uv lock --project python --check --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + ``` + + This strategy considers versions across indexes; use it only with trusted + indexes. Commit the Python requirements and lockfile together with the bump. + 5. Run validation: ```bash @@ -57,6 +74,16 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r Fix real diagnostics and rerun clippy until it succeeds. Do not skip warnings. + Run the Python dependency guards with Python 3.11+ and `packaging` installed: + + ```bash + python3 -m unittest discover -s ci/tests -p test_python_dependencies.py -v + ``` + + These checks reject older Lance majors and ensure the `pylance` and `tests` + extras remain compatible. The full CI-script suite uses + `python3 -m unittest discover -s ci/tests -v`. + 6. Inspect `git status --short` and `git diff` to ensure only the Lance dependency update and required compatibility fixes are present. 7. If the task only asks to prepare local changes, stop here and report the changed files and validation result. diff --git a/.github/workflows/ci-scripts.yml b/.github/workflows/ci-scripts.yml index f7f0dbb07..d45289c2d 100644 --- a/.github/workflows/ci-scripts.yml +++ b/.github/workflows/ci-scripts.yml @@ -5,6 +5,8 @@ on: branches: - main paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -12,6 +14,8 @@ on: - .github/workflows/ci-scripts.yml pull_request: paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -30,5 +34,7 @@ jobs: - uses: actions/setup-python@v6 with: python-version: "3.13" + - name: Install test dependencies + run: python -m pip install packaging - name: Run tests run: python -m unittest discover -s ci/tests -v diff --git a/ci/tests/test_python_dependencies.py b/ci/tests/test_python_dependencies.py new file mode 100644 index 000000000..97fb672c6 --- /dev/null +++ b/ci/tests/test_python_dependencies.py @@ -0,0 +1,47 @@ +import tomllib +import unittest +from pathlib import Path + +from packaging.requirements import Requirement +from packaging.version import Version + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +class PythonDependenciesTest(unittest.TestCase): + def setUp(self): + with (REPO_ROOT / "python/pyproject.toml").open("rb") as f: + self.project = tomllib.load(f)["project"] + self.pylance = Requirement(self.project["optional-dependencies"]["pylance"][0]) + + def test_pylance_extra_excludes_older_lance_majors(self): + with (REPO_ROOT / "Cargo.toml").open("rb") as f: + lance = tomllib.load(f)["workspace"]["dependencies"]["lance"] + if "version" not in lance: + self.skipTest("Local Lance checkout has no version requirement") + major = Version(lance["version"].lstrip("=")).major + for older_major in range(major): + for suffix in (".0.0", ".0.0b6", ".999.999"): + version = f"{older_major}{suffix}" + with self.subTest(version=version): + self.assertFalse( + self.pylance.specifier.contains(version, prereleases=True), + f"The pylance extra must exclude Lance {version}; " + f"the bundled Lance is major {major}", + ) + + def test_pylance_extra_accepts_tested_version(self): + requirements = [ + Requirement(dep) for dep in self.project["optional-dependencies"]["tests"] + ] + tested = next(dep for dep in requirements if dep.name == "pylance") + pin = next(spec for spec in tested.specifier if spec.operator == "==") + self.assertTrue( + self.pylance.specifier.contains(pin.version, prereleases=True), + f"The tests and pylance extras must be installable together: {tested}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/web/quickstart.mdx b/docs/web/quickstart.mdx index 788dd6585..3f466041d 100644 --- a/docs/web/quickstart.mdx +++ b/docs/web/quickstart.mdx @@ -58,6 +58,33 @@ uv add --prerelease allow --index https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra and +include the Lance preview index: + + +```bash Python (pip) icon="terminal" +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +```bash Python (uv) icon="terminal" +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + Pre-release builds receive the same level of testing as stable releases, but their availability is not guaranteed for more than 6 months after release. For real-world workloads, we recommend you use the latest stable release diff --git a/python/README.md b/python/README.md index f941c9e20..907f7c099 100644 --- a/python/README.md +++ b/python/README.md @@ -40,6 +40,33 @@ Stable releases are created about every 2 weeks. For the latest features and bug pip install --pre --extra-index-url https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra +and include the Lance preview index: + +```bash +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +With uv, also select the best matching version across the indexes: + +```bash +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + ### Threading in CPU-limited containers LanceDB uses separate pools for compute work and storage I/O. On a container with diff --git a/python/pyproject.toml b/python/pyproject.toml index ed54ae536..a0de6509c 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -45,7 +45,8 @@ repository = "https://github.com/lancedb/lancedb" [project.optional-dependencies] pylance = [ - "pylance>=5.0.0b5", + # Keep the minimum in the bundled Lance major and compatible with tests below. + "pylance>=15.0.0b1", ] # A library only needs the OpenTelemetry API; the application supplies and # configures the SDK (the actual exporter/reader). See @@ -63,7 +64,7 @@ tests = [ "polars>=0.19, <=1.32.3", "pyarrow<25", "pyarrow-stubs>=16.0", - "pylance==14.0.0b12", + "pylance==15.0.0b1", "requests>=2.31.0", "datafusion>=54,<55", "opentelemetry-sdk>=1.30.0", diff --git a/python/uv.lock b/python/uv.lock index e808f5f9e..99bcb439c 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -2006,8 +2006,8 @@ requires-dist = [ { name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" }, { name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" }, { name = "pydantic", specifier = ">=2.7.4,<3" }, - { name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" }, - { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" }, + { name = "pylance", marker = "extra == 'pylance'", specifier = ">=15.0.0b1" }, + { name = "pylance", marker = "extra == 'tests'", specifier = "==15.0.0b1" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" }, { name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" }, { name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" }, @@ -3938,7 +3938,7 @@ crypto = [ [[package]] name = "pylance" -version = "14.0.0b12" +version = "15.0.0b1" source = { registry = "https://pypi.fury.io/lance-format/" } dependencies = [ { name = "lance-namespace" }, @@ -3948,12 +3948,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" }, - { url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" }, - { url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" }, - { url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" }, - { url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" }, - { url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" }, + { url = "https://pypi.fury.io/lance-format/-/ver_23L1Nd/pylance-15.0.0b1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:19605d3daca1224cc0db079054be5d4eec403dcf286bf4617430bfc259ab0a52" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1Ho3RK/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be22b0b6f1e1d250ffaa82b449e982c817db33103c483d1005020285310623cf" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1QmyF0/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5b2761d60b1b9db112058213f1f7b2de82b08cbe65888fab08b0a56e702482e" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1in2Ph/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:294a0fedae312ed32d1b66cddbf7e69770ec9b4c82a873af60c619ed88450718" }, + { url = "https://pypi.fury.io/lance-format/-/ver_21xpLH/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ff20a16703d4d3f3bc675b25ce1871ba1d54137929dc46305bffdd07f7b35a6b" }, + { url = "https://pypi.fury.io/lance-format/-/ver_AAGzF/pylance-15.0.0b1-cp310-abi3-win_amd64.whl", hash = "sha256:3b134e24a5becfc940ca134cf78aa3bae8b80ec40f8927a1b856273ab76213f2" }, ] [[package]] From e74a7c15b5f322f8ac15b96bdfc391616518d711 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 24 Oct 2026 03:47:16 +0800 Subject: [PATCH 10/36] fix: bump tokio to address RUSTSEC-2026-0386 (#4642) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upgrade the locked Tokio version from 1.53.1 to 1.53.3 to fix [RUSTSEC-2026-0386](https://rustsec.org/advisories/RUSTSEC-2026-0386) and restore the Rust workflow's `deny` job. The diff changes only Tokio's version and checksum. Validation: - `cargo deny check advisories` — passed. - `cargo deny check advisories bans licenses sources` — passed with the repository's `deny.toml`. - `cargo check --quiet -p lancedb` — passed (dead-code warnings). - PR `deny` job — [passed](https://github.com/lancedb/lancedb/actions/runs/40000920646/job/120576455117). --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8bad6ca92..fac89c4da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10446,9 +10446,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.53.1" +version = "1.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "ce3335fa71841cda333a58d7615b03901380ecf09d59b3296d21f8bbac0dde4e" dependencies = [ "bytes", "libc", From 3321860e9fc7248d9c97c7513bc0ac0d258738a0 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Sat, 24 Oct 2026 04:22:30 +0800 Subject: [PATCH 11/36] fix(python): preserve Arrow map scalars in updates (#4448) Arrow map scalars returned by a table query cannot currently be passed back to `update`: the generic Arrow scalar conversion turns them into a list of tuples, which the SQL literal encoder cannot handle. Empty map scalars become an untyped SQL list and also fail to update a map column. Add a dedicated `MapScalar` conversion that renders its keys and values as a SQL map while preserving the existing null conversion. The regression tests exercise real synchronous and asynchronous table updates, quoted keys, nullable values, empty typed maps, and nested maps. Validation: the complete documented unit-test scope passed locally (1,515 passed, 37 skipped; unchanged baseline: 1,507 passed, 37 skipped). Full doctests passed (71 passed, 6 skipped), and repository-wide Python formatting/lint, Rust formatting/clippy, and Python 3.13 type checking passed. The locally built native wheel also passed all eight regression cases in a separate Python 3.13 environment. Full Linux/S3 integration and all platform matrices remain for CI. The SQL map constructor's existing backend constraints still apply: duplicate keys are rejected, and a nested map containing both nonempty and empty inner maps can fail SQL type unification. This change does not claim arbitrary Arrow map roundtrip support. AI assistance: implementation and local verification used Codex with gpt-6.1-sol agents; the diff and evidence were reviewed before submission. --------- Co-authored-by: Xuanwo --- python/python/lancedb/util.py | 9 ++++++ python/python/tests/test_table.py | 51 +++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index a043f0c8d..36536f499 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -452,6 +452,15 @@ def _(value: pa.Scalar): return value_to_sql(value.as_py()) +@value_to_sql.register(pa.MapScalar) +def _(value: pa.MapScalar): + if not value.is_valid: + return value_to_sql(None) + keys = value_to_sql(list(value.values.field(0))) + values = value_to_sql(list(value.values.field(1))) + return f"map({keys}, {values})" + + @value_to_sql.register(np.ndarray) def _(value: np.ndarray): return value_to_sql(value.tolist()) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index b55abfa7e..18d3c16db 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -2724,6 +2724,57 @@ def test_update_with_arrow_scalar(mem_db: DBConnection): assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]] +@pytest.mark.asyncio +@pytest.mark.parametrize("async_api", [False, True]) +@pytest.mark.parametrize("entries", [[("o'brien", 4), ("nil", None)], [], None]) +async def test_update_map_arrow_scalar(tmp_path, async_api, entries): + map_type = pa.map_(pa.string(), pa.int64()) + data = pa.table( + { + "id": [1, 2], + "attrs": pa.array([[("original", 1)], entries], type=map_type), + } + ) + if async_api: + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("maps", data) + source = await table.query().where("id = 2").to_arrow() + result = await table.update({"attrs": source["attrs"][0]}, where="id = 1") + actual = await table.query().to_arrow() + else: + db = lancedb.connect(tmp_path) + table = db.create_table("maps", data) + source = table.search().where("id = 2").to_arrow() + result = table.update(values={"attrs": source["attrs"][0]}, where="id = 1") + actual = table.to_arrow() + assert result.rows_updated == 1 + assert actual.schema.field("attrs").type == map_type + assert actual.sort_by("id")["attrs"].to_pylist() == [entries, entries] + + +@pytest.mark.parametrize( + "map_type, entries", + [ + (pa.map_(pa.int64(), pa.string()), []), + ( + pa.map_(pa.string(), pa.map_(pa.string(), pa.int64())), + [("outer", [("inner", 3)]), ("null", None)], + ), + ], +) +def test_update_typed_map_arrow_scalar(mem_db: DBConnection, map_type, entries): + table = mem_db.create_table( + "maps", + pa.table({"id": [1, 2], "attrs": pa.array([None, entries], type=map_type)}), + ) + scalar = table.search().where("id = 2").to_arrow()["attrs"][0] + result = table.update(values={"attrs": scalar}, where="id = 1") + actual = table.to_arrow().sort_by("id")["attrs"] + assert result.rows_updated == 1 + assert actual.type == map_type + assert actual.to_pylist() == [entries, entries] + + @pytest.mark.parametrize("use_arrow_scalar", [False, True]) @pytest.mark.parametrize( "decimal_type, value", From 9da4658d57cbdcd037827f12effd9a4dc869a095 Mon Sep 17 00:00:00 2001 From: Joaquin Hui Gomez <132194176+joaquinhuigomez@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:10:07 +0100 Subject: [PATCH 12/36] fix(python): accept with_row_id in take_offsets/take_row_ids on every Table implementation (#4374) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Table.take_offsets` and `Table.take_row_ids` declare a keyword-only `with_row_id: bool = False`, but none of the three implementations accept it — `LanceTable`, `AsyncTable` and `RemoteTable` all raise `TypeError: ... got an unexpected keyword argument 'with_row_id'` on the documented call. The builders they return already support it (`LanceTakeQueryBuilder.with_row_id()` / `AsyncTakeQuery.with_row_id()`), so this is a signature that never reached the implementations. The base docstrings also don't mention the parameter and still say `take_offsets` returns a `pa.RecordBatch` (it returns a `LanceTakeQueryBuilder`), so those are corrected too. Each implementation now accepts the kwarg and applies `.with_row_id()` to the builder it returns. Tests: sync and async assert `_rowid` appears in the result columns; the remote test uses the suite's request-capturing fixture and asserts `with_row_id: true` reaches the wire (and is absent by default). All three fail on `main` with the `TypeError`. ruff format/check clean. A related mismatch I left alone: the deprecated `LanceTable.create_scalar_index`/`create_fts_index` also drop a `wait_timeout` the base declares — deprecated since 0.25, so not worth touching here. Co-authored-by: Xuanwo --- python/python/lancedb/remote/table.py | 18 +++++++-- python/python/lancedb/table.py | 58 ++++++++++++++++++++------- python/python/tests/test_query.py | 27 +++++++++++++ python/python/tests/test_remote_db.py | 17 ++++++++ 4 files changed, 102 insertions(+), 18 deletions(-) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 86e8fb8a0..0744410b3 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1154,11 +1154,21 @@ class RemoteTable(Table): """ return LOOP.run(self._table.uri()) - def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + if with_row_id: + builder.with_row_id() + return builder - def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + if with_row_id: + builder.with_row_id() + return builder def uses_v2_manifest_paths(self) -> bool: return LOOP.run(self._table.uses_v2_manifest_paths()) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 5eea55760..b9557f219 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2103,11 +2103,14 @@ class Table(ABC): ---------- offsets: list[int] The offsets to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - pa.RecordBatch - A record batch containing the rows at the given offsets. + LanceTakeQueryBuilder + A query builder that can be executed to get the rows at the given + offsets. """ def __getitems__(self, offsets: list[int]) -> pa.RecordBatch: @@ -2188,11 +2191,13 @@ class Table(ABC): ---------- row_ids: list[int] The row ids to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - AsyncTakeQuery - A query object that can be executed to get the rows. + LanceTakeQueryBuilder + A query builder that can be executed to get the rows. """ @abstractmethod @@ -3182,11 +3187,21 @@ class LanceTable(Table): """Get the current version of the table""" return LOOP.run(self._table.version()) - def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + if with_row_id: + builder.with_row_id() + return builder - def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + if with_row_id: + builder.with_row_id() + return builder def blob_columns(self) -> list[str]: return LOOP.run(self._table.blob_columns()) @@ -7325,7 +7340,9 @@ class AsyncTable: """ await self._inner.restore(version) - def take_offsets(self, offsets: list[int]) -> AsyncTakeQuery: + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> AsyncTakeQuery: """ Take a list of offsets from the table. @@ -7343,15 +7360,23 @@ class AsyncTable: ---------- offsets: list[int] The offsets to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - pa.RecordBatch - A record batch containing the rows at the given offsets. + AsyncTakeQuery + A query object that can be executed to get the rows at the given + offsets. """ - return AsyncTakeQuery(self._inner.take_offsets(offsets), self) + query = AsyncTakeQuery(self._inner.take_offsets(offsets), self) + if with_row_id: + query.with_row_id() + return query - def take_row_ids(self, row_ids: list[int]) -> AsyncTakeQuery: + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> AsyncTakeQuery: """ Take a list of row ids from the table. @@ -7372,13 +7397,18 @@ class AsyncTable: ---------- row_ids: list[int] The row ids to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- AsyncTakeQuery A query object that can be executed to get the rows. """ - return AsyncTakeQuery(self._inner.take_row_ids(row_ids), self) + query = AsyncTakeQuery(self._inner.take_row_ids(row_ids), self) + if with_row_id: + query.with_row_id() + return query async def blob_columns(self) -> list[str]: return await self._inner.blob_columns() diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 58ea7deb5..cef985aa6 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -2324,6 +2324,33 @@ def test_take_queries(tmp_path): ] +def test_take_queries_with_row_id(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("test", pa.table({"idx": range(100)})) + + by_row_id = table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_row_id.column_names + by_offset = table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_offset.column_names + + # The column is opt-in and must stay out of the default projection. + assert "_rowid" not in table.take_row_ids([5, 2, 17]).to_arrow().column_names + + +@pytest.mark.asyncio +async def test_take_queries_with_row_id_async(tmp_path): + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("test", pa.table({"idx": range(100)})) + + by_row_id = await table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_row_id.column_names + by_offset = await table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_offset.column_names + + default = await table.take_row_ids([5, 2, 17]).to_arrow() + assert "_rowid" not in default.column_names + + def test_take_queries_to_batches(tmp_path): # Regression test for the sync take-query path: `to_batches` previously # raised ``AttributeError: 'AsyncTakeQuery' object has no attribute diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index a45e1aaef..9655f5807 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2225,6 +2225,23 @@ def test_head(): assert data == pa.table({"id": [1, 2, 3]}) +def test_remote_take_row_ids_with_row_id(): + bodies = [] + + def handler(body): + bodies.append(body) + return pa.table({"id": [1, 2], "_rowid": pa.array([0, 1], pa.uint64())}) + + with query_test_table(handler) as table: + result = table.take_row_ids([0, 1], with_row_id=True).to_arrow() + assert "_rowid" in result.column_names + assert bodies[-1]["with_row_id"] is True + + with query_test_table(handler) as table: + table.take_row_ids([0, 1]).to_arrow() + assert "with_row_id" not in bodies[-1] + + def test_query_sync_minimal(): def handler(body): assert body == { From 27b65064d0166feabe1c307eba07f66ca2601857 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:37:07 -0700 Subject: [PATCH 13/36] fix(python): support custom blob thresholds without pylance (#4634) --- python/python/lancedb/schema.py | 62 ++++++++++------ python/python/tests/test_blob.py | 37 ++++++++++ python/python/tests/test_table.py | 118 ++++++++++++++++++++++++++++++ 3 files changed, 193 insertions(+), 24 deletions(-) diff --git a/python/python/lancedb/schema.py b/python/python/lancedb/schema.py index 5edda5610..8bec17058 100644 --- a/python/python/lancedb/schema.py +++ b/python/python/lancedb/schema.py @@ -4,6 +4,7 @@ """Schema helpers for Lance blob columns.""" +import ctypes import importlib from typing import TYPE_CHECKING, Optional @@ -16,6 +17,7 @@ if TYPE_CHECKING: _BLOB_EXTENSION_NAME = "lance.blob.v2" _BLOB_V1_KEY = "lance-encoding:blob" _ARROW_EXT_NAME_KEY = "ARROW:extension:name" +_MAX_RUST_USIZE = ctypes.c_size_t(-1).value _BLOB_V2_STORAGE_TYPE = pa.struct( [ pa.field("data", pa.large_binary(), nullable=True), @@ -209,32 +211,44 @@ def blob( Maximum size in bytes of a single packed blob sidecar file before a new one is started. Must be positive. - The thresholds are forwarded to ``lance.blob_field``, which validates - them and stores them as field metadata, so pylance is required whenever - any threshold is set. When all thresholds are omitted the behavior is - unchanged. + Thresholds are validated and stored as Arrow field metadata, with the + same format as ``lance.blob_field``. They do not require pylance. """ - if ( - inline_size_threshold is None - and dedicated_size_threshold is None - and pack_file_size_threshold is None + metadata = {} + for threshold, value, allow_zero, key in ( + ( + "inline_size_threshold", + inline_size_threshold, + True, + "lance-encoding:blob-inline-size-threshold", + ), + ( + "dedicated_size_threshold", + dedicated_size_threshold, + False, + "lance-encoding:blob-dedicated-size-threshold", + ), + ( + "pack_file_size_threshold", + pack_file_size_threshold, + False, + "lance-encoding:blob-pack-file-size-threshold", + ), ): - blob_type = _resolve_blob_type() - return pa.field(name, blob_type(), nullable=nullable) - try: - lance = importlib.import_module("lance") - except ModuleNotFoundError as err: - raise ValueError( - "Setting blob tier thresholds requires pylance to be installed" - ) from err - _resolve_blob_type() - return lance.blob_field( - name, - nullable=nullable, - inline_size_threshold=inline_size_threshold, - dedicated_size_threshold=dedicated_size_threshold, - pack_file_size_threshold=pack_file_size_threshold, - ) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{threshold} must be an int, got {type(value).__name__}") + if value < 0 or (value == 0 and not allow_zero): + bound = "non-negative" if allow_zero else "positive" + raise ValueError(f"{threshold} must be {bound}") + if value > _MAX_RUST_USIZE: + raise OverflowError(f"{threshold} must fit in a Rust usize") + metadata[key] = str(value) + + blob_type = _resolve_blob_type() + field = pa.field(name, blob_type(), nullable=nullable) + return field.with_metadata(metadata) if metadata else field def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType: diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index ec1b8b27e..c2dc2a265 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import ctypes import io import subprocess import sys @@ -14,6 +15,7 @@ import pytest from lance.blob import Blob, BlobType as LanceBlobType import lancedb +from lancedb import schema as schema_module from lancedb._blob import ( _set_blob_column, blob_v2_projection_sources, @@ -139,6 +141,41 @@ def test_blob_without_thresholds_has_no_threshold_metadata(): assert not any(b"threshold" in key for key in metadata) +@pytest.mark.parametrize("nullable", [False, True]) +@pytest.mark.parametrize( + "thresholds", + [ + {}, + {"inline_size_threshold": 0}, + {"inline_size_threshold": 8192}, + {"dedicated_size_threshold": 1 << 20}, + {"pack_file_size_threshold": 1 << 26}, + { + "inline_size_threshold": 8192, + "dedicated_size_threshold": 1 << 20, + "pack_file_size_threshold": 1 << 26, + }, + { + "inline_size_threshold": ctypes.c_size_t(-1).value, + "dedicated_size_threshold": ctypes.c_size_t(-1).value, + "pack_file_size_threshold": ctypes.c_size_t(-1).value, + }, + ], +) +def test_blob_thresholds_match_pylance(monkeypatch, nullable, thresholds): + expected = lance.blob_field("payload", nullable=nullable, **thresholds) + actual = lancedb.blob("payload", nullable=nullable, **thresholds) + assert type(actual.type) is LanceBlobType + assert actual.equals(expected, check_metadata=True) + + # Both extension classes must encode the same Arrow field on the wire. + monkeypatch.setattr( + schema_module, "_resolve_blob_type", lambda: schema_module._FallbackBlobType + ) + fallback = lancedb.blob("payload", nullable=nullable, **thresholds) + assert pa.schema([fallback]).serialize() == pa.schema([expected]).serialize() + + def test_blob_thresholds_are_validated(): with pytest.raises(ValueError, match="inline_size_threshold must be non-negative"): lancedb.blob("payload", inline_size_threshold=-1) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 18d3c16db..6ccba606f 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -59,6 +59,124 @@ def _blob_v2_table(db: DBConnection, name: str): return table +@pytest.mark.parametrize( + "threshold", + [ + "inline_size_threshold", + "dedicated_size_threshold", + "pack_file_size_threshold", + ], +) +@pytest.mark.parametrize( + "value,error", + [ + (-1, ValueError), + (True, TypeError), + (False, TypeError), + (1.5, TypeError), + ("8192", TypeError), + (ctypes.c_size_t(-1).value + 1, OverflowError), + ], +) +def test_blob_thresholds_reject_invalid_values(threshold, value, error): + # Also runs in CI with pylance actually uninstalled. + with pytest.raises(error, match=f"{threshold} must"): + lancedb.blob("blob", **{threshold: value}) + + +@pytest.mark.parametrize( + "threshold", ["dedicated_size_threshold", "pack_file_size_threshold"] +) +def test_blob_thresholds_reject_zero_for_positive_thresholds(threshold): + with pytest.raises(ValueError, match=f"{threshold} must be positive"): + lancedb.blob("blob", **{threshold: 0}) + + +def test_blob_thresholds_accept_integer_boundaries(): + maximum = ctypes.c_size_t(-1).value + field = lancedb.blob( + "blob", + nullable=False, + inline_size_threshold=0, + dedicated_size_threshold=maximum, + pack_file_size_threshold=maximum, + ) + assert not field.nullable + assert field.metadata == { + b"lance-encoding:blob-inline-size-threshold": b"0", + b"lance-encoding:blob-dedicated-size-threshold": str(maximum).encode(), + b"lance-encoding:blob-pack-file-size-threshold": str(maximum).encode(), + } + + +def test_blob_thresholds_control_writes_after_reopen(tmp_path): + # Keep this in test_table.py: CI runs it without pylance or pandas installed. + field = lancedb.blob( + "blob", + inline_size_threshold=4, + dedicated_size_threshold=12, + pack_file_size_threshold=8, + ) + metadata = { + b"lance-encoding:blob-inline-size-threshold": b"4", + b"lance-encoding:blob-dedicated-size-threshold": b"12", + b"lance-encoding:blob-pack-file-size-threshold": b"8", + } + assert field.metadata == metadata + schema = pa.schema([pa.field("id", pa.int64()), field]) + db = lancedb.connect(tmp_path) + table = db.create_table("blobs", schema=schema) + # Every payload fits in the default inline tier; custom thresholds must + # instead send the larger values to packed and dedicated storage. + rows = [ + {"id": 0, "blob": b"a" * 2}, + {"id": 1, "blob": b"b" * 8}, + {"id": 2, "blob": b"c" * 16}, + {"id": 3, "blob": b"d" * 8}, + ] + + def check_writes(table, expected): + assert table.schema.field("blob").metadata.items() >= metadata.items() + hits = table.search().with_row_id(True).limit(10).to_arrow() + actual = {row["id"]: row for row in hits.to_pylist()} + assert actual.keys() == {row["id"] for row in expected} + for row in expected: + size = len(row["blob"]) + descriptor = actual[row["id"]]["blob"] + assert descriptor["size"] == size + assert descriptor["kind"] == (2 if size > 12 else 1 if size > 4 else 0) + row_ids = [actual[row["id"]]["_rowid"] for row in expected] + assert table.fetch_blobs("blob", row_ids).to_pylist() == [ + row["blob"] for row in expected + ] + # Two packed values reach the configured sidecar limit independently. + packed = [ + actual[row["id"]]["blob"] for row in expected if len(row["blob"]) == 8 + ] + assert len({desc["blob_id"] for desc in packed}) == len(packed) + + table.add(rows) + check_writes(table, rows) + table = lancedb.connect(tmp_path).open_table("blobs") + check_writes(table, rows) + updates = [ + {"id": 0, "blob": b"e" * 16}, + {"id": 1, "blob": b"f" * 2}, + {"id": 2, "blob": b"g" * 8}, + {"id": 3, "blob": b"h" * 8}, + {"id": 4, "blob": b"i" * 16}, + ] + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(updates) + ) + assert result.num_updated_rows == 4 + assert result.num_inserted_rows == 1 + check_writes(lancedb.connect(tmp_path).open_table("blobs"), updates) + + async def _blob_v2_table_async(db: AsyncConnection, name: str): schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")]) table = await db.create_table(name, schema=schema) From 8cf115b081d6a4b48db13a11c8fcbde8deadb750 Mon Sep 17 00:00:00 2001 From: alexcastgonz <82875882+pcbeingused333@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:25:44 -0400 Subject: [PATCH 14/36] fix(python): replace the retired gpt-4-turbo-preview default in OpenaiReranker (#4613) OpenAI shuts down `gpt-4-turbo` and its snapshots on 2026-10-23 (https://developers.openai.com/api/docs/deprecations), so `OpenaiReranker` fails for every user who relies on the default `gpt-4-turbo-preview` after that date. This changes the default to `gpt-5.6-sol`, the replacement OpenAI names for `gpt-4-turbo` on that page, in the class, its docstring and the integration docs page. `_rerank` also stops sending `temperature=0`: the GPT-5 models only accept the default temperature, so keeping it would fail the new default on its first call. The reranker still asks for `json_object` output and the rest of the request is unchanged. The new test calls `_rerank` with a mocked client and checks that no `temperature` is sent. It fails on main and passes with the change. No request was made to OpenAI. This change was made with Claude Code. Co-authored-by: Xuanwo --- docs/web/integrations/reranking/openai.mdx | 2 +- python/python/lancedb/rerankers/openai.py | 5 ++--- python/python/tests/test_rerankers.py | 19 +++++++++++++++++++ 3 files changed, 22 insertions(+), 4 deletions(-) diff --git a/docs/web/integrations/reranking/openai.mdx b/docs/web/integrations/reranking/openai.mdx index 6ceb8a030..71fe63117 100644 --- a/docs/web/integrations/reranking/openai.mdx +++ b/docs/web/integrations/reranking/openai.mdx @@ -26,7 +26,7 @@ credentials from environment variables or a secrets manager. ## Accepted Arguments {#accepted-arguments} | Argument | Type | Default | Description | | --- | --- | --- | --- | -| `model_name` | `str` | `"gpt-4-turbo-preview"` | The name of the reranker model to use.| +| `model_name` | `str` | `"gpt-5.6-sol"` | The name of the reranker model to use.| | `column` | `str` | `"text"` | The name of the column to use as input to the cross encoder model. | | `return_score` | `str` | `"relevance"` | Options are "relevance" or "all". The type of score to return. If "relevance", will return only the `_relevance_score. If "all" is supported, will return relevance score along with the vector and/or fts scores depending on query type. | | `api_key` | `str` | `None` | The API key to use. If None, will use the OPENAI_API_KEY environment variable. diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index da0394d1c..49b067f79 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -21,7 +21,7 @@ class OpenaiReranker(Reranker): Parameters ---------- - model_name : str, default "gpt-4-turbo-preview" + model_name : str, default "gpt-5.6-sol" The name of the cross encoder model to use. column : str, default "text" The name of the column to use as input to the cross encoder model. @@ -33,7 +33,7 @@ class OpenaiReranker(Reranker): def __init__( self, - model_name: str = "gpt-4-turbo-preview", + model_name: str = "gpt-5.6-sol", column: str = "text", return_score="relevance", api_key: Optional[str] = None, @@ -54,7 +54,6 @@ class OpenaiReranker(Reranker): response = self._client.chat.completions.create( model=self.model_name, response_format={"type": "json_object"}, - temperature=0, messages=[ { "role": "system", diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 743669367..dd6cb4ab5 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import os +from unittest.mock import MagicMock import random import lancedb @@ -547,6 +548,24 @@ def test_openai_reranker(tmp_path): _run_test_reranker(reranker, table, "single player experience", None, schema) +def test_openai_reranker_does_not_send_temperature(): + reranker = OpenaiReranker(api_key="test-key") + create = MagicMock() + create.return_value.choices = [ + MagicMock( + message=MagicMock( + content='{"documents": [{"content": "a", "relevance_score": 0.5}]}' + ) + ) + ] + reranker.__dict__["_client"] = MagicMock() + reranker.__dict__["_client"].chat.completions.create = create + + reranker._rerank(pa.table({"text": ["a"]}), "query") + + assert "temperature" not in create.call_args.kwargs + + @pytest.mark.skipif( os.environ.get("JINA_API_KEY") is None, reason="JINA_API_KEY not set" ) From 0dd78656b966233569d2f44a2d8833dffa9003e1 Mon Sep 17 00:00:00 2001 From: Kerman Lau <107780822+WitBlue6@users.noreply.github.com> Date: Sat, 24 Oct 2026 07:13:09 +0800 Subject: [PATCH 15/36] fix(python): honor explicit OpenAI reranker API keys (#4515) Fixes #4514 ## Summary Honor an explicitly supplied `OpenaiReranker(api_key=...)` before falling back to `OPENAI_API_KEY`, matching the constructor documentation. Previously the environment key always won when both were set, so an instance configured for a different account still used the process-wide account. ## Tests Added four cases using real OpenAI client construction without sending requests: explicit/environment conflict, explicit key alone, environment fallback, and missing credentials. The conflict case fails before the change; all four pass afterward. Clients are closed after each assertion. - Targeted source-module tests: 4 passed. - Ruff check and format check: passed for both changed files. - Full native source build was attempted but failed because the environment has no Rust toolchain; bootstrap then attempted a cache path outside the writable workspace. The targeted tests load the current Python reranker module against installed dependencies and the released LanceDB package; no database/native behavior is exercised. Full native/integration suite not run. This change and its tests were prepared with an AI coding agent. The diff was reviewed and the failure/pass behavior verified locally. Co-authored-by: Xuanwo --- python/python/lancedb/rerankers/openai.py | 2 +- .../tests/test_openai_reranker_credentials.py | 38 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) create mode 100644 python/python/tests/test_openai_reranker_credentials.py diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index 49b067f79..75efbbfc7 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -131,4 +131,4 @@ class OpenaiReranker(Reranker): "OPENAI_API_KEY not set. Either set it in your environment or \ pass it as `api_key` argument to the CohereReranker." ) - return openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY") or self.api_key) + return openai.OpenAI(api_key=self.api_key or os.environ.get("OPENAI_API_KEY")) diff --git a/python/python/tests/test_openai_reranker_credentials.py b/python/python/tests/test_openai_reranker_credentials.py new file mode 100644 index 000000000..aa11239b2 --- /dev/null +++ b/python/python/tests/test_openai_reranker_credentials.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pytest + +from lancedb.rerankers import OpenaiReranker + + +@pytest.mark.parametrize( + "explicit_key, environment_key, expected_key", + [ + ("explicit-test-key", "environment-test-key", "explicit-test-key"), + ("explicit-test-key", None, "explicit-test-key"), + (None, "environment-test-key", "environment-test-key"), + ], +) +def test_openai_reranker_api_key_precedence( + monkeypatch, explicit_key, environment_key, expected_key +): + pytest.importorskip("openai") + if environment_key is None: + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + else: + monkeypatch.setenv("OPENAI_API_KEY", environment_key) + + reranker = OpenaiReranker(api_key=explicit_key) + client = reranker._client + try: + assert client.api_key == expected_key + finally: + client.close() + + +def test_openai_reranker_requires_api_key(monkeypatch): + pytest.importorskip("openai") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(ValueError, match="OPENAI_API_KEY not set"): + OpenaiReranker()._client From dd4795028b390302f8edad3b1cdef9b2ed3d9b9f Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 08:25:16 +0800 Subject: [PATCH 16/36] fix(python): align sync remote methods with async behavior (#4411) Sync remote table and connection methods inherited bare base-class stubs, so unsupported conversion and bulk-drop operations raised `NotImplementedError` with an empty message and differed from the async API. Delegate `to_lance()`, `drop_all_tables()`, and `drop_database()` to the async implementation. Sync and async remote Lance conversion share #4306's native-table guard and reject before requesting a URI or opening client storage. Conversion documentation directs remote reads through server queries. Add an explicit Polars error, suggest query conversions in Arrow/Pandas/Polars errors, and give the remaining base-class errors an operation name. Preserve upstream `blob_mode` validation for remote Arrow and Pandas conversion. The shared remote bulk-drop error suggests dropping tables individually. Regression tests cover sync/async rejection parity, accessible and missing advertised locations, pinned versions, and bulk-drop rejection without requests. Sync blob query tests verify that descriptions come from the server and unsupported modes cannot open client storage. Keep the remote closed-blob-handle assertion consistent with the upstream `ValueError` contract. Validation: - All 30 focused regression cases passed. - 222 Python tests passed with rebuilt bindings in the locked uv environment: the full remote and error suites, plus selected local conversion, blob, and deletion tests. - Python lint and formatting, Rust formatting, `cargo check --quiet --features remote --tests --examples`, and `cargo clippy --quiet --features remote --tests --examples` passed. - Root Ruff checking reports 49 existing errors in upstream `docs/web-tests`; those files are unchanged relative to `main`. Fixes #4392 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/db.py | 31 +++-- python/python/lancedb/remote/db.py | 14 +++ python/python/lancedb/remote/table.py | 29 ++++- python/python/lancedb/table.py | 56 +++++---- python/python/tests/test_errors.py | 21 ++++ python/python/tests/test_remote_db.py | 157 ++++++++++++++++++++++++-- rust/lancedb/src/remote/db.rs | 4 +- 7 files changed, 271 insertions(+), 41 deletions(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 90dc1cccf..74acc28ba 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -548,7 +548,9 @@ class DBConnection(EnforceOverrides): LanceTable(name='table4', ...) """ - raise NotImplementedError + raise NotImplementedError( + "create_table is not supported for this connection type" + ) def __getitem__(self, name: str) -> LanceTable: return self.open_table(name) @@ -603,7 +605,9 @@ class DBConnection(EnforceOverrides): ------- A LanceTable object representing the table. """ - raise NotImplementedError + raise NotImplementedError( + "open_table is not supported for this connection type" + ) def create_materialized_view( self, @@ -743,13 +747,17 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_table is not supported for this connection type" + ) def drop_table_async( self, name: str, namespace_path: Optional[List[str]] = None ) -> Job: """Start dropping a table and return its cleanup job.""" - raise NotImplementedError + raise NotImplementedError( + "drop_table_async is not supported for this connection type" + ) def rename_table( self, @@ -777,19 +785,26 @@ class DBConnection(EnforceOverrides): cur_namespace_path = [] if new_namespace_path is None: new_namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "rename_table is not supported for this connection type" + ) def drop_database(self): """ Drop database This is the same thing as dropping all the tables """ - raise NotImplementedError + raise NotImplementedError( + "drop_database is not supported for this connection type" + ) def drop_all_tables(self, namespace_path: Optional[List[str]] = None): """ Drop all tables from the database + Remote connections do not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + Parameters ---------- namespace_path: List[str], optional @@ -798,7 +813,9 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_all_tables is not supported for this connection type" + ) @property def uri(self) -> str: diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 08125cc42..5e9288471 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -841,6 +841,20 @@ class RemoteDBConnection(DBConnection): job = LOOP.run(self._conn.drop_table_async(name, namespace_path=namespace_path)) return Job(job if isinstance(job, AsyncJob) else AsyncJob(job)) + @override + def drop_all_tables(self, namespace_path: Optional[List[str]] = None): + """Drop all tables from the database. + + The remote API does not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + """ + LOOP.run(self._conn.drop_all_tables(namespace_path=namespace_path)) + + @override + def drop_database(self): + """Deprecated alias for ``drop_all_tables()``.""" + LOOP.run(self._conn.drop_database()) + @override def rename_table( self, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 0744410b3..53591182e 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -7,6 +7,7 @@ import logging from functools import cached_property import os from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -82,6 +83,10 @@ from ..table import ( from ..types import BaseTokenizerType from ..util import infer_vector_column_name +if TYPE_CHECKING: + import lance + import polars as pl + class RemoteTable(Table): def __init__( @@ -243,12 +248,32 @@ class RemoteTable(Table): def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """to_arrow() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_arrow() is not supported for remote tables.") + raise NotImplementedError( + "to_arrow() is not supported for remote tables. " + "Use table.search().to_arrow() instead." + ) def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs): """to_pandas() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_pandas() is not supported for remote tables.") + raise NotImplementedError( + "to_pandas() is not supported for remote tables. " + "Use table.search().to_pandas() instead." + ) + + def to_lance(self, **kwargs) -> "lance.LanceDataset": + """Direct Lance dataset access is not supported for remote tables. + + Use ``table.search().to_arrow()`` to read query results through the server. + """ + return LOOP.run(self._table.to_lance(**kwargs)) + + def to_polars(self, **kwargs) -> "pl.LazyFrame": + """Use ``table.search().to_polars()`` to read remote query results.""" + raise NotImplementedError( + "to_polars() is not supported for remote tables. " + "Use table.search().to_polars() instead." + ) def checkout(self, version: Union[int, str]): result = LOOP.run(self._table.checkout(version)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index b9557f219..c08c6390e 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1356,13 +1356,13 @@ class Table(ABC): @abstractmethod def name(self) -> str: """The name of this Table""" - raise NotImplementedError + raise NotImplementedError("name is not supported for this table type") @property @abstractmethod def version(self) -> int: """The version of this Table""" - raise NotImplementedError + raise NotImplementedError("version is not supported for this table type") @property @abstractmethod @@ -1371,7 +1371,7 @@ class Table(ABC): of this Table """ - raise NotImplementedError + raise NotImplementedError("schema is not supported for this table type") @property @abstractmethod @@ -1400,7 +1400,7 @@ class Table(ABC): tags = table.tags.list() """ - raise NotImplementedError + raise NotImplementedError("tags is not supported for this table type") @property def branches(self) -> "Branches": @@ -1409,11 +1409,11 @@ class Table(ABC): Branches are isolated, writable lines of history forked from another branch (or version). Writes on a branch do not affect ``main``. """ - raise NotImplementedError + raise NotImplementedError("branches is not supported for this table type") def current_branch(self) -> Optional[str]: """The branch this table handle is scoped to, or ``None`` for ``main``.""" - raise NotImplementedError + raise NotImplementedError("current_branch is not supported for this table type") def __len__(self) -> int: """The number of rows in this Table""" @@ -1436,7 +1436,7 @@ class Table(ABC): filter: str, optional A SQL where clause to filter the rows to count. """ - raise NotImplementedError + raise NotImplementedError("count_rows is not supported for this table type") def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs) -> "pandas.DataFrame": """Return the table as a pandas DataFrame. @@ -1471,20 +1471,26 @@ class Table(ABC): ------- pa.Table """ - raise NotImplementedError + raise NotImplementedError("to_arrow is not supported for this table type") def to_lance(self, **kwargs) -> lance.LanceDataset: - """Return the table as a lance.LanceDataset. + """Return a local table as a lance.LanceDataset. + + Remote tables cannot be opened as Lance datasets by the client. Use + ``table.search().to_arrow()`` to read them through the server. Returns ------- lance.LanceDataset """ - raise NotImplementedError + raise NotImplementedError("to_lance is not supported for this table type") def to_polars(self, **kwargs) -> "pl.LazyFrame": """Return the table as a Polars LazyFrame. + For remote tables, use ``table.search().to_polars()`` to convert query + results to a Polars DataFrame. + Note ---- The Polars streaming engine is not supported because it does not currently @@ -1495,7 +1501,7 @@ class Table(ABC): ------- polars.LazyFrame """ - raise NotImplementedError + raise NotImplementedError("to_polars is not supported for this table type") # New unified API overload @overload @@ -1600,7 +1606,7 @@ class Table(ABC): ... "l2", vector_column_name="vector" ... ) """ - raise NotImplementedError + raise NotImplementedError("create_index is not supported for this table type") def create_index_async( self, @@ -1618,7 +1624,9 @@ class Table(ABC): be complete when returned; callers must not assume the index exists until :meth:`Job.wait` returns. """ - raise NotImplementedError + raise NotImplementedError( + "create_index_async is not supported for this table type" + ) def drop_index(self, name: str) -> None: """ @@ -1638,7 +1646,7 @@ class Table(ABC): Use [list_indices][lancedb.table.Table.list_indices] to find the names of the indices. """ - raise NotImplementedError + raise NotImplementedError("drop_index is not supported for this table type") def wait_for_index( self, index_names: Iterable[str], timeout: timedelta = timedelta(seconds=300) @@ -1655,14 +1663,14 @@ class Table(ABC): timeout: timedelta Timeout to wait for asynchronous indexing. The default is 5 minutes. """ - raise NotImplementedError + raise NotImplementedError("wait_for_index is not supported for this table type") @abstractmethod def stats(self) -> TableStatistics: """ Retrieve table and fragment statistics. """ - raise NotImplementedError + raise NotImplementedError("stats is not supported for this table type") @abstractmethod def create_scalar_index( @@ -1727,7 +1735,9 @@ class Table(ABC): ``my_col = 0 OR not_indexed = 1`` will not be able to use any scalar index on ``my_col``. """ - raise NotImplementedError + raise NotImplementedError( + "create_scalar_index is not supported for this table type" + ) def create_fts_index( self, @@ -1853,7 +1863,9 @@ class Table(ABC): ``LANCE_LANGUAGE_MODEL_HOME`` to override the default platform data directory under ``lance/language_models``. """ - raise NotImplementedError + raise NotImplementedError( + "create_fts_index is not supported for this table type" + ) @abstractmethod def add( @@ -1927,7 +1939,7 @@ class Table(ABC): AddResult An object containing the new version number of the table after adding data. """ - raise NotImplementedError + raise NotImplementedError("add is not supported for this table type") def merge_insert(self, on: Union[str, Iterable[str]]) -> LanceMergeInsertBuilder: """ @@ -2079,7 +2091,7 @@ class Table(ABC): - and also the "_distance" column which is the distance between the query vector and the returned vector. """ - raise NotImplementedError + raise NotImplementedError("search is not supported for this table type") @abstractmethod def take_offsets( @@ -2365,7 +2377,7 @@ class Table(ABC): x vector 0 3 [5.0, 6.0] """ - raise NotImplementedError + raise NotImplementedError("delete is not supported for this table type") @abstractmethod def update( @@ -2432,7 +2444,7 @@ class Table(ABC): 1 4 [5.0, 6.0] 2 3 [10.0, 10.0] """ - raise NotImplementedError + raise NotImplementedError("update is not supported for this table type") @abstractmethod def cleanup_old_versions( diff --git a/python/python/tests/test_errors.py b/python/python/tests/test_errors.py index 338b6a0d7..0e362e5d8 100644 --- a/python/python/tests/test_errors.py +++ b/python/python/tests/test_errors.py @@ -3,7 +3,28 @@ import pickle +import pytest + +from lancedb.db import DBConnection from lancedb.remote.errors import HttpError, LanceDBClientError, RetryError +from lancedb.table import Table + + +@pytest.mark.parametrize( + "base, method", + [ + (Table, "to_lance"), + (Table, "to_polars"), + (DBConnection, "drop_all_tables"), + (DBConnection, "drop_database"), + ], +) +def test_base_unsupported_methods_have_messages(base, method): + with pytest.raises(NotImplementedError) as exc_info: + getattr(base, method)(None) + message = str(exc_info.value) + assert method in message + assert "not supported" in message def test_pickle_lancedb_client_error(): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 9655f5807..c1db92948 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -857,10 +857,116 @@ def test_table_unimplemented_functions(): with mock_lancedb_connection(handler) as db: table = db.create_table("test", [{"id": 1}]) - with pytest.raises(NotImplementedError): - table.to_arrow() - with pytest.raises(NotImplementedError): - table.to_pandas() + for method in ["to_arrow", "to_pandas", "to_polars"]: + with pytest.raises(NotImplementedError) as exc_info: + getattr(table, method)() + message = str(exc_info.value) + assert method in message + assert "not" in message and "supported" in message + assert "remote tables" in message + assert "cloud" not in message.lower() + assert f"search().{method}()" in message + + +@pytest.mark.parametrize("method", ["to_arrow", "to_pandas"]) +@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes", "invalid"]) +def test_remote_conversion_blob_mode_errors(method, blob_mode): + with blob_remote_table() as table: + if blob_mode == "invalid": + with pytest.raises(ValueError, match="blob_mode must be one of"): + getattr(table, method)(blob_mode=blob_mode) + else: + with pytest.raises(NotImplementedError) as error: + getattr(table, method)(blob_mode=blob_mode) + assert f"search().{method}()" in str(error.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [None, 1]) +@pytest.mark.parametrize("expose_location", [False, True]) +async def test_remote_to_lance_rejection( + tmp_path, version, expose_location, monkeypatch +): + import lance + + location = tmp_path / "test.lance" + lance.write_dataset(pa.table({"id": [1]}), location) + lance.write_dataset(pa.table({"id": [2]}), location, mode="append") + requests = [] + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) + + def handler(request): + requests.append(request.path) + assert request.path == "/v1/table/test/describe/" + content_len = int(request.headers.get("Content-Length", 0)) + body = json.loads(request.rfile.read(content_len)) if content_len else {} + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + response = { + "version": body.get("version") or 2, + "schema": {"fields": []}, + } + if expose_location: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) + + options = {"default_scan_options": {"with_row_id": True}} + with mock_lancedb_connection(handler) as db: + table = db.open_table("test", version=version) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as sync_error: + table.to_lance(**options) + assert requests == opened_requests + + async with mock_lancedb_connection_async(handler) as db: + table = await db.open_table("test", version=version) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as async_error: + await table.to_lance(**options) + assert requests == opened_requests + + assert "to_lance() is not supported for remote tables" in str(sync_error.value) + assert "query the server" in str(sync_error.value) + assert str(sync_error.value) == str(async_error.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method, kwargs", + [ + ("drop_all_tables", {}), + ("drop_all_tables", {"namespace_path": []}), + ("drop_all_tables", {"namespace_path": ["team"]}), + ("drop_database", {}), + ], +) +async def test_remote_drop_all_tables_not_supported(method, kwargs): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(500) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError) as sync_error: + getattr(db, method)(**kwargs) + + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError) as async_error: + await getattr(db, method)(**kwargs) + + assert "Dropping all tables is not currently supported in the remote API" in str( + sync_error.value + ) + assert "drop_table" in str(sync_error.value) + assert str(sync_error.value) == str(async_error.value) + assert requests == [] def test_table_to_pandas_not_supported(): @@ -1681,10 +1787,13 @@ def test_remote_python_only_messages(caplog): table = db.open_table("test", storage_options={}, index_cache_size=1) assert "storage_options is ignored for remote tables" in caplog.text assert "index_cache_size is ignored for remote tables" in caplog.text - for method in ["to_arrow", "to_pandas"]: + for method in ["to_arrow", "to_pandas", "to_polars"]: with pytest.raises(NotImplementedError) as error: getattr(table, method)() - assert str(error.value) == f"{method}() is not supported for remote tables." + assert str(error.value) == ( + f"{method}() is not supported for remote tables. " + f"Use table.search().{method}() instead." + ) for method in ["compact_files", "cleanup_old_versions"]: with pytest.warns(UserWarning) as warning: getattr(table, method)() @@ -3490,14 +3599,17 @@ def blob_descriptor_take_table(): @contextlib.contextmanager -def blob_remote_table(*, server_version=Version("0.5.0")): +def blob_remote_table(*, server_version=Version("0.5.0"), location=None): def handler(request): if request.path == "/v1/table/test/describe/": request.send_response(200) request.send_header("Content-Type", "application/json") request.send_header("phalanx-version", str(server_version)) request.end_headers() - request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode()) + response = dict(BLOB_DESCRIBE_RESPONSE) + if location is not None: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) elif request.path.startswith("/v1/table/test/blob/image/"): path = request.path.partition("?")[0] row_id = int(path.split("/")[-2]) @@ -3532,7 +3644,6 @@ def blob_remote_table(*, server_version=Version("0.5.0")): if columns == ["id"]: response_table = blob_query_response_table().select(["id"]) else: - assert body["with_row_id"] is True if columns == ["image"]: # Blob handles are sized from a descriptor take. assert body["filter"].startswith("_rowid IN") @@ -3545,6 +3656,8 @@ def blob_remote_table(*, server_version=Version("0.5.0")): ) else: response_table = blob_query_response_table() + if not body.get("with_row_id", False): + response_table = response_table.drop(["_rowid"]) request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.file") request.send_header("x-lancedb-version", "1") @@ -3579,6 +3692,32 @@ def test_remote_blob_columns_and_fetch(): assert blobs.to_pylist() == [b"alpha", None, b"gamma"] +@pytest.mark.parametrize("blob_mode", ["descriptions", "bytes", "lazy"]) +def test_sync_remote_blob_pandas_avoids_storage_access( + tmp_path, monkeypatch, blob_mode +): + import lance + + location = tmp_path / "remote.lance" + lance.write_dataset(pa.table({"id": [7], "image": [b"server-owned"]}), location) + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) + with blob_remote_table(location=location) as table: + monkeypatch.setattr(table, "to_lance", fail_dataset) + query = table.search().select(["id", "image"]) + if blob_mode == "descriptions": + df = query.to_pandas(blob_mode=blob_mode) + assert df["id"].tolist() == [1, 2, 3] + assert "_rowid" not in df.columns + assert df["image"][0]["size"] == 5 + else: + with pytest.raises(NotImplementedError, match="stable table snapshot"): + query.to_pandas(blob_mode=blob_mode) + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True]) async def test_remote_fetch_blob_ranges(use_async): diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 8a456ab3b..cad60ab7d 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1707,7 +1707,9 @@ impl Database for RemoteDatabase { // TODO: Implement namespace-aware drop_all_tables let _namespace_path = namespace_path; // Suppress unused warning for now Err(crate::Error::NotSupported { - message: "Dropping all tables is not currently supported in the remote API".to_string(), + message: "Dropping all tables is not currently supported in the remote API. \ + List tables and call drop_table for each table instead." + .to_string(), }) } From 6777a40dc280be8fd8f55e67750cbc6ecf7a2a60 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 24 Oct 2026 09:51:10 +0800 Subject: [PATCH 17/36] fix: support remote tables in storage root (#3881) ## Summary - allow `db://` to select the remote storage root when a host override supplies the endpoint - preserve the empty `x-lancedb-database` header used by the legacy clients for root storage - reject root-storage connections without a host override before constructing an invalid Cloud URL - document the supported URI and cover both accepted and rejected configurations ## Root cause The shared Rust remote URL parser required every `db://` URI to contain a host. That validation was introduced when the Node and Python clients moved to the Rust implementation, but the legacy Node client intentionally treated the empty host in `db://` as the storage-root selector. ## Validation - `cargo fmt --all` - `cargo test --quiet --features remote -p lancedb test_connect_remote_storage_root` - `cargo check --quiet --features remote --tests --examples` - `cargo clippy --quiet --features remote --tests --examples` - `cargo test --quiet --features remote --tests` Fixes #1944 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- rust/lancedb/src/connection.rs | 34 +++++++++++++++++++++++++++++++ rust/lancedb/src/lib.rs | 1 + rust/lancedb/src/remote/client.rs | 31 +++++++++++++++++----------- rust/lancedb/src/remote/db.rs | 7 +++++++ 4 files changed, 61 insertions(+), 12 deletions(-) diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 124f6a54a..482acf89a 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -1270,6 +1270,7 @@ pub struct ConnectRequest { /// - `/path/to/database` - local database on file system. /// - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud object store /// - `db://dbname` - LanceDB Cloud + /// - `db://` with a host override - remote tables in the storage root pub uri: String, #[cfg(feature = "remote")] @@ -1382,6 +1383,8 @@ impl ConnectBuilder { /// /// This option is only used when connecting to LanceDB Cloud (db:// URIs) /// and will be ignored for other URIs. + /// Use the URI `db://` together with a host override to connect to remote + /// tables stored directly in the storage root. /// /// # Arguments /// @@ -2007,6 +2010,37 @@ mod tests { } } + #[cfg(feature = "remote")] + #[tokio::test] + async fn test_connect_remote_storage_root() { + let conn = ConnectBuilder::new("db://") + .region("us-east-1") + .api_key("my-api-key") + .host_override("https://example.com") + .execute() + .await + .unwrap(); + + let (impl_name, properties) = conn.namespace_client_config().await.unwrap(); + assert_eq!(impl_name, "rest"); + assert_eq!(properties["uri"], "https://example.com"); + assert!( + !properties.contains_key("header.x-lancedb-database"), + "storage-root connections must not send the x-lancedb-database header" + ); + + let result = ConnectBuilder::new("db://") + .region("us-east-1") + .api_key("my-api-key") + .execute() + .await; + assert!(matches!( + result, + Err(Error::InvalidInput { message }) + if message.contains("A host override is required") + )); + } + #[cfg(feature = "remote")] #[tokio::test] async fn test_connect_rejects_header_provider_with_oauth_config() { diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index e15017e7f..68521ff2c 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -54,6 +54,7 @@ //! - `/path/to/database` - local database on file system. //! - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud object store //! - `db://dbname` - Lance Cloud +//! - `db://` with a host override - remote tables in the storage root //! //! You can also use [`ConnectBuilder`] to configure the connection to the database. //! diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 783d2ef84..fe4a2072e 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -378,23 +378,30 @@ pub struct ParsedDbUrl { /// Parse a database URL and extract the database name and optional prefix. /// -/// Expected format: `db://db_name` or `db://db_name/prefix` +/// Expected format: `db://db_name`, `db://db_name/prefix`, or `db://` when +/// connecting to the storage root through a host override. pub fn parse_db_url(db_url: &str) -> Result { let parsed_url = url::Url::parse(db_url).map_err(|err| Error::InvalidInput { message: format!("db_url is not a valid URL. '{db_url}'. Error: {err}"), })?; debug_assert_eq!(parsed_url.scheme(), "db"); - if !parsed_url.has_host() { - return Err(Error::InvalidInput { - message: format!("Invalid database URL (missing host) '{}'", db_url), - }); - } - let db_name = urlencoding::decode(parsed_url.host_str().unwrap()) - .map_err(|err| Error::InvalidInput { - message: format!("Invalid encoded database name: {err}"), - })? - .into_owned(); - validate_database_name(&db_name)?; + let db_name = match parsed_url.host_str() { + Some(db_name) => { + let db_name = urlencoding::decode(db_name) + .map_err(|err| Error::InvalidInput { + message: format!("Invalid encoded database name: {err}"), + })? + .into_owned(); + validate_database_name(&db_name)?; + db_name + } + None if matches!(parsed_url.path(), "" | "/") => String::new(), + None => { + return Err(Error::InvalidInput { + message: format!("Invalid database URL (missing host) '{}'", db_url), + }); + } + }; let db_prefix = { let prefix = parsed_url.path().trim_start_matches('/'); if prefix.is_empty() { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index cad60ab7d..f07d20fc3 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -318,6 +318,13 @@ impl RemoteDatabase { read_consistency_interval: Option, ) -> Result { let parsed = super::client::parse_db_url(uri)?; + if parsed.db_name.is_empty() && host_overrides.rest.is_none() { + return Err(Error::InvalidInput { + message: + "A host override is required when connecting to the storage root with 'db://'" + .to_string(), + }); + } Self::try_new_with_identity( uri, api_key, From a8305e072246d21a97e8857c480e3bbdb9a4178f Mon Sep 17 00:00:00 2001 From: Romin Parekh <3831468+rominparekh@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:08:40 -0700 Subject: [PATCH 18/36] fix(rust): accept usize::MAX as a local query limit (#4567) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #2211. Local queries converted `limit`/`offset` with `as i64`, so `limit(usize::MAX)` wrapped to `-1` and Lance failed with `Limit must be non-negative`. The remote client already worked around this (#2354); this fixes the local path. - Saturate scanner offsets to `i64::MAX` and cap limits at `i64::MAX - offset`, keeping both conversion and full-text pagination addition safe. - Use `saturating_add` for vector `top_k = limit + offset`, then cap candidates at the dataset row count to avoid indexed-search capacity and refinement overflows. Keep a positive candidate bound for empty tables because Lance rejects `k = 0`. - Leave public APIs and remote query handling unchanged. ### Testing - Regression coverage includes plain scans, indexed IVF search, refinement, index bypass, empty vector tables, and indexed full-text pagination with finite and extreme limits/offsets. - Reproduced the indexed-search and FTS overflows before applying the follow-up fix. - `cargo test --quiet --features remote -p lancedb --lib query::` — 81 passed. - `cargo clippy --quiet --features remote -p lancedb --tests --examples` — clean. - `cargo fmt --all` and `git diff --check` — clean. Vector planning now reads the dataset row count (metadata fast path) to bound candidate allocation; this was not performance-benchmarked. --------- Co-authored-by: Romin Parekh Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 29 +++--- rust/lancedb/src/query.rs | 153 ++++++++++++++++++++++++++++++++ rust/lancedb/src/table/query.rs | 21 ++++- 3 files changed, 188 insertions(+), 15 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 9d7610f72..52f38d077 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -1561,9 +1561,10 @@ class LanceQueryBuilder(ABC): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -1619,8 +1620,8 @@ class LanceQueryBuilder(ABC): GlobalLimitExec: elapsed=..., skip=0, fetch=10, metrics=[output_rows=..., elapsed_compute=..., output_bytes=...] FilterExec: elapsed=..., _distance@2 IS NOT NULL, metrics=[...] - SortExec: elapsed=..., TopK(fetch=10), expr=[...], - preserve_partitioning=[...], + SortExec: elapsed=..., TopK(fetch=1), expr=[...], + preserve_partitioning=[...], filter=[...], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., row_replacements=...] KNNVectorDistance: elapsed=..., metric=l2, @@ -1633,6 +1634,7 @@ class LanceQueryBuilder(ABC): metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., fragments_scanned=..., ranges_scanned=1, rows_scanned=1, bytes_read=..., iops=..., requests=..., task_wait_time=...] + Parameters ---------- @@ -2959,9 +2961,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -3601,9 +3604,10 @@ class AsyncQueryBase(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters @@ -4521,9 +4525,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): LanceRead: uri=..., projection=[text], source=stream(_rowid) GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score, _rowid@0 as _rowid] LanceRead: uri=..., projection=[vector, text], source=stream(_rowid) GlobalLimitExec: skip=0, fetch=10 @@ -4877,9 +4881,10 @@ class BaseQueryBuilder(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 00a2b88f3..9bac51209 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -2714,6 +2714,159 @@ mod tests { } } + #[tokio::test] + async fn test_limit_usize_max() { + let conn = connect("memory://foo").execute().await.unwrap(); + let table = conn + .create_table("my_table", make_non_empty_batches()) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .offset(1) + .nearest_to(&[0.1; 4]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected - 1); + } + + #[tokio::test] + async fn test_indexed_vector_limit_usize_max() { + use crate::index::vector::IvfFlatIndexBuilder; + + let conn = connect("memory://indexed_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_non_empty_batches()) + .execute() + .await + .unwrap(); + table + .create_index( + &["vector"], + Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(1)), + ) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + for limit in [expected, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, usize::MAX] { + for (bypass, refine_factor) in [(false, None), (false, Some(2)), (true, None)] { + let query = table + .query() + .limit(limit) + .offset(offset) + .nearest_to(&[0.1; 4]) + .unwrap(); + let query = if bypass { + query.bypass_vector_index() + } else { + query + }; + let query = if let Some(factor) = refine_factor { + query.refine_factor(factor) + } else { + query + }; + let batches: Vec = + query.execute().await.unwrap().try_collect().await.unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected.saturating_sub(offset), + "limit={limit}, offset={offset}, bypass={bypass}, refine={refine_factor:?}" + ); + } + } + } + } + + #[tokio::test] + async fn test_empty_vector_limit_usize_max() { + let conn = connect("memory://empty_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_test_batches()) + .execute() + .await + .unwrap(); + let batches: Vec = table + .query() + .limit(usize::MAX) + .nearest_to(&[0.1; 128]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + + #[tokio::test] + async fn test_fts_limit_usize_max_with_offset() { + let conn = connect("memory://fts_limit").execute().await.unwrap(); + let data = RecordBatch::try_from_iter([( + "text", + Arc::new(StringArray::from_iter_values((0..32).map(|_| "hello"))) + as arrow_array::ArrayRef, + )]) + .unwrap(); + let table = conn + .create_table("documents", data) + .execute() + .await + .unwrap(); + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + + for limit in [32, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, i64::MAX as usize, usize::MAX] { + let batches: Vec = table + .query() + .full_text_search(FullTextSearchQuery::new("hello".into())) + .limit(limit) + .offset(offset) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 32usize.saturating_sub(offset), + "limit={limit}, offset={offset}" + ); + } + } + } + #[tokio::test] async fn test_repeated_only_if_combines_with_and() { use crate::expr::{col, lit}; diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index f625a4aff..4ad840350 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -280,7 +280,13 @@ pub async fn create_plan( let (_, element_type) = lance::index::vector::utils::get_vector_type(schema, &column)?; let is_binary = matches!(element_type, DataType::UInt8); - let top_k = query.base.limit.unwrap_or(DEFAULT_TOP_K) + query.base.offset.unwrap_or(0); + let top_k = query + .base + .limit + .unwrap_or(DEFAULT_TOP_K) + .saturating_add(query.base.offset.unwrap_or(0)) + // Lance allocates candidates from k; no search can return more rows than the table. + .min(ds_ref.count_rows(None).await?.max(1)); if is_binary { let query_vector = arrow::compute::cast(&query_vector, &DataType::UInt8)?; @@ -317,9 +323,18 @@ pub async fn create_plan( // For a batch query, `nearest` already applies k to each query vector. // Adding Scanner's global limit would truncate the combined result to k rows. if !is_batch_query { + let offset = query + .base + .offset + .map(|offset| i64::try_from(offset).unwrap_or(i64::MAX)); scanner.limit( - query.base.limit.map(|limit| limit as i64), - query.base.offset.map(|offset| offset as i64), + // Lance's FTS planner adds limit + offset as i64. + query.base.limit.map(|limit| { + i64::try_from(limit) + .unwrap_or(i64::MAX) + .min(i64::MAX - offset.unwrap_or(0)) + }), + offset, )?; } From cf5c9b5453e5abb30341d33ca1b67a7812b58ecb Mon Sep 17 00:00:00 2001 From: J Ganesh Kumar Reddy Date: Sat, 24 Oct 2026 10:01:49 +0530 Subject: [PATCH 19/36] fix: validate cached remote tables on open (#4557) ## Summary Fixes #4401. - Validate cached remote table entries with the existing `describe` endpoint before returning a handle. - Preserve cached server-version behavior after a successful validation. - Add a regression test for a table dropped by another connection after its entry has been cached. A cached entry can otherwise outlive an external drop, allowing `open_table()` to return a handle that fails only on its first later operation. ## Verification [Linux validation](https://github.com/Ganesh1729-ui/lancedb/actions/runs/39420560974) passed against the exact runtime and test sources at PR head `d39092b01f6b49921b5d66ccedf29709a7ddc23d`, using the repository's Rust toolchain and `ci` profile: - `cargo fmt --all -- --check` - `cargo test --profile ci --locked -p lancedb --features remote --lib remote::db::tests::test_open_table -- --nocapture`: **8 passed**, including cached-table deletion, independent handle checkout, branch/version handling, schema seeding and describe-response handling. - Removing only the cached-hit existence guard makes `test_open_table_checks_that_a_cached_table_still_exists` fail at its `TableNotFound` assertion; restoring the fix makes the same test pass again. - `git diff --check` The validation workflow is on a separate fork branch and is excluded from this PR. These are focused, mock-backed remote unit tests; the full upstream matrix and live Cloud integration tests still need upstream CI approval. Local Windows compilation was unavailable because the MSVC environment has no `link.exe`; the linked Linux run completed successfully. --------- Co-authored-by: Ganesh172919 Co-authored-by: Xuanwo --- deny.toml | 10 +++++++++ rust/lancedb/src/remote/db.rs | 42 +++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/deny.toml b/deny.toml index 7d60ae63f..77412709f 100644 --- a/deny.toml +++ b/deny.toml @@ -110,8 +110,18 @@ ignore = [ # compact_str/smol_str, so clearing this requires polars to migrate. # https://rustsec.org/advisories/RUSTSEC-2026-0249 { id = "RUSTSEC-2026-0249", reason = "smartstring unmaintained via polars; no fixed upstream release" }, + # crypto-bigint < 0.7.6: timing side-channels in constant-time ops, + # incl. subtraction, zero counting, division and Montgomery multiplication. + # Reached only through aws-sdk-s3 sigv4a -> aws-sigv4 (p256/ecdsa); signed + # requests carry the SigV4a credential-derived key. The patched 0.7.6 was + # published 2026-10-23, but no published aws-sigv4 release uses it yet + # (latest 1.6.3 pins crypto-bigint ^0.5), so clearing requires smithy-rs to + # ship an updated aws-sigv4. + # https://rustsec.org/advisories/RUSTSEC-2026-0387 + { id = "RUSTSEC-2026-0387", reason = "crypto-bigint via aws-sigv4 (sigv4a); no fixed upstream aws-sigv4 release yet" }, ] + # --------------------------------------------------------------------------- # Licenses: only allow licenses we've reviewed as compatible with Apache-2.0. # --------------------------------------------------------------------------- diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index f07d20fc3..768de79d2 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1614,6 +1614,17 @@ impl Database for RemoteDatabase { // Every open gets its own checkout, schema cache, and freshness state. if let Some(version) = self.table_cache.get(&cache_key).await { + // A cached version can outlive a drop from another connection. Verify that + // the table still exists before returning a handle for it. + let req = self + .client + .post(&format!("/v1/table/{}/describe/", identifier)) + .json(&serde_json::json!({})); + 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?; + Ok(Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), request.name, @@ -2557,6 +2568,37 @@ mod tests { assert_eq!(table.name(), "table1"); } + #[tokio::test] + async fn test_open_table_checks_that_a_cached_table_still_exists() { + let describe_calls = Arc::new(AtomicUsize::new(0)); + let calls = describe_calls.clone(); + let mut db = super::RemoteDatabase::new_mock(move |request| { + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + match calls.fetch_add(1, Ordering::SeqCst) { + 0 => http::Response::builder() + .status(200) + .body(r#"{"version": 1}"#.to_string()) + .unwrap(), + 1 => http::Response::builder() + .status(404) + .body("table not found".to_string()) + .unwrap(), + call => panic!("unexpected describe call: {call}"), + } + }); + db.table_cache = moka::future::Cache::new(10); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + + conn.open_table("table1").execute().await.unwrap(); + let result = conn.open_table("table1").execute().await; + + assert!(matches!(result, Err(Error::TableNotFound { .. }))); + assert_eq!(describe_calls.load(Ordering::SeqCst), 2); + } + #[tokio::test] async fn test_open_table_checkout_is_independent_per_handle() { let latest = Arc::new(AtomicUsize::new(2)); From 2c9f2e75a068a960863f3c45c36334d965bf2c1c Mon Sep 17 00:00:00 2001 From: Yeluru Nitish <141549775+Nitish-1303@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:21:48 -0500 Subject: [PATCH 20/36] fix: keep FTS indexes in list_indices when their files are missing (#4251) Fixes #4250 `list_indices()` loads each FTS index's segment params from the index files to fill in `index_details`. That call used `?`, so a single FTS index with missing files made the whole listing fail. A BTree or Bitmap index with missing files still lists fine, because it never takes that path. Now, if the params can't be loaded, it logs a warning and keeps the entry with the details from the manifest. Callers can still see the damaged index and rebuild or drop it. Added a Python test based on the repro in the issue. It fails on 0.39.0 with the same RuntimeError. --------- Co-authored-by: Xuanwo --- python/python/tests/test_fts.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index d7fa2e0ec..1dc2fd9a1 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -391,6 +391,20 @@ def test_custom_stop_words_list(table): ) +def test_list_indices_with_missing_fts_files(table): + # Regression test for https://github.com/lancedb/lancedb/issues/4250 + table.create_fts_index("text") + table.create_scalar_index("id") + + fts = next(i for i in table.list_indices() if i.index_type == "FTS") + shutil.rmtree(Path(table.uri) / "_indices" / fts.index_uuid) + + indices = {i.name: i for i in table.list_indices()} + assert set(indices) == {"text_idx", "id_idx"} + assert indices["text_idx"].index_type == "FTS" + assert indices["text_idx"].columns == ["text"] + + def test_search_fts(table): table.create_fts_index("text") results = table.search("puppy").select(["id", "text"]).limit(5).to_list() From 0ea932dafdac1ba3f86135cded890012d9733fdc Mon Sep 17 00:00:00 2001 From: Shiduo Li Date: Fri, 23 Oct 2026 23:47:41 -0700 Subject: [PATCH 21/36] fix(rust): compile jobs without remote feature (#4103) ## What Move typed job decode error construction behind a feature-aware helper so builds without the `remote` feature do not reference `Error::Http`. Local decode failures continue to surface as `Error::Runtime`, while remote builds retain HTTP errors with request IDs. Add a regression test for local typed-result decode failures. Closes #4096. ## Testing - `cargo fmt --all -- --check` - `cargo check -p lancedb --no-default-features` - `cargo test -p lancedb --no-default-features job::tests::local_typed_result_decode_errors_are_runtime_errors` Co-authored-by: Xuanwo --- rust/lancedb/src/job.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 6865e59ee..d671365f9 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -636,6 +636,14 @@ mod tests { use super::*; + #[test] + fn local_typed_result_decode_errors_are_runtime_errors() { + assert!(matches!( + TerminalResult::local(Value::Null).decode::(), + Err(Error::Runtime { .. }) + )); + } + #[tokio::test] async fn mapped_spawned_job_reuses_outcome() { let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1); From 6b6778e68495466e78cc786ebc912601c57c8b5d Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sat, 24 Oct 2026 18:50:27 -0700 Subject: [PATCH 22/36] feat: infer lazy blob reader Function inputs (#4663) Infer Blob V2 input access from the Python Function signature. `video: BlobReader` requests a lazy, seekable reader; `video: bytes` requests the complete payload. Users do not need a separate `blob_inputs` registration option. ## Python experience For a video table joined to a table of frame timestamps, the Function can pass its reader directly to a media decoder: ```python import io import av from lancedb import BlobReader, udf @udf(pip=["av", "Pillow"]) def extract_frame(video: BlobReader, timestamp_s: float) -> bytes: with av.open(video) as container: stream = container.streams.video[0] container.seek(int(timestamp_s / stream.time_base), stream=stream) for frame in container.decode(stream): if frame.time is not None and frame.time >= timestamp_s: output = io.BytesIO() frame.to_image().save(output, format="JPEG") return output.getvalue() raise ValueError("No frame at or after the requested timestamp") ``` After registering `extract_frame`, the query keeps the usual Function call: ```sql SELECT f.video_id, f.timestamp_s, extract_frame(v.video, f.timestamp_s) AS image FROM videos AS v JOIN frames AS f ON v.id = f.video_id ``` The returned JPEG can feed a subsequent image-model step. The reader provides `read`, `readinto`, `seek`, `tell`, `size()`, and `read_range(offset, length)`; range reads leave the sequential cursor unchanged. Each row invocation owns its reader, which closes after the call. Decoder seeks can request the parts of the MP4 needed to locate and decode the requested frame; the amount read depends on the file layout and codec. Functions that need the whole value can keep a bytes signature: ```python @udf def video_length(video: bytes) -> int: return len(video) ``` ## Execution contract - Signature inference maps `BlobReader` to a Blob V2 field and records `blob_access="reader"`. `bytes` keeps its binary Arrow type and records `blob_access="bytes"`; it can also bind to a Blob V2 source. Explicit schemas support the same choice, including a Blob V2 field annotated as `bytes`. - Reader execution carries a compact blob reference into the Function worker. Payload reads happen when Python calls the reader, with requests split into at most 1 MiB ranges. The full video does not need to enter the Arrow input batch just to invoke the Function. - Bytes execution materializes the full source before the Function runs and remains subject to eager-input memory admission. Calling `reader.read()` without a size also materializes the remaining contents in Python; lazy access depends on how the Function consumes it. - Input access metadata survives Python/Rust Function serialization and registration. Packaged Functions import the worker-provided reader implementation. This PR defines the SDK contract; execution requires corresponding server support. - This changes blob inputs. The example still returns ordinary image bytes; it does not add streamed outputs or run an image model. Adds coverage for signature inference, explicit schemas, nullable readers, signature round trips, seeking, bounded range reads, buffering, and invalid responses. --- docs/src/python/python.md | 8 ++ docs/web/management/function.mdx | 34 +++++ python/python/lancedb/__init__.py | 1 + python/python/lancedb/_blob_reader.py | 120 ++++++++++++++++++ python/python/lancedb/functions.py | 80 +++++++++--- python/python/tests/test_blob_reader.py | 59 +++++++++ .../tests/test_first_class_function_slice2.py | 55 ++++++++ rust/lancedb/src/function.rs | 23 ++++ 8 files changed, 359 insertions(+), 21 deletions(-) create mode 100644 python/python/lancedb/_blob_reader.py create mode 100644 python/python/tests/test_blob_reader.py diff --git a/docs/src/python/python.md b/docs/src/python/python.md index c465d0ac4..5d1393924 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -304,6 +304,14 @@ listing a storage directory. ## Functions and Jobs +Annotate a blob input with `lancedb.BlobReader` for lazy, seekable access, or +with `bytes` to materialize its contents before the Function runs. Reader +inputs infer the Blob V2 type without an explicit input schema. Readers are +scoped to each invocation and support `read`, `readinto`, `seek`, `tell`, +`size`, and cursor-independent `read_range`. + +::: lancedb.BlobReader + ::: lancedb.functions.FunctionArtifact ::: lancedb.functions.FunctionParameter diff --git a/docs/web/management/function.mdx b/docs/web/management/function.mdx index 8e5b53e4a..6df99b5fd 100644 --- a/docs/web/management/function.mdx +++ b/docs/web/management/function.mdx @@ -205,6 +205,40 @@ Function output blobs follow the same compaction and cleanup lifecycle as ingested Managed blobs. Retained versions continue to own their payloads; cleanup reclaims an object only after its retained references are gone. +### Draft: incremental blob inputs (upcoming) + +This feature requires an SDK that exports `BlobReader` and a compatible +Enterprise runtime with reader input support enabled. Coordinate the runtime +upgrade before registering or executing reader Functions. + +Annotate a row Function parameter with `BlobReader` to read a Blob v2 input +incrementally. The signature selects the access mode without another +registration option. For example, read the first MP4 box type without loading +the whole video: + +```python +from lancedb import BlobReader, udf + + +@udf +def first_box_type(video: BlobReader) -> bytes: + video.seek(4) + return video.read(4) +``` + +The worker passes a seekable reader backed by the stored blob. `read(size)` +reads from the current cursor; `read_range(offset, length)` reads a byte range +without moving it. `size()` returns the payload size without reading it. +Each row call receives its own reader, which closes when the call returns or +raises. Do not retain it for another invocation. + +Use `video: bytes` when the Function needs the complete payload. Bytes inputs +materialize before invocation and remain subject to eager input limits. +Reader inputs keep the unread payload out of the input batch, but calling +`video.read()` without a size loads all remaining bytes in Python. Reader +inputs are supported in row mode. Return values remain ordinary values such +as `bytes`; this feature does not add streamed outputs. + ## Declare the environment {#environment} The function runs remotely, so its dependencies are part of its definition: diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 6b902c466..a3aa55209 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -24,6 +24,7 @@ from .remote import ClientConfig from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector +from ._blob_reader import BlobReader as BlobReader from .job import AsyncJob, Job from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery diff --git a/python/python/lancedb/_blob_reader.py b/python/python/lancedb/_blob_reader.py new file mode 100644 index 000000000..2fd0f00ca --- /dev/null +++ b/python/python/lancedb/_blob_reader.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Seekable binary input for Functions that consume Blob V2 values lazily.""" + +import io +import operator +import urllib.request +from collections.abc import Callable +from http import HTTPStatus + +_READ_CHUNK_BYTES = 1024 * 1024 +_READ_TIMEOUT_SECONDS = 300 + + +class BlobReader(io.RawIOBase): + """Read a blob incrementally inside a Function. + + Annotating a Function input with ``BlobReader`` requests a lazy, seekable + reader. Annotating it with ``bytes`` requests materialized bytes instead. + ``read_range`` does not change the sequential cursor. Readers are scoped + to the invocation and must not be retained for a later call. + + For local calls, construct a reader from a range-reading callable and the + blob's size. The callable accepts an offset and length and returns bytes. + """ + + def __init__(self, read_range: Callable[[int, int], bytes], size: int): + super().__init__() + self._read_range = read_range + self._size = operator.index(size) + if self._size < 0: + raise ValueError("blob size must be nonnegative") + self._position = 0 + + @classmethod + def _from_uri(cls, uri: str, size: int): + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + + def read_range(offset, length): + request = urllib.request.Request( + uri, headers={"Range": f"bytes={offset}-{offset + length - 1}"} + ) + with opener.open(request, timeout=_READ_TIMEOUT_SECONDS) as response: + if response.status != HTTPStatus.PARTIAL_CONTENT: + raise OSError( + "blob range service did not return a partial response" + ) + data = response.read(length + 1) + if len(data) != length: + raise OSError( + "blob range service returned an invalid payload length" + ) + return data + + return cls(read_range, size) + + def readable(self) -> bool: + return not self.closed + + def seekable(self) -> bool: + return not self.closed + + def size(self) -> int: + """Return the blob size without reading its payload.""" + self._checkClosed() + return self._size + + def tell(self) -> int: + self._checkClosed() + return self._position + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + self._checkClosed() + offset = operator.index(offset) + if whence == io.SEEK_SET: + position = offset + elif whence == io.SEEK_CUR: + position = self._position + offset + elif whence == io.SEEK_END: + position = self._size + offset + else: + raise ValueError(f"invalid whence: {whence}") + if position < 0: + raise ValueError("negative seek position") + self._position = position + return position + + def read_range(self, offset: int, length: int) -> bytes: + """Read a byte range without changing the cursor; truncate at EOF.""" + self._checkClosed() + offset, length = operator.index(offset), operator.index(length) + if offset < 0 or length < 0: + raise ValueError("offset and length must be nonnegative") + remaining = min(length, max(0, self._size - offset)) + chunks = [] + while remaining: + count = min(remaining, _READ_CHUNK_BYTES) + chunk = self._read_range(offset, count) + if len(chunk) != count: + raise OSError("blob range reader returned an invalid payload length") + chunks.append(chunk) + offset += count + remaining -= count + return b"".join(chunks) + + def read(self, size: int = -1) -> bytes: + self._checkClosed() + size = -1 if size is None else operator.index(size) + if size < 0: + size = max(0, self._size - self._position) + data = self.read_range(self._position, size) + self._position += len(data) + return data + + def readinto(self, buffer) -> int: + view = memoryview(buffer).cast("B") + data = self.read(len(view)) + view[: len(data)] = data + return len(data) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 7f5776b32..42c4b27e1 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -33,6 +33,7 @@ from typing import ( Annotated, Any, Callable, + Literal, Optional, Union, get_args, @@ -52,7 +53,9 @@ from pydantic import ( model_validator, ) +from .schema import blob as _blob_field from .schema import is_blob_v2_field as _is_blob_v2_field +from ._blob_reader import BlobReader from .secrets import EnvVarSecret _Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1) @@ -213,6 +216,7 @@ class FunctionParameter(_RemoteValue): name: str arrow_type: str nullable: bool + blob_access: Optional[Literal["bytes", "reader"]] = None class FunctionResultField(_OpenRemoteValue): @@ -1096,6 +1100,43 @@ def _infer_signature( if (input_schema is None) != (output_schema is None): raise ValueError("input_schema and output_schema must be provided together") + try: + annotations = get_type_hints(function, include_extras=True) + except Exception as error: + raise TypeError(f"failed to resolve Function annotations: {error}") from error + + def input_parameter(field): + annotation = annotations.get(field.name) + if get_origin(annotation) in (Union, types.UnionType): + choices = [ + value for value in get_args(annotation) if value is not type(None) + ] + if len(choices) == 1: + annotation = choices[0] + arrow_type = _canonical_arrow_field(field) + if annotation is BlobReader: + if arrow_type != _FUNCTION_BLOB_V2_TYPE: + raise TypeError( + f"BlobReader input '{field.name}' requires a Blob V2 field" + ) + access = "reader" + elif annotation is bytes: + if arrow_type != _FUNCTION_BLOB_V2_TYPE and not ( + pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type) + ): + raise TypeError( + f"bytes input '{field.name}' requires binary or Blob V2" + ) + access = "bytes" + else: + access = "bytes" if arrow_type == _FUNCTION_BLOB_V2_TYPE else None + return FunctionParameter( + name=field.name, + arrow_type=arrow_type, + nullable=field.nullable, + blob_access=access, + ) + if input_schema is not None: if not isinstance(input_schema, pa.Schema): raise TypeError("input_schema must be a PyArrow Schema") @@ -1110,20 +1151,9 @@ def _infer_signature( "input_schema fields must exactly match callable parameters in order: " f"expected {expected!r}, got {actual!r}" ) - inputs = tuple( - FunctionParameter( - name=field.name, - arrow_type=_canonical_arrow_field(field), - nullable=field.nullable, - ) - for field in input_schema - ) + inputs = tuple(input_parameter(field) for field in input_schema) return FunctionSignature(inputs=inputs, output=_function_output(output_schema)) - try: - annotations = get_type_hints(function, include_extras=True) - except Exception as error: - raise TypeError(f"failed to resolve Function annotations: {error}") from error missing = [ parameter.name for parameter in parameters if parameter.name not in annotations ] @@ -1132,15 +1162,21 @@ def _infer_signature( raise TypeError(f"missing Function annotations: {names!r}") inputs = [] for parameter in parameters: - data_type, nullable = _annotation_type(annotations[parameter.name]) - inputs.append( - FunctionParameter( - name=parameter.name, - arrow_type=_canonical_arrow_field( - pa.field(parameter.name, data_type, nullable=nullable) - ), - nullable=nullable, + annotation = annotations[parameter.name] + base = annotation + nullable = False + if get_origin(base) in (Union, types.UnionType): + choices = [value for value in get_args(base) if value is not type(None)] + if len(choices) == 1: + base, nullable = choices[0], True + if base is BlobReader: + inputs.append( + input_parameter(_blob_field(parameter.name, nullable=nullable)) ) + continue + data_type, nullable = _annotation_type(annotation) + inputs.append( + input_parameter(pa.field(parameter.name, data_type, nullable=nullable)) ) output_type, output_nullable = _annotation_type(annotations["return"]) if output_nullable: @@ -1239,7 +1275,7 @@ _PYTHON_CALLABLE_ARTIFACT = "python_callable" _PYTHON_BUNDLE_ARTIFACT = "python_bundle" # Modules the Function image places beside user code. _RESERVED_SOURCE_MODULES = frozenset( - {"user_function", "lance_source", "lance_callable"} + {"user_function", "lance_source", "lance_callable", "lance_blob"} ) # The service accepts bundle paths made of these components only. _SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") @@ -1499,6 +1535,8 @@ class _SourcePackager: """One module-level line that rebinds `name` to `value` in the artifact: an import for modules and importable classes/functions, a literal otherwise.""" + if value is BlobReader: + return f"from lance_blob import BlobReader as {name}" if isinstance(value, types.ModuleType): if value.__name__.split(".")[0] in _NAMESPACE_MODULES: raise ValueError( diff --git a/python/python/tests/test_blob_reader.py b/python/python/tests/test_blob_reader.py new file mode 100644 index 000000000..dd12409af --- /dev/null +++ b/python/python/tests/test_blob_reader.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import io + +import pytest + +from lancedb import BlobReader + + +def test_reader_is_lazy_and_seekable(): + requests = [] + size = 64 * 1024 * 1024 + + def read_range(offset, length): + requests.append((offset, length)) + return bytes((index % 251 for index in range(offset, offset + length))) + + with BlobReader(read_range, size) as reader: + assert requests == [] + assert reader.size() == size + reader.seek(-4, io.SEEK_END) + assert reader.read(10) == bytes(index % 251 for index in range(size - 4, size)) + assert reader.read(1) == b"" + assert requests == [(size - 4, 4)] + assert reader.read_range(1, 3) == b"\x01\x02\x03" + assert reader.tell() == size + reader.seek(2) + target = bytearray(3) + assert reader.readinto(target) == 3 + assert target == b"\x02\x03\x04" + reader.seek(-1, io.SEEK_CUR) + assert reader.tell() == 4 + with pytest.raises(ValueError, match="closed"): + reader.read(1) + + +def test_reader_supports_buffered_io_and_bounds_range_requests(): + requests = [] + payload = b"x" * (3 * 1024 * 1024 + 7) + + def read_range(offset, length): + requests.append((offset, length)) + return payload[offset : offset + length] + + with io.BufferedReader(BlobReader(read_range, len(payload))) as reader: + assert reader.read() == payload + assert requests + assert all(length <= 1024 * 1024 for _, length in requests) + + +def test_reader_rejects_invalid_ranges_and_short_reads(): + reader = BlobReader(lambda offset, length: b"", 10) + with pytest.raises(ValueError): + reader.seek(-1) + with pytest.raises(ValueError): + reader.read_range(-1, 1) + with pytest.raises(OSError, match="payload length"): + reader.read(1) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 9ee3e388a..874034564 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -23,6 +23,7 @@ import pyarrow as pa import pytest import lancedb +from lancedb import BlobReader from lancedb.functions import ( PythonRuntimeSpec, SecretBinding, @@ -966,6 +967,60 @@ def test_blob_fields_use_the_scalar_function_semantic_type(): assert signature.output.arrow_type == "blob_v2" +def test_blob_input_access_is_inferred_from_annotations(): + @udf + def read_blob(video: BlobReader, timestamp_s: float) -> bytes: + video.seek(int(timestamp_s)) + return video.read(4) + + @udf + def read_bytes(video: bytes) -> int: + return len(video) + + reader = read_blob.registration_request.signature + assert reader.inputs[0].arrow_type == "blob_v2" + assert reader.inputs[0].blob_access == "reader" + assert reader.inputs[1].blob_access is None + assert reader.output.arrow_type == "binary" + eager = read_bytes.registration_request.signature.inputs[0] + assert eager.arrow_type == "binary" + assert eager.blob_access == "bytes" + encoded = json.loads(read_blob.registration_request.to_canonical_json()) + assert encoded["signature"]["inputs"][0]["blob_access"] == "reader" + + +def test_nullable_blob_reader_and_explicit_blob_bytes(): + @udf + def nullable(video: Optional[BlobReader]) -> int: + return 0 if video is None else video.size() + + @udf( + input_schema=pa.schema([lancedb.blob("video", nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def eager(video: bytes) -> int: + return len(video) + + parameter = nullable.registration_request.signature.inputs[0] + assert parameter.nullable is True + assert parameter.blob_access == "reader" + parameter = eager.registration_request.signature.inputs[0] + assert parameter.arrow_type == "blob_v2" + assert parameter.blob_access == "bytes" + + +def test_blob_reader_rejects_an_incompatible_explicit_schema(): + def read(video: BlobReader) -> int: + return video.size() + + with pytest.raises(TypeError, match="requires a Blob V2 field"): + udf( + read, + input_schema=pa.schema([pa.field("video", pa.int64(), nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + + def test_whole_named_struct_function_can_include_a_blob_result_field(): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=False)]), diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 98d915118..36e6c30c8 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -156,11 +156,20 @@ pub struct FunctionArtifact { } /// One ordered Arrow input parameter. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BlobAccess { + Bytes, + Reader, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionParameter { pub name: String, pub arrow_type: String, pub nullable: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub blob_access: Option, } /// One field of an ordered named-struct result. @@ -956,6 +965,20 @@ mod conda_environment_tests { mod tests { use super::*; + #[test] + fn blob_reader_access_survives_signature_round_trip() { + let parameter: FunctionParameter = serde_json::from_value(serde_json::json!({ + "name": "video", "arrow_type": "blob_v2", "nullable": false, + "blob_access": "reader" + })) + .unwrap(); + assert_eq!(parameter.blob_access, Some(BlobAccess::Reader)); + assert_eq!( + serde_json::to_value(parameter).unwrap()["blob_access"], + "reader" + ); + } + /// Canonical form is what the FunctionVersion hash is taken over, so key /// order must come from the keys and not from however serde happened to /// emit them. Nesting is included because the sort is recursive. From 24c2669a50d31ad14fb763bbdd1e236a5087b0b7 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 25 Oct 2026 00:47:29 -0700 Subject: [PATCH 23/36] feat(sql): follow adaptive streaming and polling responses (#4664) SQL execution now starts with one Flight call. The server can stream rows immediately or return a polling continuation, which the SDK follows through the existing query handle and reader APIs. The same behavior applies with and without parameters. Both initial transports share direct-stream lifecycle handling. Handoffs are explicitly negotiated; existing servers can continue returning ordinary streams. Add coverage for streaming, deferred results, cancellation, malformed decisions, and legacy responses. ```mermaid flowchart TD A[execute_query or execute_query_async] --> B{Parameters supplied?} B -->|No| C[DoGet: SQL command ticket] B -->|Yes| D[DoExchange: SQL and Arrow parameter row] C --> E[Server chooses execution path] D --> E E -->|Stream| F[Initial response: streaming decision and Arrow batches] E -->|Defer| G[Initial response: PollInfo continuation] G --> H[SDK follows PollFlightInfo] H -->|Still running| H H -->|Result endpoints available| I[SDK fetches tickets with DoGet] F --> J[Query reader yields Arrow batches] I --> J ``` `execute_query` opens the reader; `execute_query_async` returns its query handle. The caller does not select the transport or execution path. --- docs/src/python/python.md | 13 +- docs/web/basics/sql.mdx | 2 + python/python/lancedb/db.py | 13 +- rust/lancedb/src/connection.rs | 10 +- rust/lancedb/src/remote/sql.rs | 86 ++++--- .../src/remote/sql/{exchange.rs => stream.rs} | 191 +++++++++++----- rust/lancedb/src/remote/sql_test.rs | 212 +++++++++++++++++- rust/lancedb/src/sql.rs | 20 +- 8 files changed, 414 insertions(+), 133 deletions(-) rename rust/lancedb/src/remote/sql/{exchange.rs => stream.rs} (69%) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 5d1393924..41d3e16d7 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -221,10 +221,19 @@ The connected database and `default_namespace_path=["public"]` are used for unqualified tables. Fully qualified references can still query other databases and namespaces available to the same deployment. `execute_query` returns a reader as soon as its initial result stream is available. `execute_query_async` -returns a query handle immediately; use it to inspect progress, open a reader, +returns a query handle once submission is accepted; use it to inspect progress, open a reader, or cancel the query. The SQL client is initialized by the first query and retained for the lifetime of the remote connection. Query ids are random, -connection-scoped references rather than encoded SQL or durable resume tokens: +connection-scoped references rather than encoded SQL or durable resume tokens. + +Both APIs start with one Flight call. The server can stream results on that +call or return a polling continuation that the SDK follows automatically. +This behavior is the same with and without parameters. Open the reader directly; +a directly streamed query needs its reader to make progress. Dropping an unread +handle or its reader stops direct streaming. Deferred execution continues until +explicitly cancelled. + +For example: ```python import lancedb diff --git a/docs/web/basics/sql.mdx b/docs/web/basics/sql.mdx index 1c4ad46cd..c89bd7a20 100644 --- a/docs/web/basics/sql.mdx +++ b/docs/web/basics/sql.mdx @@ -82,6 +82,8 @@ In application code, load credentials from environment variables or a secrets manager. The [connection guide](/basics/connect#enterprise) shows that setup. `execute_query()` returns a blocking Arrow reader when results start arriving. +The server can stream rows on the first call or return result tickets that the +SDK follows automatically, with the same behavior for parameterized queries. The example uses `read_all()` to collect the small result in memory. For larger results, consume the reader in batches. To submit a query and check its status separately, use `execute_query_async()`. See [Queries](/basics/query) for query diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 74acc28ba..8e53eac04 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1253,12 +1253,13 @@ class DBConnection(EnforceOverrides): parameters=[np.array([0.1, 0.2, 0.3], dtype=np.float32), 10], ) - Parameters are supported in queries, not in DDL or DML. A - parameterized query runs on the call that returns its rows: - cancelling it, or dropping its handle before reading, stops it on the - server. Its status follows that stream as well: ``describe()`` - reports ``running`` until the reader has received every row, so open - ``reader()`` directly rather than waiting for ``finished``. + Parameters are supported in queries, not in DDL or DML. With or + without parameters, the server chooses whether to stream rows on the + initial call or return a polling continuation. The reader follows + either path transparently. Dropping an unread handle or its reader + stops a directly streamed query; deferred execution requires explicit + cancellation. Open ``reader()`` directly: a streamed query needs its + reader to make progress and reports ``running`` until the stream ends. Local connections do not support SQL. """ diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 482acf89a..26afef5e8 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -372,12 +372,10 @@ impl ExecuteQueryAsyncBuilder { /// Do not follow a named placeholder directly with `$`: `$a$` starts a /// dollar-quoted string. /// - /// A parameterized statement runs on the call that returns its rows rather - /// than detached from it: cancelling the query, or dropping its handle - /// before reading it, stops the statement on the server. Its status follows - /// that stream too -- it reports `Running` until the reader has received - /// every row -- so open the reader directly rather than waiting for - /// `Finished`. + /// With or without parameters, the server may stream rows on the initial + /// call or return a polling continuation. The reader follows either path + /// transparently. Open it directly instead of waiting for `Finished`: + /// directly streamed queries need their reader to make progress. /// /// # Example /// diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index a58a9afad..d8f04119f 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -34,9 +34,9 @@ use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; -mod exchange; +mod stream; -use exchange::{ExchangeQuery, ExchangeQueryHandle}; +use stream::{StreamingQuery, StreamingQueryHandle, Submission}; const DEFAULT_SQL_PORT: u16 = 10025; const DEFAULT_SQL_TLS_PORT: u16 = 10026; @@ -191,20 +191,39 @@ impl SqlClient { } pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result { - match request.parameters { - Some(parameters) => { - self.submit_with_parameters( - &request.query, - parameters, - &request.default_namespace_path, - ) - .await + let timeout = self.inner.overall_timeout()?; + with_overall_timeout(timeout, "SQL query submission", async { + validate_namespace_path(&request.default_namespace_path)?; + let query_id = Uuid::now_v7(); + let submission = StreamingQuery::open( + self.inner.clone(), + &request.query, + request.parameters, + &request.default_namespace_path, + ) + .await?; + match submission { + Submission::Stream(response) => { + let query = + Arc::new(StreamingQuery::new(query_id, self.inner.clone(), response)); + self.queries + .insert(query_id, RegisteredQuery::Streamed(query.clone())); + Ok(Query::new(Arc::new(StreamingQueryHandle::new(query)))) + } + Submission::Deferred(poll_info) => { + let query = Arc::new(RemoteQuery::new( + query_id, + self.inner.clone(), + request.default_namespace_path, + poll_info, + )?); + self.queries + .insert(query_id, RegisteredQuery::Polled(query.clone())); + Ok(Query::new(Arc::new(RemoteQueryHandle::new(query)))) + } } - None => { - self.submit(&request.query, &request.default_namespace_path) - .await - } - } + }) + .await } pub(super) async fn submit( @@ -271,35 +290,6 @@ impl SqlClient { }) } - /// Values cannot ride on a poll, so a statement that has them is sent on - /// its own exchange instead. - pub(super) async fn submit_with_parameters( - &self, - query: &str, - parameters: RecordBatch, - default_namespace_path: &[String], - ) -> Result { - let timeout = self.inner.overall_timeout()?; - with_overall_timeout(timeout, "SQL query submission", async { - validate_namespace_path(default_namespace_path)?; - let query_id = Uuid::now_v7(); - let query = Arc::new( - ExchangeQuery::open( - query_id, - self.inner.clone(), - query, - parameters, - default_namespace_path, - ) - .await?, - ); - self.queries - .insert(query_id, RegisteredQuery::Exchanged(query.clone())); - Ok(Query::new(Arc::new(ExchangeQueryHandle::new(query)))) - }) - .await - } - pub(super) async fn describe(&self, query_id: Uuid) -> Result { let query = self .queries @@ -586,14 +576,14 @@ impl SqlClientInner { #[derive(Clone)] enum RegisteredQuery { Polled(Arc), - Exchanged(Arc), + Streamed(Arc), } impl RegisteredQuery { fn touch(&self) { match self { Self::Polled(query) => query.touch(), - Self::Exchanged(query) => query.touch(), + Self::Streamed(query) => query.touch(), } } @@ -602,14 +592,14 @@ impl RegisteredQuery { fn registry_expired(&self) -> bool { match self { Self::Polled(query) => query.registry_expired(Arc::strong_count(query) == 1), - Self::Exchanged(query) => query.registry_expired(Arc::strong_count(query) == 1), + Self::Streamed(query) => query.registry_expired(Arc::strong_count(query) == 1), } } async fn describe(&self) -> Result { match self { Self::Polled(query) => query.describe().await, - Self::Exchanged(query) => query.describe(), + Self::Streamed(query) => query.describe(), } } } diff --git a/rust/lancedb/src/remote/sql/exchange.rs b/rust/lancedb/src/remote/sql/stream.rs similarity index 69% rename from rust/lancedb/src/remote/sql/exchange.rs rename to rust/lancedb/src/remote/sql/stream.rs index 523aaf5d7..ed85c6c92 100644 --- a/rust/lancedb/src/remote/sql/exchange.rs +++ b/rust/lancedb/src/remote/sql/stream.rs @@ -1,26 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! SQL statements sent with parameters. -//! -//! A statement without parameters is polled: `PollFlightInfo` starts it -//! detached from the connection and `DoGet` reads its endpoints. Neither call -//! carries Arrow data, so a statement's values would have to be spelled into -//! its text. A statement with parameters is sent on one `DoExchange` instead: -//! the command and the parameter schema, then one row, and the result rows -//! come back on the same call. The call is the statement's lifetime, so -//! cancelling it -- or dropping the handle before reading -- stops the -//! statement on the server. +//! SQL submitted on one call, with either streamed rows or a polling handoff. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::Instant; use arrow_array::RecordBatch; -use arrow_flight::FlightDescriptor; +use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::encode::FlightDataEncoderBuilder; -use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; +use arrow_flight::sql::{Any, CommandStatementQuery, ProstMessageExt}; +use arrow_flight::{FlightDescriptor, PollInfo, Ticket}; use arrow_schema::SchemaRef; +use futures::TryStreamExt; use prost::Message; use tokio::sync::{Notify, mpsc}; use uuid::Uuid; @@ -51,8 +44,8 @@ fn failure(error: &Error) -> Lifecycle { }) } -/// A statement running on its own `DoExchange` call. -pub(super) struct ExchangeQuery { +/// A statement streaming on its initial `DoGet` or `DoExchange` call. +pub(super) struct StreamingQuery { id: Uuid, client: Arc, /// The call's rows, until a reader takes them or a cancellation drops them. @@ -63,7 +56,7 @@ pub(super) struct ExchangeQuery { last_accessed: StdMutex, } -struct PreparedExchange { +struct PreparedStream { schema: SchemaRef, first: Option, rest: Option, @@ -76,7 +69,7 @@ struct PreparedExchange { /// this one did. A query cancelled meanwhile has its rows dropped instead, /// which ends the call. struct Borrowed<'a> { - query: &'a ExchangeQuery, + query: &'a StreamingQuery, rows: Option, } @@ -104,18 +97,37 @@ impl Drop for Borrowed<'_> { } } -impl ExchangeQuery { - /// Send the statement and its parameters, returning once the server has - /// planned it: a statement it refuses is refused here, as a polled one is - /// on its first poll. +pub(super) const PROTOCOL_HEADER: &str = "lancedb-sql-protocol"; +pub(super) const PROTOCOL_VERSION: &str = "adaptive-v1"; +const RESPONSE_TYPE_URL: &str = "type.googleapis.com/lancedb.sql.QueryResponse"; + +#[derive(Clone, PartialEq, Message)] +pub(super) struct QueryResponse { + #[prost(message, optional, tag = "1")] + pub poll_info: Option, +} + +#[derive(Clone, PartialEq, Message)] +struct CommandTicket { + #[prost(bytes = "vec", tag = "2")] + command: Vec, +} + +pub(super) enum Submission { + Stream(ResultEndpointStream), + Deferred(PollInfo), +} + +impl StreamingQuery { pub(super) async fn open( - id: Uuid, client: Arc, query: &str, - parameters: RecordBatch, + parameters: Option, default_namespace_path: &[String], - ) -> Result { - if parameters.num_rows() != 1 { + ) -> Result { + if let Some(parameters) = ¶meters + && parameters.num_rows() != 1 + { return Err(Error::InvalidInput { message: format!( "SQL query parameters must be exactly one row, got {}", @@ -130,39 +142,102 @@ impl ExchangeQuery { Some(DEFAULT_READ_TIMEOUT), )? .unwrap(); - let descriptor = FlightDescriptor::new_cmd( - CommandStatementQuery { - query: query.to_string(), - transaction_id: None, + let command = CommandStatementQuery { + query: query.to_string(), + transaction_id: None, + } + .as_any() + .encode_to_vec(); + let connection = client.connection(&request_id).await?; + let headers = client.headers(default_namespace_path, &request_id).await?; + let mut metadata = super::client_with_headers(connection.client.clone(), &headers)? + .metadata() + .clone(); + metadata.insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + let mut flight = connection.client.clone(); + let response = if let Some(parameters) = parameters { + let messages = FlightDataEncoderBuilder::new() + .with_flight_descriptor(Some(FlightDescriptor::new_cmd(command))) + .with_schema(parameters.schema()) + .build(futures::stream::iter([Ok(parameters)])) + .try_collect::>() + .await + .map_err(|error| sql_error(&request_id, error))?; + let mut request = tonic::Request::new(futures::stream::iter(messages)); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_exchange(request)).await + } else { + let mut request = + tonic::Request::new(Ticket::new(CommandTicket { command }.encode_to_vec())); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_get(request)).await + } + .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? + .map_err(|error| sql_error(&request_id, error))?; + let negotiated = response + .metadata() + .get(PROTOCOL_HEADER) + .is_some_and(|value| value == PROTOCOL_VERSION); + let mut response_stream = response.into_inner(); + if negotiated { + let first = tokio::time::timeout(read_timeout, response_stream.message()) + .await + .map_err(|_| sql_error(&request_id, "SQL query handoff timed out"))? + .map_err(|error| sql_error(&request_id, error))? + .ok_or_else(|| { + sql_error(&request_id, "SQL response omitted its execution decision") + })?; + if !first.data_header.is_empty() || !first.data_body.is_empty() { + return Err(sql_error( + &request_id, + "SQL response sent rows before its execution decision", + )); } - .as_any() - .encode_to_vec(), + let envelope = + Any::decode(first.app_metadata).map_err(|error| sql_error(&request_id, error))?; + if envelope.type_url != RESPONSE_TYPE_URL { + return Err(sql_error( + &request_id, + "SQL response returned an unknown execution decision", + )); + } + let decision = QueryResponse::decode(envelope.value) + .map_err(|error| sql_error(&request_id, error))?; + if let Some(info) = decision.poll_info { + if info.info.is_none() && info.flight_descriptor.is_none() { + return Err(sql_error( + &request_id, + "SQL handoff omitted result and continuation information", + )); + } + return Ok(Submission::Deferred(info)); + } + } + let stream = FlightRecordBatchStream::new_from_flight_data( + response_stream + .map_err(|status| arrow_flight::error::FlightError::Tonic(Box::new(status))), ); - // The descriptor rides on the schema message, and the row follows it. - let request = FlightDataEncoderBuilder::new() - .with_flight_descriptor(Some(descriptor)) - .with_schema(parameters.schema()) - .build(futures::stream::iter([Ok(parameters)])); - let mut flight = client - .client_with_headers(default_namespace_path, &request_id) - .await?; - let stream = tokio::time::timeout(read_timeout, flight.do_exchange(request)) - .await - .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? - .map_err(|err| sql_error(&request_id, err))?; - Ok(Self { + Ok(Submission::Stream(ResultEndpointStream { + stream, + request_id, + read_timeout, + })) + } + + pub(super) fn new( + id: Uuid, + client: Arc, + response: ResultEndpointStream, + ) -> Self { + Self { id, client, - response: StdMutex::new(Some(ResultEndpointStream { - stream, - request_id, - read_timeout, - })), + response: StdMutex::new(Some(response)), lifecycle: StdMutex::new(Lifecycle::Running), cancelled: Notify::new(), terminal_at: OnceLock::new(), last_accessed: StdMutex::new(Instant::now()), - }) + } } pub(super) fn touch(&self) { @@ -242,7 +317,7 @@ impl ExchangeQuery { } /// Take the rows and wait for the first batch, which carries the schema. - async fn prepare(&self) -> Result { + async fn prepare(&self) -> Result { let Some(rows) = self.response.lock().unwrap().take() else { // Taken by an earlier attempt that failed, or dropped by a cancel. return Err(match &*self.lifecycle.lock().unwrap() { @@ -275,7 +350,7 @@ impl ExchangeQuery { if first.is_none() { self.settle(Lifecycle::Finished); } - Ok(PreparedExchange { + Ok(PreparedStream { schema, rest: first.is_some().then_some(rows), first, @@ -284,7 +359,7 @@ impl ExchangeQuery { async fn run( self: Arc, - mut prepared: PreparedExchange, + mut prepared: PreparedStream, sender: mpsc::Sender>, ) -> Result<()> { if let Some(batch) = prepared.first.take() @@ -343,13 +418,13 @@ impl ExchangeQuery { } } -pub(super) struct ExchangeQueryHandle { - query: Arc, +pub(super) struct StreamingQueryHandle { + query: Arc, result_started: AtomicBool, } -impl ExchangeQueryHandle { - pub(super) fn new(query: Arc) -> Self { +impl StreamingQueryHandle { + pub(super) fn new(query: Arc) -> Self { Self { query, result_started: AtomicBool::new(false), @@ -357,7 +432,7 @@ impl ExchangeQueryHandle { } } -impl Drop for ExchangeQueryHandle { +impl Drop for StreamingQueryHandle { // Rows nobody took are rows nobody will read. Keeping the call open would // hold the statement running on the server, stalled on flow control, // until the registry forgot it. @@ -369,7 +444,7 @@ impl Drop for ExchangeQueryHandle { } #[async_trait::async_trait] -impl QueryHandle for ExchangeQueryHandle { +impl QueryHandle for StreamingQueryHandle { fn id(&self) -> Uuid { self.query.touch(); self.query.id diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 79ca339db..4837baab7 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -19,6 +19,7 @@ use futures::{StreamExt, TryStreamExt}; use tonic::metadata::MetadataMap; use tonic::{Request, Response, Status, Streaming}; +use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse}; use super::*; use crate::database::Database; use crate::remote::RemoteCatalogOptions; @@ -70,6 +71,9 @@ type ReceivedExchanges = Arc)> #[derive(Clone)] struct TestSqlService { + adaptive: bool, + defer: bool, + invalid_decision: bool, query_count: Arc, do_get_count: Arc, cancel_count: Arc, @@ -120,6 +124,51 @@ impl TestSqlService { .unwrap_or_default(), }); } + fn adaptive_response( + &self, + query: &str, + rows: BoxStream<'static, std::result::Result>, + ) -> Response>> { + if !self.adaptive { + return Response::new(rows); + } + let poll_info = + self.defer.then(|| PollInfo { + info: Some(FlightInfo::new().with_endpoint( + FlightEndpoint::new().with_ticket(Ticket::new(query.to_string())), + )), + flight_descriptor: Some(FlightDescriptor::new_cmd(format!("poll:1:{query}"))), + progress: Some(0.25), + expiration_time: None, + }); + let first = FlightData { + app_metadata: if self.invalid_decision { + b"invalid protobuf".to_vec().into() + } else { + Any { + type_url: "type.googleapis.com/lancedb.sql.QueryResponse".to_string(), + value: QueryResponse { poll_info }.encode_to_vec().into(), + } + .encode_to_vec() + .into() + }, + ..Default::default() + }; + let rest = if self.defer { + futures::stream::empty().boxed() + } else { + rows + }; + let mut response = Response::new( + futures::stream::once(async { Ok(first) }) + .chain(rest) + .boxed(), + ); + response + .metadata_mut() + .insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + response + } } impl Default for TestSqlService { @@ -154,6 +203,9 @@ impl Default for TestSqlService { let dictionary_result = RecordBatch::try_new(dictionary_schema, vec![Arc::new(dictionary)]).unwrap(); Self { + adaptive: false, + defer: false, + invalid_decision: false, query_count: Arc::new(AtomicUsize::new(0)), do_get_count: Arc::new(AtomicUsize::new(0)), cancel_count: Arc::new(AtomicUsize::new(0)), @@ -308,7 +360,22 @@ impl FlightService for TestSqlService { request: Request, ) -> std::result::Result::DoGetStream>, Status> { self.do_get_count.fetch_add(1, Ordering::SeqCst); - let ticket = request.get_ref().ticket.as_ref(); + let initial = Any::decode(request.get_ref().ticket.clone()) + .ok() + .and_then(|envelope| Any::decode(envelope.value).ok()) + .and_then(|command| command.unpack::().ok().flatten()); + if initial.is_some() { + self.capture_headers(request.metadata()); + assert_eq!( + request.metadata().get(PROTOCOL_HEADER).unwrap(), + PROTOCOL_VERSION + ); + } + let query = initial + .as_ref() + .map(|command| command.query.as_bytes()) + .unwrap_or(request.get_ref().ticket.as_ref()); + let ticket = query; let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; @@ -337,7 +404,11 @@ impl FlightService for TestSqlService { encoder = encoder.with_max_flight_data_size(8 * 1024 * 1024); } let stream = encoder.build(input).map_err(Status::from); - Ok(Response::new(Box::pin(stream))) + if let Some(initial) = initial { + Ok(self.adaptive_response(&initial.query, stream.boxed())) + } else { + Ok(Response::new(stream.boxed())) + } } async fn do_put( @@ -486,7 +557,7 @@ impl FlightService for TestSqlService { let _dropped = &dropped; message }); - Ok(Response::new(Box::pin(stream))) + Ok(self.adaptive_response(&query, stream.boxed())) } } @@ -1609,3 +1680,138 @@ fn validates_metadata_with_header_map() { assert!(insert_header(&mut headers, "bad header", "value").is_err()); assert!(insert_header(&mut headers, "valid-header", "bad\nvalue").is_err()); } + +#[tokio::test] +async fn adaptive_streams_start_on_one_rpc_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = parameters.clone().unwrap_or_else(|| service.result.clone()); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT echo") + }) + .await + .unwrap(); + assert_eq!(query.describe().await.unwrap().status, QueryStatus::Running); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(gets.load(Ordering::SeqCst), usize::from(!with_parameters)); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_share_polling_and_survive_handle_drop() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = service.result.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT adaptive") + }) + .await + .unwrap(); + let id = query.id(); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + drop(query); + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 0, + "the original SQL must not be resubmitted through polling" + ); + assert_eq!( + gets.load(Ordering::SeqCst), + 1 + usize::from(!with_parameters) + ); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_use_server_cancellation_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let cancels = service.cancel_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT cancel") + }) + .await + .unwrap(); + query.cancel().await.unwrap(); + assert_eq!(cancels.load(Ordering::SeqCst), 1); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Cancelled + ); + assert!(query.reader().await.is_err()); + } +} + +#[tokio::test] +async fn malformed_adaptive_decision_never_resubmits_the_statement() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + invalid_decision: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + assert!( + client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT invalid") + }) + .await + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + } +} diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql.rs index c2ef13ba5..fee2d5d79 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql.rs @@ -18,9 +18,9 @@ pub enum QueryStatus { Running, /// The server has made the complete result available. /// - /// A query submitted with parameters streams its rows on the call that - /// runs it, so it reports `Finished` only once its reader has received - /// the last batch. Read it rather than waiting for this status. + /// Directly streamed queries report `Finished` when their stream reaches + /// its end. Deferred queries report the server's execution state. Open the + /// reader directly instead of waiting for this status before reading. Finished, /// The server accepted cancellation but has not confirmed it yet. Cancelling, @@ -64,9 +64,10 @@ pub(crate) trait QueryHandle: Send + Sync { /// A handle to a submitted SQL query. /// /// The handle can be inspected, opened as an Arrow reader, or cancelled. -/// Dropping it does not cancel the server-side query, unless the query was -/// submitted with parameters: those run on the call that returns their rows, -/// so dropping the handle before reading, or dropping the reader, stops them. +/// The server may stream results on the initial call or defer execution and +/// return result tickets. Dropping an unread handle or its reader stops a +/// directly streamed query; deferred execution continues until explicitly +/// cancelled. This applies equally with and without parameters. /// Identifier lookup is scoped to the connection that submitted the query and /// is not a durable resume mechanism. pub struct Query { @@ -95,10 +96,9 @@ impl Query { /// Get a point-in-time description of the query. /// - /// For a query submitted with parameters the status follows its result - /// stream: it stays [`QueryStatus::Running`] until the reader has received - /// every row, so open [`Self::reader`] directly instead of polling for - /// [`QueryStatus::Finished`] first. + /// Directly streamed queries stay [`QueryStatus::Running`] until the result + /// stream ends. Deferred queries expose the server's execution state. Open + /// [`Self::reader`] directly instead of waiting for [`QueryStatus::Finished`]. pub async fn describe(&self) -> Result { self.handle.describe().await } From 9af77570532c57029a47c7d47481b5cc12a9e2fe Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 25 Oct 2026 01:20:13 -0700 Subject: [PATCH 24/36] fix: execute remote materialized view creation through SQL (#4665) Remote materialized-view creation currently expects a REST response containing a registry job ID, while refresh already submits SQL. Route both through the SQL query backend and preserve Python's synchronous and asynchronous return types. ### Query to Job mapping `SqlClient::submit_as_job()` submits an `ExecuteQueryRequest` through the same backend as `execute_query_async()`. Once submission succeeds, it retains the returned `Query` and starts one client-side task that consumes its result reader to completion. `Job::from_query_task()` wraps that task, its cached outcome, a single ID and a cancellation callback. Python exposes this as `Job[T]` or `AsyncJob[T]`. The existing `job_type` property maps SQL-backed handles to the new `"query"` kind. The shared `id` contains that query's UUID string, so Python callers can inspect it with `db.describe_query(job.id)` on the submitting connection. No separate `query_id`, backend enum or `type()` API is added. | Handle | `job_type` | `id` | Available lookup | | --- | --- | --- | --- | | SQL-backed | `"query"`, known at submission | Connection-scoped query UUID string | `describe_query(job.id)` | | Registry-backed | Existing server-reported kind, after `status()` or `refresh()` | Registry job ID string | Registry job APIs, descriptions and events | | Local | `None` | `None` | The current in-process handle | | Job concept | SQL-backed behavior | | --- | --- | | Submission | Await initial SQL submission before returning the handle. Submission errors propagate immediately. | | `status()` | Tracks the adapter task: running during result consumption/finalization, finished after both succeed, failed if either fails, or cancelled after cancellation. | | `wait()` | Returns the operation's typed terminal result. One task consumes the query reader once; repeated waits reuse its cached outcome without resubmitting SQL or reopening the reader. | | `cancel()` | Delegates to `Query.cancel()`. Rejected or unsupported cancellation leaves the task running; an accepted request stops the task. Cancelling a finished operation is a no-op. | | Registry details | Query-backed handles expose no registry description, event history, pause/resume or reattachment after a client restart. `refresh()` updates cached lifecycle state. | The Job tracks command completion rather than exposing SQL batches. A wait timeout does not cancel the query. Query IDs remain scoped to the submitting connection and cannot be used with registry APIs or as durable resume tokens. Python `describe_query()` accepts both UUID objects and the string `id` exposed by query-backed jobs. Rust, Python and TypeScript expose the same classification through their existing job-type accessors. ### Materialized-view operations - **CREATE:** Render `CREATE MATERIALIZED VIEW` with the view namespace and optional auto-refresh settings. After command completion, resolve the job to `None`. Synchronous creation waits on the job, then opens and returns the materialized view. - **REFRESH:** Submit `REFRESH MATERIALIZED VIEW` through the same adapter. After command completion, read the refreshed table's version and row count to construct the existing typed refresh result. Synchronous refresh waits on the same path. Remote creation requires the SQL service. Execution placement remains a server decision. REST endpoints are unchanged. The job guide documents the ID/job-type contract and query-backed limitations. Regression coverage exercises SQL submission, namespaces, refresh options, job-type/ID preservation through typed result mapping, repeated waits, failures and cancellation rejection. No local builds or tests were run. --- docs/src/js/classes/Job.md | 13 +- docs/web/basics/jobs.mdx | 37 ++++- nodejs/__test__/materialized_view.test.ts | 1 + nodejs/lancedb/job.ts | 13 +- nodejs/src/job.rs | 12 +- python/python/lancedb/_lancedb.pyi | 2 + python/python/lancedb/db.py | 25 ++- python/python/lancedb/job.py | 30 ++-- python/python/lancedb/namespace.py | 2 +- python/python/lancedb/remote/db.py | 2 +- .../tests/test_first_class_function_slice2.py | 4 + .../python/tests/test_materialized_views.py | 87 +++++++--- python/python/tests/test_sql.py | 20 ++- python/src/job.rs | 5 + rust/lancedb/src/job.rs | 76 +++++++-- rust/lancedb/src/materialized_view.rs | 3 +- rust/lancedb/src/remote/db.rs | 144 ++++++----------- rust/lancedb/src/remote/sql.rs | 45 +++--- rust/lancedb/src/remote/sql_test.rs | 152 +++++++++++++++++- rust/lancedb/src/remote/table.rs | 8 +- 20 files changed, 465 insertions(+), 216 deletions(-) diff --git a/docs/src/js/classes/Job.md b/docs/src/js/classes/Job.md index cfbe454ec..74e1c00e2 100644 --- a/docs/src/js/classes/Job.md +++ b/docs/src/js/classes/Job.md @@ -17,7 +17,8 @@ extra round trip on every call: - [Job.refresh](Job.md#refresh) and [Job.status](Job.md#status) fetch the whole record. - [Job.wait](Job.md#wait) records the terminal state it establishes, but not the rest of the record. -- Everything is null until one of those runs. +- Registry details are null until one of those runs. Query handles report + jobType "query" immediately. ## Accessors @@ -57,8 +58,10 @@ get id(): null | string Identifies the operation on the server that is running it. -Operations that run in this process have no server id. The value is -opaque: parsing it or storing it to resume the job later is not supported. +Registry handles expose a registry job ID; query handles expose a +connection-scoped SQL query UUID string. Local handles have no ID. +Query handles report jobType "query". Their IDs are not registry IDs or +durable resume tokens. #### Returns @@ -72,8 +75,8 @@ opaque: parsing it or storing it to resume the job later is not supported. get jobType(): null | string ``` -The job's type, as the server names it. Null for an in-process job, which -has no server-side record. +"query" immediately for SQL-backed handles; otherwise the cached +server-reported job type. Null for in-process handles. #### Returns diff --git a/docs/web/basics/jobs.mdx b/docs/web/basics/jobs.mdx index 0577cd120..a87663a51 100644 --- a/docs/web/basics/jobs.mdx +++ b/docs/web/basics/jobs.mdx @@ -20,19 +20,20 @@ failures. Methods such as Python's `create_index_async` and TypeScript's `createIndexAsync` return job handles. In OSS, the work runs in your application -process; in Enterprise, it runs on server workers. A returned handle can -represent an operation that has already finished, so check its status before +process; in Enterprise, it runs on server workers or through the SQL query +backend. A returned handle can represent an operation that has already finished, +so check its status before waiting or requesting cancellation. See the operation's API reference for when the method returns. ## Submit and wait for a job {#waiting} 1. **Submit the operation.** Keep the returned job handle. For an Enterprise - job, also keep its ID if another process needs to look it up later. + registry job, also keep its ID if another process needs to look it up later. 2. **Check progress.** Read the current status while the operation runs. 3. **Wait for completion.** Check that the job succeeded before using its - result. If it failed or was cancelled, inspect the failure details and any - results it committed before stopping. + result. If it failed or was cancelled, inspect the failure details and any + results it committed before stopping. In Python, `job.status()` reports the current state and `job.wait()` waits for completion, raising if the job fails or is cancelled. If `job.wait()` times out, the client stops waiting. @@ -59,13 +60,37 @@ history. ## OSS and Enterprise execution {#in-process} -| | OSS | Enterprise | +The following table describes local operations and Enterprise registry jobs. +SQL-backed job handles follow the query lifecycle described below. + +| | OSS | Enterprise registry job | | --- | --- | --- | | Execution | In your application process | On server workers | | Job identity | No server ID | A server ID for looking up the job | | Event history | No server event history | Available through the jobs API | | Reattaching after a client restart | No server job to reopen | Reopen by ID while the record is retained | +### SQL-backed job handles + +Remote materialized-view CREATE and REFRESH use the SQL query backend. Python's +`create_materialized_view_async()` and `mv.refresh_async()` still return job +handles with `status()`, `wait()` and `cancel()`. Synchronous creation and refresh +wait on the same SQL execution before returning. + +The existing `job.job_type` identifies query-backed handles as `"query"` from +submission, and `job.id` contains the connection-scoped SQL query UUID string. +Use `describe_query(job.id)` on the submitting connection to inspect the query. +Registry handles retain their server-reported job types (available after +`status()` or `refresh()`) and registry IDs. Local handles have no ID or job type. + +SQL-backed handles do not expose registry records or event history and cannot +be reopened through `open_job()` or reattached after a client restart. Their IDs +are not durable resume tokens. Cancellation follows the SQL server's +capabilities and can be unsupported. + +The server decides whether SQL executes on its serving fleet or starts SQL job +workers. That placement does not change the handle's query-backed identity. + ## Learn more {#next} - [Build and repair indexes](/indexing/reindexing): submit index work. diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index 3e47a00e9..3852af717 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -129,6 +129,7 @@ describe("materialized views", () => { const job = await db.dropMaterializedViewAsync("adults"); expect(job.id).toBeNull(); + expect(job.jobType).toBeNull(); await job.wait(); expect(await collect(db.listMaterializedViews())).toEqual([]); }); diff --git a/nodejs/lancedb/job.ts b/nodejs/lancedb/job.ts index 0baa0f571..1c8d74ad2 100644 --- a/nodejs/lancedb/job.ts +++ b/nodejs/lancedb/job.ts @@ -24,7 +24,8 @@ export interface JobEventsOptions { * - {@link Job.refresh} and {@link Job.status} fetch the whole record. * - {@link Job.wait} records the terminal state it establishes, but not the * rest of the record. - * - Everything is null until one of those runs. + * - Registry details are null until one of those runs. Query handles report + * jobType "query" immediately. * * @hideconstructor */ @@ -38,8 +39,10 @@ export class Job { /** * Identifies the operation on the server that is running it. * - * Operations that run in this process have no server id. The value is - * opaque: parsing it or storing it to resume the job later is not supported. + * Registry handles expose a registry job ID; query handles expose a + * connection-scoped SQL query UUID string. Local handles have no ID. + * Query handles report jobType "query". Their IDs are not registry IDs or + * durable resume tokens. */ get id(): string | null { return this.inner.id ?? null; @@ -51,8 +54,8 @@ export class Job { } /** - * The job's type, as the server names it. Null for an in-process job, which - * has no server-side record. + * "query" immediately for SQL-backed handles; otherwise the cached + * server-reported job type. Null for in-process handles. */ get jobType(): string | null { return this.inner.jobType ?? null; diff --git a/nodejs/src/job.rs b/nodejs/src/job.rs index 214e687b4..320812128 100644 --- a/nodejs/src/job.rs +++ b/nodejs/src/job.rs @@ -29,8 +29,8 @@ impl Job { #[napi] impl Job { - /// Identifies the operation on the server that is running it. Operations - /// that run in this process have no server id. The value is opaque. + /// The operation's opaque ID. Interpret registry and query IDs using job_type(). + /// Local handles have no ID. #[napi(getter)] pub fn id(&self) -> Option { self.inner.id().map(str::to_string) @@ -62,8 +62,8 @@ impl Job { /// Ask the backend for this job's current state, and for a server-side job /// its full record, then cache it for the getters below. /// - /// They are all null until this runs, because submitting an operation - /// returns only a job id. {@link Job.status} fetches the whole record too; + /// Registry details are null until this runs; query handles already expose + /// job_type "query". {@link Job.status} fetches the whole registry record too; /// {@link Job.wait} records only the terminal state it establishes. #[napi(catch_unwind)] pub async fn refresh(&self) -> napi::Result<()> { @@ -76,8 +76,8 @@ impl Job { self.inner.state() } - /// The job's type, as the server names it. Null for an in-process job, - /// which has no server-side record. + /// "query" immediately for SQL-backed handles; otherwise the cached + /// server-reported job type. Null for in-process handles. #[napi(getter)] pub fn job_type(&self) -> Option { self.inner.job_type() diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 6ef23abc7..9e602cce8 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -320,6 +320,8 @@ class Job: @property def id(self) -> Optional[str]: ... @property + def job_type(self) -> Optional[str]: ... + @property def _state(self) -> Optional[str]: ... @property def _description(self) -> Optional[JobDescription]: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 8e53eac04..1558dd52a 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -674,10 +674,11 @@ class DBConnection(EnforceOverrides): ) -> Job[None]: """Submit materialized-view creation and return its job. - The job may already be complete for a local database. On LanceDB - Cloud and Enterprise, its ``id`` is the server job identifier from - the ``202 Accepted`` create response. Wait for the job before opening - or querying the view. + The job may already be complete for a local database. Remote creation + submits CREATE MATERIALIZED VIEW through the SQL server. Its job has + ``job_type == "query"`` and ``id`` identifies the connection-scoped + query. Registry event history is unavailable. + Wait for the job before opening or querying the view. """ raise NotImplementedError( "materialized views are not supported on this connection type" @@ -1265,7 +1266,7 @@ class DBConnection(EnforceOverrides): """ raise NotImplementedError("SQL is not supported for this connection type") - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" raise NotImplementedError("SQL is not supported for this connection type") @@ -2869,7 +2870,9 @@ class AsyncConnection(object): ) -> AsyncJob[None]: """Submit materialized-view creation and return its job. - Wait for the returned job before opening or querying the view. + Wait for the returned job before opening or querying the view. Remote + creation uses SQL; ``job_type`` is ``"query"`` and ``id`` + identifies that query on this connection. """ source, native_options = await self._materialized_view_source( source, select, where, limit @@ -3409,8 +3412,14 @@ class AsyncConnection(object): ) ) - async def describe_query(self, query_id: UUID) -> QueryDescription: - """Describe a submitted SQL query by its connection-scoped id.""" + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: + """Describe a submitted SQL query by its UUID or string ID. + + For a handle with ``job_type == "query"``, pass its ``id`` on the + connection that submitted it. Registry job IDs are not query IDs. + """ + if isinstance(query_id, str): + query_id = UUID(query_id) return await self._inner.describe_query(query_id) async def namespace_client(self) -> LanceNamespace: diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index e57fde0ea..e61a188f2 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -"""Handles to operations a server may run asynchronously.""" +"""Handles to asynchronous registry jobs, SQL queries, and local operations.""" import asyncio import json @@ -44,12 +44,12 @@ class AsyncJob(Generic[T]): @property def id(self) -> Optional[str]: - """Identifies the operation on the server that is running it. + """The operation's opaque ID, interpreted according to :attr:`job_type`. - Returned for correlating with server logs or the jobs API. Operations - that run in this process have no server id and return `None`. The value - is opaque: parsing it or storing it to resume the job later is not - supported. + Registry jobs expose a registry job ID; SQL-backed jobs expose a + connection-scoped query ID as a string. Local operations return ``None``. + Query IDs cannot be used with registry job APIs or to resume after a + client restart. """ return self._inner.id if self._inner is not None else None @@ -94,9 +94,11 @@ class AsyncJob(Generic[T]): """Ask the backend for this job's current state, and for a server-side job its full record, then cache it for the properties below. - The properties are all `None` until this runs, because submitting an - operation returns only a job id. `status` fetches the whole record too; - `wait` records only the terminal state it establishes. + Registry details are `None` until this runs, because submitting an + operation does not fetch its full registry record. Query-backed jobs + expose only lifecycle state through this method. For registry jobs, + `status` fetches the whole record too; `wait` records only the terminal + state it establishes. """ if self._inner is None: return @@ -114,11 +116,13 @@ class AsyncJob(Generic[T]): @property def job_type(self) -> Optional[str]: - """The job's type, as the server names it. + """The job's kind of work, without contacting the backend. - `None` for an in-process job, which has no server-side record. + SQL-backed handles report ``"query"`` immediately. Registry handles + report the server's job type after :meth:`refresh` or :meth:`status`. + In-process handles return ``None``. """ - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: @@ -294,7 +298,7 @@ class Job(Generic[T]): @property def job_type(self) -> Optional[str]: """The job's type. See :attr:`AsyncJob.job_type`.""" - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 55baa39a3..992926769 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -1683,7 +1683,7 @@ class AsyncLanceNamespaceDBConnection: parameters=parameters, ) - async def describe_query(self, query_id: UUID) -> QueryDescription: + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query when supported.""" return await self._inner.describe_query(query_id) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 5e9288471..62d00f284 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -1134,7 +1134,7 @@ class RemoteDBConnection(DBConnection): ) @override - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" return LOOP.run( self._conn.describe_query( diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 874034564..61e511345 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1729,7 +1729,11 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip(): ) registration = db.create_function_async(normalize_score) assert registration.id == "job-register" + assert registration.job_type is None + assert registration.status() == "finished" + assert registration.job_type == "create_function" created = registration.wait() + assert registration.job_type == "create_function" reopened = db.get_function("normalize_score", version=created.version) assert created == reopened diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index 46972a15c..e2be9c0d4 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -8,6 +8,7 @@ import threading import lancedb import pytest +from lancedb.db import AsyncConnection from lancedb.materialized_view import MaterializedViewDefinition from lancedb.remote.db import RemoteDBConnection @@ -56,7 +57,7 @@ def mock_remote_materialized_views(): @contextlib.contextmanager -def mock_remote_materialized_view_create(): +def mock_remote_materialized_view_drop(): requests = [] class Handler(http.server.BaseHTTPRequestHandler): @@ -67,7 +68,8 @@ def mock_remote_materialized_view_create(): length = int(self.headers.get("Content-Length", 0)) body = json.loads(self.rfile.read(length) or b"{}") requests.append((self.path, body)) - job_id = "mv-drop-123" if self.path.endswith("/drop") else "mv-create-123" + assert self.path.endswith("/drop") + job_id = "mv-drop-123" encoded = json.dumps({"job_id": job_id}).encode() self.send_response(202) self.send_header("Content-Type", "application/json") @@ -98,7 +100,7 @@ def test_remote_list_uses_namespace_route(): def test_remote_drop_async_returns_server_job(): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect( "db://dev", api_key="fake", @@ -192,11 +194,13 @@ def test_create_and_refresh_jobs(tmp_path): "adults", "people", where="age >= 18", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert create_job.wait() is None view = db.open_materialized_view("adults") refresh_job = view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 2 @@ -337,11 +341,13 @@ async def test_async_create_and_refresh_jobs(tmp_path): "adults", "people", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert await create_job.wait() is None view = await db.open_materialized_view("adults") refresh_job = await view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = await refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 1 @@ -567,6 +573,27 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): assert list(resumed) == [] +class QueryBackedCreateJob: + id = "0198f1b2-c3d4-7e5f-8123-456789abcdef" + + job_type = "query" + + async def wait(self): + return None + + async def status(self): + return "finished" + + +class CapturingCreateConnection: + def __init__(self): + self.calls = [] + + async def create_materialized_view_async(self, name, source, **options): + self.calls.append((name, source, options)) + return QueryBackedCreateJob() + + @pytest.mark.parametrize( "options", [ @@ -576,16 +603,28 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): ], ) def test_remote_create_forwards_refresh_options_without_interval_default(options): - with mock_remote_materialized_view_create() as (host, requests): - db = lancedb.connect("db://dev", api_key="fake", host_override=host) - job = db.create_materialized_view_async("adults", "people", **options) - assert job.id == "mv-create-123" - assert requests == [ + native = CapturingCreateConnection() + db = RemoteDBConnection.__new__(RemoteDBConnection) + db._conn = AsyncConnection(native) + job = db.create_materialized_view_async("adults", "people", **options) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert job.wait() is None + assert job.status() == "finished" + assert native.calls == [ ( - "/v1/materialized_view/adults/create", + "adults", + "people", { - "query": 'SELECT * FROM "people"', - **options, + "projections": None, + "filter": None, + "limit": None, + "with_no_data": False, + "auto_refresh_enabled": options.get("auto_refresh_enabled"), + "auto_refresh_interval_seconds": options.get( + "auto_refresh_interval_seconds" + ), }, ) ] @@ -593,18 +632,18 @@ def test_remote_create_forwards_refresh_options_without_interval_default(options @pytest.mark.asyncio async def test_async_remote_create_keeps_interval_unspecified(): - with mock_remote_materialized_view_create() as (host, requests): - db = await lancedb.connect_async("db://dev", api_key="fake", host_override=host) - job = await db.create_materialized_view_async( - "adults", - "people", - auto_refresh_enabled=True, - ) - assert job.id == "mv-create-123" - assert requests[0][1] == { - "query": 'SELECT * FROM "people"', - "auto_refresh_enabled": True, - } + native = CapturingCreateConnection() + db = AsyncConnection(native) + job = await db.create_materialized_view_async( + "adults", "people", auto_refresh_enabled=True + ) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert await job.wait() is None + assert await job.status() == "finished" + assert native.calls[0][2]["auto_refresh_enabled"] is True + assert native.calls[0][2]["auto_refresh_interval_seconds"] is None @pytest.mark.parametrize( @@ -625,7 +664,7 @@ def test_embedded_create_refuses_automatic_refresh_before_creating_view( @pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf")]) def test_remote_create_rejects_invalid_interval_without_request(seconds): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect("db://dev", api_key="fake", host_override=host) with pytest.raises(ValueError, match="positive, finite"): db.create_materialized_view_async( diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index e6f5fd3c7..539b88647 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -137,9 +137,23 @@ async def test_async_namespace_connection_rejects_sql(tmp_path): await connection.describe_query(NIL_QUERY_ID) -def test_describe_query_requires_uuid(): - with pytest.raises(TypeError, match="UUID"): - remote_connection().describe_query(str(NIL_QUERY_ID)) +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", [NIL_QUERY_ID, str(NIL_QUERY_ID)]) +async def test_describe_query_accepts_uuid_and_job_id_string(identity): + calls = [] + + class NativeConnection: + async def describe_query(self, query_id): + calls.append(query_id) + + connection = AsyncConnection(NativeConnection()) + await connection.describe_query(identity) + assert calls == [NIL_QUERY_ID] + + +def test_describe_query_rejects_registry_job_id(): + with pytest.raises(ValueError, match="UUID"): + remote_connection().describe_query("registry-job-id") @pytest.mark.parametrize( diff --git a/python/src/job.rs b/python/src/job.rs index e22b2f897..1740b3e00 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -81,6 +81,11 @@ impl Job { self.inner.id().map(str::to_string) } + #[getter] + pub fn job_type(&self) -> Option { + self.inner.job_type() + } + pub fn status(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); future_into_py( diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index d671365f9..5c27587dc 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -13,6 +13,7 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; use tokio::sync::watch; use tokio::task::{AbortHandle, JoinHandle}; +use uuid::Uuid; use crate::database::JobDescription; use crate::error::{Error, JobFailure, Result}; @@ -55,10 +56,13 @@ fn job_detail_not_supported(what: &str) -> Result { /// Backend-specific tracking for an asynchronous operation. #[async_trait] pub(crate) trait JobHandle: Send + Sync { - /// Server-assigned id, when the backend has one. + /// Backend-specific id, when the operation has one. fn id(&self) -> Option<&str> { None } + fn job_type(&self) -> Option<&str> { + None + } async fn status(&self) -> Result; async fn wait(&self) -> Result; async fn cancel(&self) -> Result<()>; @@ -167,14 +171,15 @@ struct JobCache { /// the endpoint's successful terminal result; unit-result operations use the /// default `Job<()>`. /// -/// The detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the +/// Registry detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the /// handle last observed. Submitting an operation returns only a job id, so /// populating them eagerly would cost an extra round trip on every call: /// /// - [`Job::refresh`] and [`Job::status`] fetch the whole record. /// - [`Job::wait`] records the terminal state it establishes, but not the rest /// of the record; call [`Job::refresh`] for that. -/// - Everything is `None` until one of those runs. +/// - Registry details are `None` until one of those runs. Query-backed handles +/// report [`Job::job_type`] as `"query"` immediately. pub struct Job where T: Clone + Send + Sync + 'static, @@ -273,13 +278,18 @@ where Self::new_typed(Box::new(SpawnedJob::new(task))) } - pub(crate) fn spawned_with_cancellation(task: JoinHandle>, cancel: F) -> Self + pub(crate) fn from_query_task( + query_id: Uuid, + task: JoinHandle>, + cancel: F, + ) -> Self where F: Fn() -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { let cancellation: Cancellation = Arc::new(move || Box::pin(cancel())); Self::new_typed(Box::new(SpawnedJob::new_with_cancellation( + query_id, task, cancellation, ))) @@ -290,12 +300,10 @@ impl Job where T: Clone + Send + Sync + 'static, { - /// Identifies the operation on the server that is running it. - /// - /// Returned for correlating with server logs or the jobs API. Operations - /// that run in this process have no server id and return `None`. The - /// value is opaque: parsing it or storing it to resume the job later is - /// not supported. + /// The operation's opaque identity, interpreted according to [`Job::job_type`]. + /// Registry handles expose a registry job ID; query handles expose the + /// connection-scoped SQL query UUID as a string. Local handles return None. + /// A query ID cannot be used with registry APIs or to resume after restart. pub fn id(&self) -> Option<&str> { match &self.inner { JobInner::Handle { handle, .. } => handle.id(), @@ -369,8 +377,14 @@ where self.cache_read().description.clone() } - /// The job's type, as the server names it. `None` for an in-process job. + /// The job's type: `"query"` for SQL-backed handles, otherwise the cached + /// server-reported kind of work. `None` for an in-process job. pub fn job_type(&self) -> Option { + if let JobInner::Handle { handle, .. } = &self.inner + && let Some(job_type) = handle.job_type() + { + return Some(job_type.to_string()); + } self.with_description(|description| description.job_type.clone()) } @@ -536,6 +550,7 @@ struct SpawnedJob { outcome: watch::Receiver>, abort: AbortHandle, cancellation: Option, + id: Option, } type CancellationFuture = Pin> + Send>>; @@ -549,11 +564,17 @@ impl SpawnedJob { Self::new_inner(task, None) } - fn new_with_cancellation(task: JoinHandle>, cancellation: Cancellation) -> Self + fn new_with_cancellation( + query_id: Uuid, + task: JoinHandle>, + cancellation: Cancellation, + ) -> Self where T: Serialize + Send + 'static, { - Self::new_inner(task, Some(cancellation)) + let mut handle = Self::new_inner(task, Some(cancellation)); + handle.id = Some(query_id.to_string()); + handle } fn new_inner(task: JoinHandle>, cancellation: Option) -> Self @@ -582,12 +603,21 @@ impl SpawnedJob { outcome, abort, cancellation, + id: None, } } } #[async_trait] impl JobHandle for SpawnedJob { + fn id(&self) -> Option<&str> { + self.id.as_deref() + } + + fn job_type(&self) -> Option<&str> { + self.id.as_ref().map(|_| "query") + } + async fn status(&self) -> Result { let label = match &*self.outcome.borrow() { None => "running", @@ -648,11 +678,31 @@ mod tests { async fn mapped_spawned_job_reuses_outcome() { let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1); + assert_eq!(job.job_type(), None); + assert_eq!(job.id(), None); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.status().await.unwrap(), "finished"); } + #[tokio::test] + async fn mapped_query_job_preserves_identity_and_type() { + let query_id = Uuid::new_v4(); + let job = Job::from_query_task(query_id, tokio::spawn(async { Ok(41_u64) }), || async { + Ok(()) + }) + .map(|value| value + 1); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + let completed = Job::new_done(); + assert_eq!(completed.job_type(), None); + assert_eq!(completed.id(), None); + } + #[tokio::test] async fn mapped_spawned_job_preserves_cancellation() { let job = Job::spawned(tokio::spawn(async { pending::>().await })) diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index ecdd1783b..7ee76a160 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -1716,7 +1716,8 @@ impl CreateMaterializedViewBuilder { } /// Submit creation and initial population, returning a [`Job`] that - /// settles when the view is ready. + /// settles when the view is ready. Remote creation submits SQL; the handle's + /// [`Job::job_type`] is `"query"` and its ID identifies that query. pub async fn execute_async(self) -> Result { if self.connection.uri().starts_with("db://") { return self diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 768de79d2..d33565653 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -63,15 +63,6 @@ struct RemoteCloneTableRequest { is_shallow: Option, } -#[derive(serde::Serialize)] -struct RemoteCreateMaterializedViewRequest { - query: String, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_enabled: Option, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_interval_seconds: Option, -} - // the versions of the server that we support // for any new feature that we need to change the SDK behavior, we should bump the server version, // and add a feature flag as method of `ServerVersion` here. @@ -594,6 +585,11 @@ mod test_utils { } } + pub(in crate::remote) fn with_sql_client(mut self, client: SqlClient) -> Self { + self.sql_client = Some(client); + self + } + pub fn new_mock_with_config(handler: F, config: ClientConfig) -> Self where F: Fn(reqwest::Request) -> http::Response + Send + Sync + 'static, @@ -885,8 +881,7 @@ impl Database for RemoteDatabase { ) -> Result { if request.with_no_data { return Err(Error::NotSupported { - message: "remote materialized views are always populated by their SQL job" - .to_string(), + message: "remote materialized views must be populated during creation".to_string(), }); } if let Some(seconds) = request.auto_refresh_interval_seconds @@ -896,36 +891,37 @@ impl Database for RemoteDatabase { message: "auto_refresh_interval_seconds must be positive, finite seconds".into(), }); } - let identifier = build_table_identifier(&request.name, &request.namespace_path)?; - let body = RemoteCreateMaterializedViewRequest { - query: request.query, - auto_refresh_enabled: request.auto_refresh_enabled, - auto_refresh_interval_seconds: request.auto_refresh_interval_seconds, - }; - let request = self - .client - .post(&format!("/v1/materialized_view/{identifier}/create")) - .json(&body); - let (request_id, response) = self.client.send(request).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - if status != StatusCode::ACCEPTED { - return Err(Error::Http { - source: "materialized-view creation must return 202 Accepted".into(), - request_id, - status_code: Some(status), - }); + validate_table_name(&request.name)?; + validate_namespace(&request.namespace_path)?; + let client = self + .sql_client + .as_ref() + .ok_or_else(|| Error::NotSupported { + message: "SQL is unavailable for this remote database client".into(), + })?; + let mut options = Vec::new(); + if let Some(enabled) = request.auto_refresh_enabled { + options.push(format!("auto_refresh_enabled = {enabled}")); } - let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { - source: "materialized-view create response did not contain a valid job_id".into(), - request_id, - status_code: Some(status), - })?; - Ok(Job::new(Box::new(RemoteJob::new( - self.client.clone(), - job_id, - )))) + if let Some(seconds) = request.auto_refresh_interval_seconds { + options.push(format!("auto_refresh_interval_seconds = {seconds}")); + } + let name = request.name.replace('"', "\"\""); + let options = if options.is_empty() { + String::new() + } else { + format!(" WITH ({})", options.join(", ")) + }; + client + .submit_as_job( + format!( + "CREATE MATERIALIZED VIEW \"{name}\"{options} AS {}", + request.query + ), + request.namespace_path, + || async { Ok(()) }, + ) + .await } async fn drop_materialized_view_async( @@ -2016,72 +2012,22 @@ mod tests { } #[tokio::test] - async fn test_create_materialized_view_uses_rest_and_preserves_optional_settings() { - for (enabled, interval) in [(None, None), (Some(true), None), (Some(false), Some(0.5))] { - let db = super::RemoteDatabase::new_mock(move |request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/materialized_view/analytics$adults/create" - ); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(body["query"], "SELECT * FROM people"); - assert_eq!( - body.get("auto_refresh_enabled"), - enabled.map(serde_json::Value::from).as_ref() - ); - assert_eq!( - body.get("auto_refresh_interval_seconds"), - interval.map(serde_json::Value::from).as_ref() - ); - http::Response::builder() - .status(202) - .body(serde_json::json!({"job_id": "j1-mv-create"}).to_string()) - .unwrap() - }); - let job = db - .create_materialized_view_async(CreateMaterializedViewRequest { - name: "adults".into(), - namespace_path: vec!["analytics".into()], - query: "SELECT * FROM people".into(), - with_no_data: false, - auto_refresh_enabled: enabled, - auto_refresh_interval_seconds: interval, - }) - .await - .unwrap(); - assert_eq!(job.id(), Some("j1-mv-create")); - } - } - - #[tokio::test] - async fn test_create_materialized_view_rejects_invalid_acceptance() { - for (status, body) in [ - (200, r#"{"job_id":"j1"}"#), - (202, "{}"), - (202, r#"{"job_id":""}"#), - ] { - let db = super::RemoteDatabase::new_mock(move |_| { - http::Response::builder().status(status).body(body).unwrap() - }); - let error = db + async fn test_create_materialized_view_validates_before_submission() { + let db = super::RemoteDatabase::new_mock(|_| -> http::Response { + panic!("materialized-view creation must not send REST requests") + }); + for seconds in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let result = db .create_materialized_view_async(CreateMaterializedViewRequest { name: "adults".into(), namespace_path: vec![], query: "SELECT * FROM people".into(), with_no_data: false, auto_refresh_enabled: None, - auto_refresh_interval_seconds: None, + auto_refresh_interval_seconds: Some(seconds), }) - .await - .err() - .unwrap(); - assert!(error.to_string().contains(if status == 200 { - "202 Accepted" - } else { - "job_id" - })); + .await; + assert!(matches!(result, Err(Error::InvalidInput { .. }))); } } diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index d8f04119f..846e852f5 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::fs; -use std::future::{Future, pending}; +use std::future::Future; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::{Duration, Instant}; @@ -12,6 +12,7 @@ use arrow_array::RecordBatch; use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::error::FlightError; use arrow_flight::flight_service_client::FlightServiceClient; +#[cfg(test)] use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; use arrow_flight::{ Action, CancelFlightInfoRequest, CancelFlightInfoResult, CancelStatus, FlightClient, @@ -22,7 +23,7 @@ use futures::TryStreamExt; use http::header::{HeaderMap, HeaderName, HeaderValue}; use prost::Message; use serde::{Serialize, de::DeserializeOwned}; -use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch}; +use tokio::sync::{Mutex, Notify, OnceCell, mpsc}; use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}; use uuid::Uuid; @@ -226,6 +227,7 @@ impl SqlClient { .await } + #[cfg(test)] pub(super) async fn submit( &self, query: &str, @@ -254,40 +256,36 @@ impl SqlClient { .await } - pub(super) fn submit_as_job( + pub(super) async fn submit_as_job( &self, statement: String, namespace: Vec, finish: F, - ) -> Job + ) -> Result> where T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, F: FnOnce() -> Fut + Send + 'static, Fut: Future> + Send + 'static, { - let client = self.clone(); - let (query_tx, query_rx) = watch::channel(None::>); + let query = Arc::new( + self.execute(ExecuteQueryRequest { + query: statement, + default_namespace_path: namespace, + parameters: None, + }) + .await?, + ); + let query_id = query.id(); + let execution = query.clone(); let task = tokio::spawn(async move { - let query = Arc::new(client.submit(&statement, &namespace).await?); - query_tx.send_replace(Some(query.clone())); - let mut reader = query.reader().await?; + let mut reader = execution.reader().await?; while reader.try_next().await?.is_some() {} finish().await }); - Job::spawned_with_cancellation(task, move || { - let mut query_rx = query_rx.clone(); - async move { - loop { - let query = query_rx.borrow().clone(); - if let Some(query) = query { - return query.cancel().await; - } - if query_rx.changed().await.is_err() { - pending::<()>().await; - } - } - } - }) + Ok(Job::from_query_task(query_id, task, move || { + let query = query.clone(); + async move { query.cancel().await } + })) } pub(super) async fn describe(&self, query_id: Uuid) -> Result { @@ -315,6 +313,7 @@ impl SqlClientInner { ) } + #[cfg(test)] async fn poll( &self, descriptor: FlightDescriptor, diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 4837baab7..1368dd753 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -22,6 +22,7 @@ use tonic::{Request, Response, Status, Streaming}; use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse}; use super::*; use crate::database::Database; +use crate::materialized_view::CreateMaterializedViewRequest; use crate::remote::RemoteCatalogOptions; use crate::remote::client::HeaderProvider; use crate::remote::client::test_utils::client_with_handler; @@ -75,6 +76,7 @@ struct TestSqlService { defer: bool, invalid_decision: bool, query_count: Arc, + statements: Arc>>, do_get_count: Arc, cancel_count: Arc, cancel_denied_count: Arc, @@ -207,6 +209,7 @@ impl Default for TestSqlService { defer: false, invalid_decision: false, query_count: Arc::new(AtomicUsize::new(0)), + statements: Arc::new(std::sync::Mutex::new(Vec::new())), do_get_count: Arc::new(AtomicUsize::new(0)), cancel_count: Arc::new(AtomicUsize::new(0)), cancel_denied_count: Arc::new(AtomicUsize::new(0)), @@ -269,6 +272,7 @@ impl FlightService for TestSqlService { .and_then(|any| any.unpack::().ok().flatten()); let (query, stage) = if let Some(command) = command { self.query_count.fetch_add(1, Ordering::SeqCst); + self.statements.lock().unwrap().push(command.query.clone()); (command.query, 0_u8) } else { let continuation = std::str::from_utf8(request.get_ref().cmd.as_ref()) @@ -289,6 +293,9 @@ impl FlightService for TestSqlService { .ok_or_else(|| Status::invalid_argument("invalid continuation"))?; (query.to_string(), stage) }; + if query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } if (query == "SELECT slow" || query == "SELECT cancelling") && stage > 0 { tokio::time::sleep(Duration::from_millis(250)).await; } @@ -364,7 +371,11 @@ impl FlightService for TestSqlService { .ok() .and_then(|envelope| Any::decode(envelope.value).ok()) .and_then(|command| command.unpack::().ok().flatten()); - if initial.is_some() { + if let Some(command) = &initial { + self.statements.lock().unwrap().push(command.query.clone()); + if command.query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } self.capture_headers(request.metadata()); assert_eq!( request.metadata().get(PROTOCOL_HEADER).unwrap(), @@ -376,7 +387,8 @@ impl FlightService for TestSqlService { .map(|command| command.query.as_bytes()) .unwrap_or(request.get_ref().ticket.as_ref()); let ticket = query; - let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; + let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"" + || ticket == b"SELECT job cancel denied"; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; let large = ticket == b"SELECT large message"; @@ -458,7 +470,9 @@ impl FlightService for TestSqlService { { return Err(Status::not_found("query cancellation completed")); } - let status = if query == "SELECT cancel unspecified" { + let status = if query == "SELECT job cancel denied" { + CancelStatus::NotCancellable + } else if query == "SELECT cancel unspecified" { CancelStatus::Unspecified } else if query == "SELECT cancelling" { if self @@ -880,6 +894,132 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { server.await.unwrap().unwrap(); } +#[tokio::test] +async fn remote_create_uses_sql_and_preserves_options_and_namespace() { + for defer in [false, true] { + let service = TestSqlService { + adaptive: true, + defer, + ..Default::default() + }; + let statements = service.statements.clone(); + let headers = service.headers.clone(); + let count = service.query_count.clone(); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + panic!("CREATE must use SQL, not REST: {}", request.url().path()); + }) + .with_sql_client(sql_client.clone()); + for (enabled, interval, expected) in [ + (None, None, ""), + (Some(true), None, " WITH (auto_refresh_enabled = true)"), + ( + Some(false), + Some(0.5), + " WITH (auto_refresh_enabled = false, auto_refresh_interval_seconds = 0.5)", + ), + ] { + let job = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["team".into(), "reports".into()], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: enabled, + auto_refresh_interval_seconds: interval, + }) + .await + .unwrap(); + assert_eq!(job.job_type().as_deref(), Some("query")); + let query_id = + Uuid::parse_str(job.id().expect("SQL-backed job must expose its ID")).unwrap(); + job.wait().await.unwrap(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + assert!(job.events(Default::default()).await.is_err()); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!( + sql_client.describe(query_id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + statements.lock().unwrap().last().unwrap(), + &format!("CREATE MATERIALIZED VIEW \"adults\"{expected} AS SELECT * FROM people") + ); + } + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(statements.lock().unwrap().len(), 3); + assert!( + headers + .lock() + .unwrap() + .iter() + .all(|header| header.namespace_path == "team$reports") + ); + shutdown.send(()).unwrap(); + } +} + +#[tokio::test] +async fn query_backed_jobs_propagate_submission_and_execution_failures() { + let (client, shutdown) = start_sql_service( + TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }, + ClientConfig::default(), + ) + .await; + let rejected = client + .submit_as_job("SELECT rejected".into(), vec![], || async { Ok(()) }) + .await; + assert!(rejected.is_err()); + let finished = Arc::new(AtomicBool::new(false)); + let callback_finished = finished.clone(); + let job = client + .submit_as_job("SELECT fail".into(), vec![], move || async move { + callback_finished.store(true, Ordering::SeqCst); + Ok(()) + }) + .await + .unwrap(); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert_eq!(job.status().await.unwrap(), "failed"); + assert!(!finished.load(Ordering::SeqCst)); + shutdown.send(()).unwrap(); +} + +#[tokio::test] +async fn rejected_job_cancellation_does_not_abort_sql_completion() { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let do_get_count = service.do_get_count.clone(); + let release = service.gate_refresh_release.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let job = client + .submit_as_job("SELECT job cancel denied".into(), vec![], || async { + Ok(()) + }) + .await + .unwrap(); + wait_for(&do_get_count, 1, "job did not begin reading its SQL result").await; + assert!(matches!( + job.cancel().await, + Err(Error::NotSupported { .. }) + )); + assert_eq!(job.status().await.unwrap(), "running"); + release.notify_one(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + job.cancel().await.unwrap(); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn remote_refresh_twice_keeps_latest_selector() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -969,7 +1109,11 @@ async fn cancelling_remote_refresh_cancels_sql_query() { let (socket, _) = listener.accept().await?; Ok::<_, std::io::Error>(Some((socket, listener))) }); - let service = TestSqlService::default(); + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; let cancel_count = service.cancel_count.clone(); let do_get_count = service.do_get_count.clone(); let gate_refresh_release = service.gate_refresh_release.clone(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index b722dc4a1..4cbba04df 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2247,8 +2247,8 @@ impl BaseTable for RemoteTable { let freshness = self.freshness.clone(); let freshness_request = self.snapshot_freshness_headers(); let job_sql_client = sql_client.clone(); - Ok( - sql_client.submit_as_job(statement, namespace.clone(), move || async move { + sql_client + .submit_as_job(statement, namespace.clone(), move || async move { let table = Self::new_with_sql_client( client, name, @@ -2270,8 +2270,8 @@ impl BaseTable for RemoteTable { source_version: 0, version, }) - }), - ) + }) + .await } async fn query_snapshot(&self) -> Result> { let description = self.describe().await?; From dcfef7c5ab259d2ce36003b6d6bf10c6de22cd9e Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 25 Oct 2026 02:31:21 -0700 Subject: [PATCH 25/36] feat(python): open lazy blob files from SQL descriptors (#4668) SQL blob descriptors need pinned source identity to reopen a value after the query finishes. Adds `open_blob_file(descriptor)` and `open_blob_files(descriptors)` so Python can read the original blob incrementally without materializing it in the SQL result. ### Python experience ```python rows = db.execute_query(""" SELECT f.timestamp, blob_descriptor(v.video) AS video FROM videos v JOIN frames f ON v.id = f.video_id """).read_all() videos = db.open_table("videos") for row in rows.to_pylist(): if row["video"] is not None: with videos.open_blob_file(row["video"]) as video: frame = extract_frame(video, row["timestamp"]) ``` The singular API accepts the descriptor dictionary or an Arrow struct scalar directly and returns a seekable `BlobFile`. `videos.open_blob_files(rows["video"])` accepts an Arrow array/chunked array or another iterable, opens handles as iteration advances, and preserves order, duplicates, and nulls as `None`. Async tables provide `await videos.open_blob_file(...)` and `async for video in videos.open_blob_files(...)`. The iterable API can be used directly with the Arrow result column. Zip it with the rows to keep each frame timestamp: ```python for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): if video is not None: with video: frame = extract_frame(video, row["timestamp"]) ``` Each iteration opens one handle. Null descriptors yield `None`; handles preserve input order and duplicates, and the `with` block closes each non-null file. For an async table: ```python async for video in async_videos.open_blob_files(rows["video"]): if video is not None: with video: payload = await video.aread() ``` ### Execution The descriptor retains the canonical `kind`, `position`, `size`, `blob_id`, and `blob_uri` fields and carries a `source` struct with its dataset URI, version, commit timestamp, branch, column, row reference, and whether that reference is a physical address. Native opening resolves that exact snapshot without changing the table handle's checkout and rejects a descriptor belonging to another table. Remote handles send authenticated byte-range requests with the captured version/branch, commit timestamp, and source URI, so an old physical address is never interpreted against the latest table. Known-size remote descriptors open without a payload request. External blobs with an unspecified size use a metadata-only HEAD request. Payload bytes load only when the file is read; full reads still allocate the full value. Existing row-ID-based `fetch_blob_files` remains available when a descriptor is not already present. Coverage includes native snapshot preservation after overwrite, both row IDs and addresses, dictionary/Arrow inputs, null/empty values, reordered and duplicate descriptors, lazy iterable behavior, async iteration, and remote range pinning for all six storage kinds. Descriptor reads isolate their pinned snapshot from the live table handle's freshness state. Both native and remote reads verify the source commit timestamp to reject a replaced snapshot after table recreation. --- docs/src/python/python.md | 54 +++++++- python/README.md | 49 +++++++ python/python/lancedb/_blob.py | 14 +- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/remote/table.py | 5 + python/python/lancedb/table.py | 76 ++++++++++ python/python/tests/test_blob.py | 132 ++++++++++++++++++ python/src/table.rs | 19 ++- rust/lancedb/src/blob.rs | 111 ++++++++++++++- rust/lancedb/src/remote/client.rs | 5 + rust/lancedb/src/remote/table.rs | 4 + rust/lancedb/src/remote/table/blobs.rs | 183 ++++++++++++++++++++++++- rust/lancedb/src/table.rs | 33 ++++- 13 files changed, 674 insertions(+), 12 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 41d3e16d7..02ccced2c 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -538,12 +538,64 @@ instead of being materialized with the rest of the row. `lancedb.BlobType` is `lance.blob.BlobType` when pylance is installed. Without pylance, LanceDB uses a matching `lance.blob.v2` extension type so blob columns -still work. Queries return descriptors. Call +still work. Native table queries return descriptors. Call [`fetch_blob_files`][lancedb.table.Table.fetch_blob_files] for lazy reads or [`fetch_blobs`][lancedb.table.Table.fetch_blobs] for eager bytes. One missing or deleted row ID rejects the entire blob-read batch; the APIs do not return partial results. +Remote SQL returns blob bytes by default. Select `blob_descriptor(video)` to +return an openable reference instead, then use +[`open_blob_file`][lancedb.table.Table.open_blob_file] or the lazy iterable +[`open_blob_files`][lancedb.table.Table.open_blob_files]: + +```python +rows = db.execute_query(""" + SELECT f.timestamp, blob_descriptor(v.video) AS video + FROM videos v JOIN frames f ON v.id = f.video_id +""").read_all() +videos = db.open_table("videos") +for row in rows.to_pylist(): + if row["video"] is not None: + with videos.open_blob_file(row["video"]) as video: + frame = extract_frame(video, row["timestamp"]) +``` + +This example uses a remote SQL connection that produces source-bearing +descriptors. The descriptor contains `kind`, `position`, `size`, `blob_id`, +`blob_uri`, and a `source` struct with `uri`, `version`, `committed_at_micros`, +`branch`, `column`, `row_reference`, and `row_is_address`. The opener verifies +the source table and commit and reads the captured version without changing +the table handle's checkout. Raw native storage descriptors omit `source`; +use `fetch_blob_files(column, row_ids)` for those query results. + +The singular API accepts a dictionary or Arrow struct scalar. The iterable +accepts Arrow arrays/chunked arrays or another iterable, opens handles as +iteration advances, and preserves order, duplicates, and nulls as `None`. +Close each non-null handle after use. Async tables provide +`await table.open_blob_file(...)` and +`async for file in table.open_blob_files(...)`. + +```python +for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): + if video is not None: + with video: + frame = extract_frame(video, row["timestamp"]) +``` + +For an async table, the iterable is consumed with `async for`: + +```python +async for video in async_videos.open_blob_files(rows["video"]): + if video is not None: + with video: + payload = await video.aread() +``` + +Payload bytes load on reads. Opening an external blob with unspecified size +can issue a metadata-only HEAD request. Remote reads use authenticated range +requests that validate the captured URI and commit timestamp. + Blobs inside lists can be written and queried as descriptors. They are omitted from `blob_columns()` because row ids do not identify individual list elements, so the fetch APIs cannot read their bytes yet. diff --git a/python/README.md b/python/README.md index 907f7c099..0b640d27c 100644 --- a/python/README.md +++ b/python/README.md @@ -101,3 +101,52 @@ print(results) ### Development See [CONTRIBUTING.md](./CONTRIBUTING.md) for information on how to contribute to LanceDB. + +### Opening SQL blob descriptors + +A SQL `blob_descriptor(video)` result keeps the Blob V2 storage metadata and a +`source` struct containing its table URI, version, commit timestamp, branch, +column and row identity. This example uses a remote SQL connection whose SQL +service produces those source-bearing descriptors. Open the reference on the +source table to read incrementally: + +```python +rows = db.execute_query(""" + SELECT f.timestamp, blob_descriptor(v.video) AS video + FROM videos v JOIN frames f ON v.id = f.video_id +""").read_all() + +videos = db.open_table("videos") +for row in rows.to_pylist(): + if row["video"] is not None: + with videos.open_blob_file(row["video"]) as video: + frame = extract_frame(video, row["timestamp"]) +``` + +`open_blob_file` accepts the dictionary or an Arrow struct scalar directly. It +keeps the query's version and branch without changing the table handle's checkout. +Payload bytes load on reads. An external blob with an unspecified size may need a +metadata-only HEAD request to open. + +`videos.open_blob_files(rows["video"])` returns an iterator over the same handles, +preserving order, duplicates and nulls (`None`). It accepts Arrow arrays/chunked +arrays or another iterable and opens each handle as iteration advances. Close each +non-null handle after use. Async tables offer `await videos.open_blob_file(...)` +and `async for video in videos.open_blob_files(...)`. + +```python +for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + header = video.read(4096) +``` + +To extract a frame while retaining each row's timestamp, iterate over the +handles and rows together: + +```python +for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): + if video is not None: + with video: + frame = extract_frame(video, row["timestamp"]) +``` diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 1e9110dd0..46462553c 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -7,7 +7,7 @@ from __future__ import annotations import io import operator -from collections.abc import Awaitable, Callable, Iterable +from collections.abc import Awaitable, Callable, Iterable, Mapping from typing import TYPE_CHECKING, Optional, Union import pyarrow as pa @@ -31,6 +31,18 @@ FetchBlobsSync = Callable[[str, list[int]], pa.Array | pa.ChunkedArray] FetchBlobsAsync = Callable[[str, list[int]], Awaitable[pa.Array | pa.ChunkedArray]] +def _normalize_blob_descriptor(value): + if isinstance(value, pa.Scalar): + value = value.as_py() + if value is None: + return None + if not isinstance(value, Mapping): + raise TypeError( + "blob descriptor must be a mapping, an Arrow struct scalar, or None" + ) + return dict(value) + + class BlobFile(io.RawIOBase): """Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`. diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 9e602cce8..0b4fc82f1 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -580,6 +580,7 @@ class Table: async def fetch_blob_files( self, column: str, row_ids: list[int] ) -> list[Optional[BlobFile]]: ... + async def open_blob_file(self, descriptor: str) -> BlobFile: ... def vector_search(self) -> VectorQuery: ... class Tags: diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 53591182e..dd22dea48 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1221,6 +1221,11 @@ class RemoteTable(Table): ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + return LOOP.run(self._table.open_blob_file(descriptor)) + def head(self, n=5) -> pa.Table: """ Return the first `n` rows of the table. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index c08c6390e..665520479 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -15,9 +15,11 @@ from functools import cached_property from typing import ( TYPE_CHECKING, Any, + AsyncIterator, Callable, Dict, Iterable, + Iterator, List, Literal, Optional, @@ -33,6 +35,7 @@ from lancedb.scannable import _hf_batches, _register_optional_converters, to_sca from . import __version__ from ._blob import ( BlobFile, + _normalize_blob_descriptor, _normalize_blob_row_ids, _wrap_blob_files, strip_auto_row_ids, @@ -2289,6 +2292,44 @@ class Table(ABC): handles are returned. """ + @abstractmethod + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + """Open a seekable file from a SQL ``blob_descriptor`` result. + + Accepts the descriptor dictionary or an Arrow struct scalar. Its source + pins the table version, branch, column and row identity, so opening does + not change this table's checkout. Null descriptors return ``None``. + Payload bytes load only when the returned handle is read. An external + blob with an unspecified size may require a metadata request on opening. + + ```python + rows = db.execute_query( + "SELECT blob_descriptor(video) AS video FROM videos" + ).read_all() + with videos.open_blob_file(rows["video"][0]) as video: + header = video.read(4096) + ``` + """ + + def open_blob_files(self, descriptors: Iterable) -> Iterator[Optional[BlobFile]]: + """Iterate over lazy files from descriptor dictionaries or Arrow scalars. + + Accepts an Arrow array/chunked array or any iterable. Opens each handle + as iteration advances, preserving input order and nulls as ``None``. + Close each non-null handle after use, normally with a ``with`` block. + + ```python + for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + header = video.read(4096) + ``` + """ + for descriptor in descriptors: + yield self.open_blob_file(descriptor) + @abstractmethod def _execute_query( self, @@ -3239,6 +3280,11 @@ class LanceTable(Table): ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + return LOOP.run(self._table.open_blob_file(descriptor)) + @property def tags(self) -> Tags: """Tag management for the table. @@ -7451,6 +7497,36 @@ class AsyncTable: ) return _wrap_blob_files(handles) + async def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + """Open a snapshot-bound descriptor as a seekable file. + + See [Table.open_blob_file][lancedb.table.Table.open_blob_file]. + """ + descriptor = _normalize_blob_descriptor(descriptor) + if descriptor is None: + return None + return BlobFile(await self._inner.open_blob_file(json.dumps(descriptor))) + + async def open_blob_files( + self, descriptors: Iterable + ) -> AsyncIterator[Optional[BlobFile]]: + """Asynchronously iterate over files in descriptor order, preserving nulls. + + Use ``async for file in table.open_blob_files(descriptors)``. Each + non-null file must be closed after use. + + ```python + async for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + payload = await video.aread() + ``` + """ + for descriptor in descriptors: + yield await self.open_blob_file(descriptor) + @property def tags(self) -> AsyncTags: """Tag management for the dataset. diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index c2dc2a265..da362f053 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -2696,3 +2696,135 @@ def test_create_table_with_python_blob_dict_and_object(): b"from-dict", b"from-blob", ] + + +def _openable_descriptors(table, *, by_address): + hits = ( + table.search() + .select(["id", "image", "_rowaddr"]) + .with_row_id(True) + .limit(100) + .to_arrow(blob_mode="descriptions") + ) + version = table.version + committed_at = next( + item["timestamp"] + for item in table.list_versions() + if item["version"] == version + ) + descriptors = {} + for row in hits.to_pylist(): + descriptor = row["image"] + if descriptor is not None and descriptor.get("kind") is not None: + descriptor = dict(descriptor) + descriptor["source"] = { + "uri": table.to_lance().uri, + "version": version, + "committed_at_micros": round(committed_at.timestamp() * 1_000_000), + "branch": None, + "column": "image", + "row_reference": row["_rowaddr" if by_address else "_rowid"], + "row_is_address": by_address, + } + else: + descriptor = None + descriptors[row["id"]] = descriptor + return descriptors + + +@pytest.mark.parametrize("by_address", [False, True]) +def test_open_blob_files_preserve_snapshot_order_nulls_and_empty_values( + tmp_path, by_address +): + db = lancedb.connect(str(tmp_path)) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("videos", schema=schema) + table.add( + [ + {"id": 1, "image": b"alpha"}, + {"id": 2, "image": None}, + {"id": 3, "image": b"beta"}, + {"id": 4, "image": b""}, + ] + ) + descriptors = _openable_descriptors(table, by_address=by_address) + table.add([{"id": 1, "image": b"replacement"}], mode="overwrite") + latest_version = table.version + values = pa.array( + [descriptors[3], descriptors[2], descriptors[1], descriptors[3], descriptors[4]] + ) + files = table.open_blob_files( + pa.chunked_array([values.slice(0, 2), values.slice(2)]) + ) + assert iter(files) is files + actual = [] + for file in files: + if file is None: + actual.append(None) + else: + with file: + actual.append(file.read()) + assert file.closed + assert actual == [b"beta", None, b"alpha", b"beta", b""] + with table.open_blob_file(descriptors[1]) as file: + file.seek(2) + assert file.read_range(0, 2) == b"al" + assert file.tell() == 2 + assert file.read(2) == b"ph" + assert table.version == latest_version + assert table.open_blob_file(pa.scalar(None, type=values.type)) is None + other = db.create_table("other", schema=schema) + with pytest.raises(ValueError, match="different table"): + other.open_blob_file(descriptors[1]) + stale = dict(descriptors[1], source=dict(descriptors[1]["source"])) + stale["source"]["committed_at_micros"] += 1 + with pytest.raises(ValueError, match="snapshot has been replaced"): + table.open_blob_file(stale) + raw = {key: value for key, value in descriptors[1].items() if key != "source"} + with pytest.raises(ValueError, match="source"): + table.open_blob_file(raw) + with pytest.raises(TypeError, match="descriptor"): + table.open_blob_file("not a descriptor") + + +def test_open_blob_files_does_not_open_ahead(monkeypatch): + table = _blob_table("lazy_descriptor_iteration", [{"id": 1, "image": b"x"}]) + opened = [] + + def open_file(descriptor): + opened.append(descriptor) + return descriptor + + monkeypatch.setattr(table, "open_blob_file", open_file) + files = table.open_blob_files(iter([1, None, 2])) + assert opened == [] + assert next(files) == 1 + assert opened == [1] + assert list(files) == [None, 2] + assert opened == [1, None, 2] + + +@pytest.mark.asyncio +async def test_open_blob_files_async_iterates_snapshot_handles(tmp_path): + db = lancedb.connect(str(tmp_path)) + table = db.create_table( + "videos", + schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), + ) + table.add([{"id": 1, "image": b"alpha"}, {"id": 2, "image": None}]) + descriptors = _openable_descriptors(table, by_address=True) + async_db = await lancedb.connect_async(str(tmp_path)) + videos = await async_db.open_table("videos") + file = await videos.open_blob_file(pa.scalar(descriptors[1])) + with file: + assert file.read(2) == b"al" + actual = [] + async for file in videos.open_blob_files( + [descriptors[2], descriptors[1], descriptors[1]] + ): + if file is None: + actual.append(None) + else: + with file: + actual.append(await file.aread()) + assert actual == [None, b"alpha", b"alpha"] diff --git a/python/src/table.rs b/python/src/table.rs index e77c58578..e58b8a139 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -24,7 +24,7 @@ use arrow::{ ffi_stream::ArrowArrayStreamReader, pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; -use lancedb::blob::{BlobFile, BlobRangeRequest}; +use lancedb::blob::{BlobDescriptor, BlobFile, BlobRangeRequest}; use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, @@ -1551,6 +1551,23 @@ impl Table { }) } + #[pyo3(signature = (descriptor))] + pub fn open_blob_file<'py>( + self_: PyRef<'py, Self>, + descriptor: &str, + ) -> PyResult> { + let descriptor: BlobDescriptor = serde_json::from_str(descriptor) + .map_err(|error| PyValueError::new_err(format!("invalid blob descriptor: {error}")))?; + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let file = inner.open_blob_file(&descriptor).await.infer_error()?; + Ok(PyBlobFile { + inner: Arc::new(file), + closed: AtomicBool::new(false), + }) + }) + } + /// Optimize the on-disk data by compacting and pruning old data, for better performance. #[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))] pub fn optimize( diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index 05650d7c2..142b2264e 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -20,11 +20,120 @@ use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams}; use lance_arrow::FieldExt; use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; -use lance_io::object_store::ObjectStore; +use lance_io::object_store::{ObjectStore, uri_to_url}; use object_store::path::Path; +use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; +/// A stored Blob V2 descriptor with the source needed to reopen its snapshot. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct BlobDescriptor { + /// Canonical Blob V2 storage kind (0 through 5). + pub kind: u8, + /// Stored offset within the backing object. + pub position: u64, + /// Stored length; zero on an external blob can mean unspecified. + pub size: u64, + /// Stored sidecar or base identifier. + pub blob_id: u32, + /// Stored external URI or managed path. + pub blob_uri: String, + /// Source identity captured by the query. + pub source: BlobDescriptorSource, +} + +/// The pinned source of a descriptor returned by SQL. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct BlobDescriptorSource { + /// URI of the dataset on the source branch. + pub uri: String, + /// Exact table version read by the query. + pub version: u64, + /// Commit timestamp in microseconds, distinguishing recreated tables. + pub committed_at_micros: u64, + /// Source branch; `None` identifies main. + pub branch: Option, + /// Blob column path in the source table. + pub column: String, + /// Row ID or row address, as specified by `row_is_address`. + pub row_reference: u64, + /// Whether the row reference is a physical address. + pub row_is_address: bool, +} + +impl BlobDescriptor { + pub(crate) fn validate(&self) -> Result<()> { + if self.kind > 5 + || self.source.version == 0 + || self.source.uri.is_empty() + || self.source.column.is_empty() + { + return Err(Error::InvalidInput { + message: "invalid stored blob descriptor or missing pinned source".into(), + }); + } + Ok(()) + } +} + +pub(crate) async fn open_native_blob_file( + dataset: &Dataset, + descriptor: &BlobDescriptor, +) -> Result { + descriptor.validate()?; + let source = &descriptor.source; + let expected_uri = dataset + .branch_location() + .find_branch(source.branch.as_deref())? + .uri; + if uri_to_url(&expected_uri)? != uri_to_url(&source.uri)? { + return Err(Error::InvalidInput { + message: "blob descriptor belongs to a different table".into(), + }); + } + let dataset = Arc::new( + dataset + .checkout_version(lance::dataset::refs::Ref::Version( + source.branch.clone(), + Some(source.version), + )) + .await?, + ); + if uri_to_url(dataset.uri())? != uri_to_url(&source.uri)? { + return Err(Error::InvalidInput { + message: "blob descriptor belongs to a different table".into(), + }); + } + if dataset.manifest().timestamp_nanos / 1_000 != u128::from(source.committed_at_micros) { + return Err(Error::InvalidInput { + message: "blob descriptor source snapshot has been replaced".into(), + }); + } + ensure_blob_v2_column(dataset.schema(), &source.column)?; + let mut files = if source.row_is_address { + dataset + .take_blobs_by_addresses(&[source.row_reference], &source.column) + .await? + } else { + dataset + .take_blobs(&[source.row_reference], &source.column) + .await? + }; + if files.len() != 1 { + return Err(Error::InvalidInput { + message: "blob descriptor does not resolve to one stored value".into(), + }); + } + files + .pop() + .flatten() + .map(Into::into) + .ok_or_else(|| Error::InvalidInput { + message: "non-null blob descriptor resolved to a null value".into(), + }) +} + /// Seekable handle for one blob value, backed by local storage or a remote /// HTTP byte-range endpoint. #[derive(Debug)] diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index fe4a2072e..f54e2fe3e 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -809,6 +809,11 @@ impl RestfulLanceDbClient { self.client.get(full_uri) } + pub fn head(&self, uri: &str) -> RequestBuilder { + let full_uri = format!("{}{}", self.host, uri); + self.client.head(full_uri) + } + pub fn post(&self, uri: &str) -> RequestBuilder { let full_uri = format!("{}{}", self.host, uri); self.client.post(full_uri) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 4cbba04df..181ce61f9 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2927,6 +2927,10 @@ impl BaseTable for RemoteTable { self.fetch_blob_files_impl(column, row_ids).await } + async fn open_blob_file(&self, descriptor: &crate::blob::BlobDescriptor) -> Result { + self.open_blob_file_impl(descriptor).await + } + async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result { if let AnyQuery::Query(request) = query && let Some(offsets) = &request.take_offsets diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index bcaf47765..e3cdfc6a6 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -21,7 +21,7 @@ use reqwest::{Response, StatusCode, header}; use tokio::sync::{Mutex, OwnedSemaphorePermit}; use crate::Error; -use crate::blob::{BlobFile, BlobRangeRequest}; +use crate::blob::{BlobDescriptor, BlobDescriptorSource, BlobFile, BlobRangeRequest}; use crate::error::Result; use crate::query::{QueryFilter, QueryRequest, Select}; use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; @@ -49,6 +49,7 @@ fn is_fetch_blobs_byte_limit_error(error: &Error) -> bool { enum RangeRequestMode { SizeProbe, DataRead, + MetadataRead, } #[async_trait::async_trait] @@ -86,6 +87,9 @@ struct TableBlobRangeRequester { path: String, version: Option, branch: Option, + source_uri: Option, + source_commit_micros: Option, + row_is_address: bool, freshness: Arc>, parent_freshness: Arc>, parent_freshness_request: FreshnessHeaders, @@ -109,15 +113,27 @@ impl BlobRangeRequester for TableBlobRangeRequester { .expect("blob request semaphore is never closed"); let freshness_request = freshness_headers_snapshot(&self.freshness, self.read_consistency_interval); - let mut request = freshness_request - .apply(self.client.get(&self.path)) - .header(header::RANGE, range_header); + let mut request = if matches!(mode, RangeRequestMode::MetadataRead) { + freshness_request.apply(self.client.head(&self.path)) + } else { + freshness_request + .apply(self.client.get(&self.path)) + .header(header::RANGE, range_header) + }; if let Some(version) = version.or(self.version) { request = request.query(&[("version", version)]); } if let Some(branch) = &self.branch { request = request.query(&[("branch", branch)]); } + if let Some(source_uri) = &self.source_uri { + request = request + .query(&[("source_uri", source_uri)]) + .query(&[("row_is_address", self.row_is_address)]); + } + if let Some(timestamp) = self.source_commit_micros { + request = request.query(&[("source_commit_micros", timestamp)]); + } let (request_id, response) = self.client.send_with_retry(request, None, true).await?; // Preserve 416 size-probe responses so the caller can detect empty blobs. if response.status() == StatusCode::RANGE_NOT_SATISFIABLE @@ -131,8 +147,10 @@ impl BlobRangeRequester for TableBlobRangeRequester { } let response = self.client.check_response(&request_id, response).await?; freshness_request.observe_headers(&self.freshness, response.headers()); - self.parent_freshness_request - .observe_headers(&self.parent_freshness, response.headers()); + if self.source_uri.is_none() { + self.parent_freshness_request + .observe_headers(&self.parent_freshness, response.headers()); + } Ok(BlobRangeResponse { request_id, response, @@ -630,6 +648,51 @@ impl RemoteTable { Ok(files) } + pub(super) async fn open_blob_file_impl( + &self, + descriptor: &BlobDescriptor, + ) -> Result { + descriptor.validate()?; + self.ensure_blob_files_supported()?; + let source = &descriptor.source; + let snapshot = ReadSnapshot { + version: Some(source.version), + freshness_state: FreshnessState { + pinned: true, + ..Default::default() + }, + ..Default::default() + }; + let requester = self.blob_range_requester_with_source( + &source.column, + source.row_reference, + &snapshot, + Some(source), + ); + let size = if descriptor.kind == 3 && descriptor.size == 0 { + let metadata = requester + .request_range("", RangeRequestMode::MetadataRead, Some(source.version)) + .await?; + if metadata.response.status() != StatusCode::OK { + return Err(Error::InvalidInput { + message: "non-null blob descriptor did not resolve to a stored value".into(), + }); + } + metadata + .response + .headers() + .get(header::CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .ok_or_else(|| Error::InvalidInput { + message: "blob metadata response has no valid Content-Length".into(), + })? + } else { + descriptor.size + }; + Ok(RemoteBlobFile::new(requester, size, Some(source.version)).into()) + } + /// Probe each row once for its size. /// /// Serves callers that must not issue a query, such as a single blob that @@ -793,6 +856,16 @@ impl RemoteTable { column: &str, row_id: u64, read_snapshot: &ReadSnapshot, + ) -> Arc { + self.blob_range_requester_with_source(column, row_id, read_snapshot, None) + } + + fn blob_range_requester_with_source( + &self, + column: &str, + row_id: u64, + read_snapshot: &ReadSnapshot, + source: Option<&BlobDescriptorSource>, ) -> Arc { let path = format!( "/v1/table/{}/blob/{}/{row_id}/bytes", @@ -803,7 +876,10 @@ impl RemoteTable { client: self.client.clone(), path, version: read_snapshot.version, - branch: self.branch.clone(), + branch: source.map_or_else(|| self.branch.clone(), |source| source.branch.clone()), + source_uri: source.map(|source| source.uri.clone()), + source_commit_micros: source.map(|source| source.committed_at_micros), + row_is_address: source.is_some_and(|source| source.row_is_address), freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)), parent_freshness: self.freshness.clone(), parent_freshness_request: read_snapshot.freshness, @@ -1024,6 +1100,99 @@ mod tests { .unwrap() } + #[tokio::test] + async fn open_descriptor_files_are_lazy_and_pin_ranges_for_every_storage_kind() { + for kind in 0..=5 { + let requests = Arc::new(StdMutex::new(Vec::new())); + let captured = requests.clone(); + let table = RemoteTable::new_mock( + "videos".to_string(), + move |request| { + for name in [ + "x-lancedb-min-version", + "x-lancedb-min-read-version", + "x-lancedb-min-timestamp", + ] { + assert!(request.headers().get(name).is_none()); + } + assert_eq!( + request.url().path(), + "/v1/table/videos/blob/video/12884901889/bytes" + ); + let parameters = request + .url() + .query_pairs() + .into_owned() + .collect::>(); + assert_eq!(parameters["version"], "7"); + assert_eq!(parameters["source_commit_micros"], "1234"); + assert_eq!(parameters["branch"], "dev"); + assert_eq!(parameters["row_is_address"], "true"); + assert_eq!( + parameters["source_uri"], + "s3://bucket/videos.lance/tree/dev" + ); + captured.lock().unwrap().push(request.method().clone()); + if request.method() == reqwest::Method::HEAD { + assert!(request.headers().get(header::RANGE).is_none()); + http::Response::builder() + .status(200) + .header(header::CONTENT_LENGTH, PAYLOAD.len().to_string()) + .body(Vec::new()) + .unwrap() + } else { + range_response(&request, PAYLOAD) + } + }, + Some(Version::new(0, 5, 0)), + ); + { + let mut freshness = table.freshness.lock().unwrap(); + freshness.min_version = Some(99); + freshness.min_read_version = Some(2); + } + let descriptor = BlobDescriptor { + kind, + position: 0, + size: if kind == 3 { 0 } else { PAYLOAD.len() as u64 }, + blob_id: 0, + blob_uri: String::new(), + source: BlobDescriptorSource { + uri: "s3://bucket/videos.lance/tree/dev".into(), + version: 7, + committed_at_micros: 1234, + branch: Some("dev".into()), + column: "video".into(), + row_reference: 12884901889, + row_is_address: true, + }, + }; + let file = table.open_blob_file(&descriptor).await.unwrap(); + let expected = if kind == 3 { + vec![reqwest::Method::HEAD] + } else { + vec![] + }; + assert_eq!(*requests.lock().unwrap(), expected); + assert_eq!(file.size(), PAYLOAD.len() as u64); + file.seek(4).await.unwrap(); + assert_eq!(file.read_up_to(5).await.unwrap().as_ref(), &PAYLOAD[4..9]); + assert_eq!(file.tell().await.unwrap(), 9); + assert_eq!(table.freshness.lock().unwrap().min_read_version, Some(2)); + assert_eq!(file.read_range(0..3).await.unwrap().as_ref(), &PAYLOAD[..3]); + assert_eq!(file.tell().await.unwrap(), 9); + assert_eq!( + requests + .lock() + .unwrap() + .iter() + .filter(|method| **method == reqwest::Method::GET) + .count(), + 2 + ); + } + } + fn mock_remote_blob_table( requests: Arc>>, ) -> RemoteTable { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 3123eb776..3fcf4bdb7 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3,7 +3,7 @@ //! LanceDB Table APIs -use crate::blob::BlobFile; +use crate::blob::{BlobDescriptor, BlobFile}; use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; use arrow_schema::{Schema, SchemaRef}; use async_trait::async_trait; @@ -796,6 +796,12 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "fetch_blob_files is not supported on this table type".into(), }) } + /// Open a descriptor at its pinned source version. See [`Table::open_blob_file`]. + async fn open_blob_file(&self, _descriptor: &BlobDescriptor) -> Result { + Err(Error::NotSupported { + message: "open_blob_file is not supported on this table type".into(), + }) + } /// Gets the table tag manager. async fn tags(&self) -> Result>; /// Optimize the dataset. @@ -1377,6 +1383,26 @@ impl Table { self.inner.fetch_blob_files(column.as_ref(), row_ids).await } + /// Open a seekable file from a SQL `blob_descriptor` result. + /// + /// The descriptor carries its table version, branch, column and row identity. + /// Opening does not materialize the payload or change this table's checkout. + /// Remote reads use the table's authenticated HTTP byte-range endpoint. + /// + /// ``` + /// use lancedb::{Table, blob::BlobDescriptor}; + /// + /// # async fn read_header(videos: &Table, descriptor: &BlobDescriptor) -> lancedb::Result<()> { + /// let video = videos.open_blob_file(descriptor).await?; + /// let header = video.read_range(0..video.size().min(4096)).await?; + /// video.close().await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn open_blob_file(&self, descriptor: &BlobDescriptor) -> Result { + self.inner.open_blob_file(descriptor).await + } + /// Insert new records into this Table /// /// # Arguments @@ -3764,6 +3790,11 @@ impl BaseTable for NativeTable { crate::blob::take_blob_files_aligned(&dataset, column, row_ids).await } + async fn open_blob_file(&self, descriptor: &BlobDescriptor) -> Result { + let dataset = self.dataset.get().await?; + crate::blob::open_native_blob_file(&dataset, descriptor).await + } + /// Delete rows from the table async fn delete(&self, predicate: Predicate<'_>) -> Result { let result = delete::execute_delete(self, predicate).await?; From de0c5edc7d05cea0a79adb180a5ae374514c69e4 Mon Sep 17 00:00:00 2001 From: Lance Release Date: Sun, 25 Oct 2026 11:22:16 +0000 Subject: [PATCH 26/36] =?UTF-8?q?Bump=20version:=200.41.0-beta.13=20?= =?UTF-8?q?=E2=86=92=200.41.0-beta.14?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .bumpversion.toml | 2 +- Cargo.lock | 6 +++--- docs/src/java/java.md | 2 +- java/lancedb-core/pom.xml | 2 +- java/pom.xml | 2 +- nodejs/Cargo.toml | 2 +- nodejs/npm/darwin-arm64/package.json | 2 +- nodejs/npm/linux-arm64-gnu/package.json | 2 +- nodejs/npm/linux-arm64-musl/package.json | 2 +- nodejs/npm/linux-x64-gnu/package.json | 2 +- nodejs/npm/linux-x64-musl/package.json | 2 +- nodejs/npm/win32-arm64-msvc/package.json | 2 +- nodejs/npm/win32-x64-msvc/package.json | 2 +- nodejs/package.json | 2 +- python/Cargo.toml | 2 +- rust/lancedb/Cargo.toml | 2 +- 16 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.bumpversion.toml b/.bumpversion.toml index 0909579cb..7dcefa69c 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.13" +current_version = "0.41.0-beta.14" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/Cargo.lock b/Cargo.lock index fac89c4da..03f4b0864 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5700,7 +5700,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "ahash", "anyhow", @@ -5798,7 +5798,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "arrow-array", "arrow-buffer", @@ -5823,7 +5823,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "arc-swap", "arrow", diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 0c887f66f..9dbccbc81 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.41.0-beta.13 + 0.41.0-beta.14 ``` diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 87c5d4cb8..7c34a801a 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.13 + 0.41.0-beta.14 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 1f3e0efb0..2dc959887 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.13 + 0.41.0-beta.14 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 1b9d47e46..e343d68ca 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index bd4b95463..343483a8e 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index 6e4e57684..4c3b93c81 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index f77c50ee2..7679a422d 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index b671b7ebf..6a62f8e8a 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index b75738ff1..d68341382 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 7158748f9..42a517310 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 83d0fb6f2..b116f2819 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 21f903235..f746f6daf 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index 076299a7d..82fb7efae 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 486eb1e8e..006886647 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true From 197403e8913fc1976d72ede8a73b7e0c69b6ff10 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 25 Oct 2026 07:19:55 -0700 Subject: [PATCH 27/36] fix: honor downstream Lance dependency overrides (#4671) Consumers that provide a local Lance checkout cannot resolve the dependency graph because the SDK directly selects a vendored `lance-arrow`, producing two path packages with the same name and version. The Rust core now declares its upstream Git dependency separately so the consuming workspace can override that source consistently. Python continues inheriting the vendored workspace dependency, keeping the sliced-struct validity fix in Maturin source distributions. Standalone SDK builds still select that same fix through the existing root patch. Adds a downstream consumer CI regression that resolves a local Lance override and checks that exactly one canonical `lance-arrow` package is selected. The Lance version updater also updates the core declaration in stable, preview and local modes. Public APIs and query behavior are unchanged. The workspace declaration lint excludes path dependencies because the Python binding intentionally uses the bundled source. Duplicate and unused external declarations remain errors; advisory, license, banned-package and source checks continue covering the full dependency graph. --------- Co-authored-by: Lance Release --- .github/workflows/pypi-publish.yml | 3 ++ .github/workflows/rust.yml | 57 ++++++++++++++++++++++++++++++ ci/set_lance_version.py | 14 ++++++++ ci/tests/test_set_lance_version.py | 28 +++++++++++++++ deny.toml | 4 ++- python/Cargo.toml | 1 - rust/lancedb/Cargo.toml | 2 +- 7 files changed, 106 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml index e71564473..1784c4c8a 100644 --- a/.github/workflows/pypi-publish.yml +++ b/.github/workflows/pypi-publish.yml @@ -15,6 +15,9 @@ on: - ci/check_cargo_lock.py - ci/check_sdist.py - python/pyproject.toml + - python/Cargo.toml + - python/vendor/** + - rust/lancedb/Cargo.toml - Cargo.toml # Change in dependency frequently breaks builds - Cargo.lock diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 930ae67e2..17fe3a7b4 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -29,6 +29,63 @@ env: RUST_BACKTRACE: "1" jobs: + dependency-overrides: + name: Downstream Lance override + runs-on: ubuntu-24.04 + timeout-minutes: 15 + steps: + - uses: actions/checkout@v6 + - uses: actions-rust-lang/setup-rust-toolchain@v1 + - name: Read the pinned Lance tag + id: lance + run: | + python3 - <<'PYTHON' + import os + import tomllib + from pathlib import Path + + manifest = tomllib.loads(Path("Cargo.toml").read_text()) + tag = manifest["workspace"]["dependencies"]["lance"]["tag"] + with open(os.environ["GITHUB_OUTPUT"], "a") as output: + output.write(f"tag={tag}\n") + PYTHON + - uses: actions/checkout@v6 + with: + repository: lance-format/lance + ref: ${{ steps.lance.outputs.tag }} + path: local-lance + - name: Resolve a downstream workspace with a local Lance override + run: | + mkdir -p downstream/src + touch downstream/src/lib.rs + cat > downstream/Cargo.toml <<'TOML' + [package] + name = "lancedb-local-lance-consumer" + version = "0.1.0" + edition = "2024" + + [workspace] + + [dependencies] + lancedb = { path = "../rust/lancedb", default-features = false } + lance-arrow = { path = "../local-lance/rust/lance-arrow" } + + [patch."https://github.com/lance-format/lance.git"] + lance-arrow = { path = "../local-lance/rust/lance-arrow" } + TOML + cargo metadata --manifest-path downstream/Cargo.toml --format-version 1 > downstream/graph.json + python3 - <<'PYTHON' + import json + from pathlib import Path + + graph = json.loads(Path("downstream/graph.json").read_text()) + packages = [p for p in graph["packages"] if p["name"] == "lance-arrow"] + assert len(packages) == 1, [p["id"] for p in packages] + assert Path(packages[0]["manifest_path"]).resolve() == Path( + "local-lance/rust/lance-arrow/Cargo.toml" + ).resolve() + PYTHON + lint: timeout-minutes: 30 runs-on: ubuntu-24.04 diff --git a/ci/set_lance_version.py b/ci/set_lance_version.py index c7bf41f55..bef97f8ad 100644 --- a/ci/set_lance_version.py +++ b/ci/set_lance_version.py @@ -2,6 +2,8 @@ import argparse import re import sys import json +import tomllib +from pathlib import Path LANCE_GIT_URL = "https://github.com/lance-format/lance.git" @@ -148,6 +150,18 @@ def update_cargo_toml(line_updater): with open("Cargo.toml", "w") as f: f.writelines(new_lines) + core_manifest = Path("rust/lancedb/Cargo.toml") + if core_manifest.exists(): + lines = core_manifest.read_text().splitlines(keepends=True) + for index, line in enumerate(lines): + if line.startswith("lance-arrow ="): + updated = tomllib.loads(line_updater(line))["lance-arrow"] + config = {"version": updated} if isinstance(updated, str) else updated + if "path" in config: + config["path"] = "../../" + config["path"] + lines[index] = dict_to_toml_line("lance-arrow", config) + core_manifest.write_text("".join(lines)) + def set_stable_version(version: str): """ diff --git a/ci/tests/test_set_lance_version.py b/ci/tests/test_set_lance_version.py index 1493fc59d..1173eab84 100644 --- a/ci/tests/test_set_lance_version.py +++ b/ci/tests/test_set_lance_version.py @@ -4,6 +4,7 @@ import subprocess import sys import tempfile import textwrap +import tomllib import unittest from pathlib import Path @@ -83,6 +84,13 @@ lance_datafusion = {{ "version" = "=8.2.0-beta.4", "features" = ["substrait"], " with self.subTest(version=version), tempfile.TemporaryDirectory() as tmp: workdir = Path(tmp) (workdir / "Cargo.toml").write_text(CARGO_TOML) + core_manifest = workdir / "rust/lancedb/Cargo.toml" + core_manifest.parent.mkdir(parents=True) + core_manifest.write_text( + "[dependencies]\nlance-arrow = { " + 'version = "=1.0.0", git = "https://github.com/lance-format/lance.git", ' + 'tag = "v1.0.0" }\narrow.workspace = true\n' + ) command_log = workdir / "commands.log" fake_bin = workdir / "bin" fake_bin.mkdir() @@ -111,6 +119,26 @@ lance_datafusion = {{ "version" = "=8.2.0-beta.4", "features" = ["substrait"], " + updated_dependencies + UNTOUCHED_DEPENDENCIES, ) + core = tomllib.loads(core_manifest.read_text())["dependencies"] + dependency = core["lance-arrow"] + if isinstance(dependency, str): + dependency = {"version": dependency} + if version == "local": + self.assertEqual( + dependency["path"], "../../../lance/rust/lance-arrow" + ) + else: + expected_version = { + "stable": "9.9.9", + "preview": "10.0.0-beta.3", + }.get(version, version.removeprefix("v")) + self.assertEqual(dependency["version"], f"={expected_version}") + if "beta" in expected_version: + self.assertEqual(dependency["git"], LANCE_GIT_URL) + self.assertEqual(dependency["tag"], f"v{expected_version}") + else: + self.assertNotIn("git", dependency) + self.assertEqual(core["arrow"], {"workspace": True}) commands = command_log.read_text().splitlines() for command in expected_commands: self.assertTrue( diff --git a/deny.toml b/deny.toml index 77412709f..e4982c149 100644 --- a/deny.toml +++ b/deny.toml @@ -177,7 +177,9 @@ wildcards = "deny" # `[workspace.dependencies]` table: any crate used by more than one member must # go through `workspace = true`, and entries nothing uses are an error. This # keeps versions from drifting between the core crate and the bindings. -workspace-dependencies = { duplicates = "deny", unused = "deny" } +# Python bundles a patched path crate while the Rust SDK uses the overridable +# upstream source. Check shared external declarations, not these path bindings. +workspace-dependencies = { duplicates = "deny", unused = "deny", include-path-dependencies = false } # Internal workspace crates reference each other via `path = "..."`, which # cargo-deny sees as a wildcard version. That's fine for private workspace # members (not published to crates.io), so allow it specifically for paths. diff --git a/python/Cargo.toml b/python/Cargo.toml index 82fb7efae..265e4d270 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -22,7 +22,6 @@ bytes.workspace = true lancedb.workspace = true datafusion-common.workspace = true lance-core.workspace = true -# The workspace path dependency includes the patched crate in Maturin sources. lance-arrow.workspace = true lance-namespace.workspace = true lance-namespace-impls.workspace = true diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 006886647..c18db11fe 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -46,7 +46,7 @@ lance-index = { workspace = true, features = ["tokenizer-jieba", "tokenizer-lind lance-table = { workspace = true } lance-linalg = { workspace = true } lance-encoding = { workspace = true } -lance-arrow = { workspace = true } +lance-arrow = { version = "=15.0.0-beta.1", tag = "v15.0.0-beta.1", git = "https://github.com/lance-format/lance.git" } lance-namespace = { workspace = true } lance-namespace-impls = { workspace = true } metrics = { workspace = true, optional = true } From e6789a4118e037824e7f3891691891841ee2773b Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:55:37 +0800 Subject: [PATCH 28/36] fix: support system columns in DataFusion table adapters (#4667) DataFusion plans SQL against the adapter's declared schema, which omitted `_rowid` and `_rowaddr` even though Lance can project both. SQL therefore rejected these columns before reaching the scan. Add `BaseTableAdapter::with_system_columns` to opt into either column using Lance's nullable `UInt64` fields. Scans without a projection now select the full declared column list so their output matches the provider schema. Default adapters retain data-only `SELECT *` and remain compatible with `INSERT INTO ... SELECT *`; opted-in adapters include the requested system columns in `SELECT *`. Regression tests compare SQL results with native queries and cover projection order, system-only projections, row-ID predicates, scans without a projection, duplicate/unsupported names, and default SQL inserts. Validation: - DataFusion unit suite: 436 tests passed with `cargo test --locked --quiet --profile ci -p lancedb --features remote --lib table::datafusion`. - `cargo clippy --locked --quiet --profile ci -p lancedb --features remote --tests --examples -- -D warnings` passed. - The new builder's documentation example passed its doctest. - `cargo fmt --all -- --check` and `git diff --check` passed. Fixes #4666 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- rust/lancedb/src/table/datafusion.rs | 222 ++++++++++++++++++++++++++- 1 file changed, 221 insertions(+), 1 deletion(-) diff --git a/rust/lancedb/src/table/datafusion.rs b/rust/lancedb/src/table/datafusion.rs index 39f3cd4c6..e6a76c3f5 100644 --- a/rust/lancedb/src/table/datafusion.rs +++ b/rust/lancedb/src/table/datafusion.rs @@ -25,10 +25,11 @@ use datafusion_physical_plan::{ }; use futures::{TryFutureExt, TryStreamExt}; use lance::dataset::{WriteMode, WriteParams}; +use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD}; use super::{AnyQuery, BaseTable}; use crate::{ - Result, + Error, Result, query::{QueryExecutionOptions, QueryFilter, QueryRequest, Select}, }; use arrow_schema::{DataType, Field}; @@ -186,6 +187,46 @@ impl BaseTableAdapter { }) } + /// Expose the requested Lance system columns to DataFusion. + /// + /// Supports [`ROW_ID`] (`_rowid`) and [`ROW_ADDR`] (`_rowaddr`), both nullable + /// `UInt64` fields. Columns are appended in the requested order, skipping any + /// already exposed by this adapter. Other names return [`Error::InvalidInput`]. + /// + /// System columns are hidden by default so existing `SELECT *` queries keep + /// their data-only schema. After opting in, `SELECT *` includes these columns. + /// + /// ``` + /// use lance_core::{ROW_ADDR, ROW_ID}; + /// use lancedb::table::datafusion::BaseTableAdapter; + /// + /// # async fn adapter(table: &lancedb::Table) -> lancedb::Result { + /// BaseTableAdapter::try_new(table.base_table().clone()) + /// .await? + /// .with_system_columns(&[ROW_ID, ROW_ADDR]) + /// # } + /// ``` + pub fn with_system_columns(mut self, columns: &[impl AsRef]) -> Result { + let mut fields = self.schema.fields().to_vec(); + for column in columns { + let column = column.as_ref(); + let field = match column { + ROW_ID => &*ROW_ID_FIELD, + ROW_ADDR => &*ROW_ADDR_FIELD, + _ => { + return Err(Error::InvalidInput { + message: format!("Unsupported system column: {column}"), + }); + } + }; + if !fields.iter().any(|field| field.name() == column) { + fields.push(Arc::new(field.clone())); + } + } + self.schema = Arc::new(ArrowSchema::new(fields)); + Ok(self) + } + /// Set the RTE read routing of every scan; see [`QueryRequest::use_rte`]. pub fn with_use_rte(mut self, use_rte: Option) -> Self { self.use_rte = use_rte; @@ -242,6 +283,15 @@ impl TableProvider for BaseTableAdapter { .map(|i| self.schema.field(*i).name().clone()) .collect(); query.select = Select::Columns(field_names); + } else { + // Select::All excludes system columns, which may be part of our schema. + query.select = Select::Columns( + self.schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(), + ); } if !filters.is_empty() { let first = filters.first().unwrap().clone(); @@ -319,6 +369,7 @@ pub mod tests { BinaryArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, UInt32Array, }; use arrow_schema::{DataType, Field, Schema}; + use arrow_select::concat::concat_batches; use datafusion::{ datasource::provider_as_source, prelude::{SessionConfig, SessionContext}, @@ -328,11 +379,13 @@ pub mod tests { use datafusion_execution::SendableRecordBatchStream; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col, lit}; use futures::{StreamExt, TryStreamExt}; + use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD}; use tempfile::tempdir; use crate::{ connect, index::{Index, scalar::BTreeIndexBuilder}, + query::{ExecutableQuery, Query, QueryBase, Select}, table::datafusion::BaseTableAdapter, }; @@ -509,6 +562,173 @@ pub mod tests { } } + #[tokio::test] + async fn test_system_columns_opt_in() { + let fixture = TestFixture::new().await; + let ctx = SessionContext::new(); + ctx.register_table("foo", fixture.adapter.clone()).unwrap(); + + let batches = ctx + .sql("SELECT * FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].num_columns(), 2); + assert_eq!(batches[0].schema(), fixture.adapter.schema()); + for column in [ROW_ID, ROW_ADDR] { + assert!(ctx.sql(&format!("SELECT {column} FROM foo")).await.is_err()); + } + + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ADDR, ROW_ADDR]) + .unwrap(); + assert_eq!(adapter.schema().fields().len(), 3); + assert_eq!(adapter.schema().field(2), &*ROW_ADDR_FIELD); + assert!(adapter.schema().index_of(ROW_ID).is_err()); + + let adapter = adapter + .with_system_columns(&[ROW_ID.to_string(), ROW_ADDR.to_string()]) + .unwrap(); + assert_eq!(adapter.schema().fields().len(), 4); + assert_eq!(adapter.schema().field(3), &*ROW_ID_FIELD); + assert!(adapter.schema().metadata().is_empty()); + + let err = adapter.with_system_columns(&["_unknown"]).unwrap_err(); + assert!(matches!(err, crate::Error::InvalidInput { .. })); + + // Default adapters must remain compatible with INSERT ... SELECT *. + ctx.sql("INSERT INTO foo SELECT * FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + let batches = ctx + .sql("SELECT COUNT(*) FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + batches[0] + .column(0) + .as_primitive::() + .value(0), + 20 + ); + } + + #[tokio::test] + async fn test_system_columns_sql() { + let fixture = TestFixture::new().await; + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ID, ROW_ADDR]) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("foo", Arc::new(adapter)).unwrap(); + + // Compare SQL results with the native query's row IDs and row addresses. + let native = Query::new(fixture.adapter.table.clone()) + .select(Select::columns(&["i", "indexed", ROW_ADDR])) + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let schema = Arc::new( + native[0] + .schema() + .as_ref() + .clone() + .with_metadata(HashMap::new()), + ); + let expected = concat_batches(&schema, &native).unwrap(); + + for columns in [ + vec!["i", ROW_ID], + vec![ROW_ID, "i"], + vec!["i", ROW_ADDR], + vec![ROW_ADDR, "i"], + vec![ROW_ADDR, ROW_ID], + vec![ROW_ID], + vec![ROW_ADDR], + vec!["i", "indexed", ROW_ID, ROW_ADDR], + ] { + let sql = format!("SELECT {} FROM foo ORDER BY i", columns.join(", ")); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let actual = concat_batches(&batches[0].schema(), &batches).unwrap(); + let indices = columns + .iter() + .map(|name| schema.index_of(name).unwrap()) + .collect::>(); + assert_eq!(actual, expected.project(&indices).unwrap(), "{sql}"); + } + + let batches = ctx + .sql("SELECT * FROM foo ORDER BY i") + .await + .unwrap() + .collect() + .await + .unwrap(); + let indices = ["i", "indexed", ROW_ID, ROW_ADDR].map(|name| schema.index_of(name).unwrap()); + assert_eq!(batches[0], expected.project(&indices).unwrap()); + + // Filtering on the ID must work even when the ID is not in the output. + let row_id = expected + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .value(3); + let batches = ctx + .sql(&format!("SELECT i FROM foo WHERE _rowid = {row_id}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0], expected.project(&[0]).unwrap().slice(3, 1)); + } + + #[tokio::test] + async fn test_system_columns_scan_without_projection() { + let fixture = TestFixture::new().await; + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ADDR, ROW_ID]) + .unwrap(); + let ctx = SessionContext::new(); + let plan = adapter.scan(&ctx.state(), None, &[], None).await.unwrap(); + assert_eq!(plan.schema(), adapter.schema()); + let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx()) + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 10); + for batch in batches { + assert_eq!(batch.schema(), adapter.schema()); + assert_eq!( + batch + .column(2) + .as_primitive::() + .values(), + batch + .column(3) + .as_primitive::() + .values() + ); + } + } + #[tokio::test] async fn test_batch_size() { let fixture = TestFixture::new().await; From 256db55a5195bba35b9e240932d17d580589009b Mon Sep 17 00:00:00 2001 From: Joaquin Hui Gomez <132194176+joaquinhuigomez@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:02:36 +0000 Subject: [PATCH 29/36] fix(python)!: align RemoteDBConnection.create_table parameter order (#4608) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `RemoteDBConnection.create_table` orders its parameters `(name, data, schema, on_bad_vectors, fill_value, mode, exist_ok, ...)`, while `DBConnection`, `LanceDBConnection` and `LanceNamespaceDBConnection` all use `(name, data, schema, mode, exist_ok, on_bad_vectors, fill_value, ...)`. Since `RemoteDBConnection` subclasses `DBConnection`, identical positional calls bind to different parameters: switching a working script from `lancedb.connect(path)` to `lancedb.connect("db://...")` silently turns `mode="overwrite"` into `on_bad_vectors="overwrite"` and `exist_ok=True` into `fill_value=True`, and nothing client-side validates either. `inspect.signature(...).bind(None, "t", data, None, "overwrite", True)` shows it deterministically, no cloud needed. This reorders the remote signature to match the other three; every parameter keeps its own existing default (including the remote's `mode=None`, which its body relies on), and the body already forwards by keyword so nothing else changes. The new test asserts positional parity with `DBConnection` through `embedding_functions` and fails on `main`. Blast radius, stated plainly: only remote callers passing four or more positional arguments are affected — positions 4–7 move; positions 1–3 and 8 are unchanged and keyword callers are unaffected. No in-repo caller passes more than two positionals. Anyone relying on the current order was relying on behaviour that contradicts the documented base contract and diverges from local connections. Not added here: the `storage_options`/`data_storage_version`/`enable_v2_manifest_paths` parameters the remote signature still lacks. --------- Co-authored-by: Xuanwo --- python/python/lancedb/remote/db.py | 4 ++-- python/python/tests/test_remote_db.py | 10 ++++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 62d00f284..c11e01ff4 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -569,10 +569,10 @@ class RemoteDBConnection(DBConnection): name: str, data: DATA = None, schema: Optional[Union[pa.Schema, LanceModel]] = None, - on_bad_vectors: str = "error", - fill_value: float = 0.0, mode: Optional[str] = None, exist_ok: bool = False, + on_bad_vectors: str = "error", + fill_value: float = 0.0, embedding_functions: Optional[List[EmbeddingFunctionConfig]] = None, *, namespace_path: Optional[List[str]] = None, diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index c1db92948..95da3eba2 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -5,6 +5,7 @@ import contextlib from datetime import timedelta import http.server import io +import inspect import json import multiprocessing as mp import pickle @@ -765,6 +766,15 @@ def test_create_table_embedding_functions(with_schema): ] +def test_create_table_positional_parameters_match_base_contract(): + from lancedb.db import DBConnection + from lancedb.remote.db import RemoteDBConnection + + expected = list(inspect.signature(DBConnection.create_table).parameters)[:9] + actual = list(inspect.signature(RemoteDBConnection.create_table).parameters)[:9] + assert actual == expected + + def test_create_table_exist_ok(): def handler(request): if request.path == "/v1/table/test/create/?mode=exist_ok": From 7ed03e1c4006f662e5da4b8e69c7b72f7904bd60 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:46:35 +0800 Subject: [PATCH 30/36] fix(python): honor bad-vector policies for pandas NaNs (#4503) Pandas converts NaN vector components to Arrow nulls. Both Python's NaN mask and Rust's default append validator ignored null children, so `add()` silently stored incomplete vectors despite `on_bad_vectors="error"`. Reject null components in the existing Rust ingestion validator, and include null components in Python's bad-vector detection and filling. `drop` removes affected rows, `fill` repairs missing components while preserving valid ones, and `null` replaces affected vectors. Explicit null components in Arrow and list inputs follow the same policies. Rust allows whole null vectors and retains the explicit `NaNVectorBehavior::Keep` opt-out. The ingestion guide and synchronous/asynchronous `create_table()` option docs describe null components, all four policies, and filling that preserves valid components. Validation: - Reproduced the default pandas append failure with the regression test before changing the Rust validator. - Rebuilt the native Python extension from the merged sources. The Python table and utility suites passed: 481 passed, 1 skipped. Coverage includes create/append/overwrite, pandas/rows/Arrow inputs, every bad-vector policy, float16/32/64, sliced and chunked arrays, all-null components, and hidden children of null lists. - All 43 focused Rust tests passed: 7 validator tests and 36 ingestion tests. - Rust Clippy with `remote`, tests, and examples passed. Ruff lint passes for the Python package, formatting checks pass for modified Python files, and `cargo fmt --all` is clean. - Documentation follow-up: all 56 focused vector-policy tests passed; Python package lint and creation-docstring formatting checks passed. Repository-wide Ruff reports 49 existing errors in untouched documentation examples. Create/overwrite retain their existing `RuntimeError` wrapping for errors raised while consuming lazy Arrow readers. Fixes #4435 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- docs/web/ingestion/create-table.mdx | 8 ++- python/python/lancedb/db.py | 14 ++-- python/python/lancedb/table.py | 48 +++++-------- python/python/tests/test_table.py | 65 +++++++++++++++++ python/python/tests/test_util.py | 54 ++++++++++++-- rust/lancedb/src/table/add_data.rs | 31 +++++--- .../src/table/datafusion/reject_nan.rs | 71 +++++++++++++------ 7 files changed, 214 insertions(+), 77 deletions(-) diff --git a/docs/web/ingestion/create-table.mdx b/docs/web/ingestion/create-table.mdx index 7d4363be7..0921de22e 100644 --- a/docs/web/ingestion/create-table.mdx +++ b/docs/web/ingestion/create-table.mdx @@ -148,9 +148,11 @@ non-nullable column, or provide actual nulls for it, ingestion fails; nullable c omitted or written with null values. Without an explicit schema, Python infers list-like vector values as fixed-size `float32` vector fields from the observed dimension. -For Python ingest, malformed vector values fail by default. If you expect occasional wrong-length, -null, or NaN vectors, choose an `on_bad_vectors` policy: `"drop"` removes those rows, `"fill"` writes -`fill_value`, and `"null"` writes nulls. +For Python ingestion, malformed vectors fail by default (`on_bad_vectors="error"`). This includes +vectors with incorrect sizes, NaNs, or null components, including pandas NaNs converted to Arrow +nulls. Use `"drop"` to remove affected rows, or `"null"` to replace affected vectors with nulls. +Use `"fill"` to replace NaN or null components with `fill_value` while preserving valid components. +This also pads short vectors with `fill_value` and truncates long vectors to the expected dimension. ### From an Arrow Table {#from-an-arrow-table} You can also create LanceDB tables directly from Arrow tables. diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 1558dd52a..9ed772e1c 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -428,10 +428,13 @@ class DBConnection(EnforceOverrides): it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved, and vectors longer than the expected + dimension are truncated. storage_options: dict, optional Additional options for the storage backend. Options already set on the connection will be inherited by the table, but can be overridden here. @@ -2550,10 +2553,13 @@ class AsyncConnection(object): it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved, and vectors longer than the expected + dimension are truncated. storage_options: dict, optional Additional options for the storage backend. Options already set on the connection will be inherited by the table, but can be overridden here. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 665520479..2ad7f821d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -547,10 +547,11 @@ def _sanitize_data( metadata : Optional[dict], default None The embedding metadata to add to the schema. on_bad_vectors : Literal["error", "drop", "fill", "null"], default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. fill_value : float, default 0.0 The value to use when filling vectors. Only used if on_bad_vectors="fill". - All entries in the vector will be set to this value. + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved. """ # At this point, the table might not match the schema we are targeting: # 1. There might be embedding columns missing that will be added @@ -1898,8 +1899,8 @@ class Table(ABC): "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". progress: bool, callable, or tqdm-like, optional @@ -4254,7 +4255,7 @@ class LanceTable(Table): "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -4571,7 +4572,7 @@ class LanceTable(Table): otherwise just open the table, it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -5260,7 +5261,7 @@ def _handle_bad_vector_column( vector_column_name: str The name of the vector column. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0.0 The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -5279,7 +5280,8 @@ def _handle_bad_vector_column( data = data.set_column(position, vector_column_name, vec_arr) if expected_value_type is not None and ( - pa.types.is_integer(vec_arr.type.value_type) + pa.types.is_null(vec_arr.type.value_type) + or pa.types.is_integer(vec_arr.type.value_type) or pa.types.is_unsigned_integer(vec_arr.type.value_type) ): vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type)) @@ -5323,7 +5325,7 @@ def _handle_bad_vector_column( ) else: raise ValueError( - f"Vector column '{vector_column_name}' has NaNs. " + f"Vector column '{vector_column_name}' has NaNs or null values. " "Set on_bad_vectors='drop' to remove them, " "set on_bad_vectors='fill' and fill_value= to replace them, " "or set on_bad_vectors='null' to replace them with null." @@ -5359,29 +5361,11 @@ def _fill_bad_vector_values( arr = arr.combine_chunks() # A fixed-size slice truncates long vectors and pads short vectors with nulls. - # Slice an array marking the original child nulls in parallel so padding nulls - # can be distinguished from null values that were already present. + # Both padding and existing null elements need to be filled. sliced = pc.list_slice(arr, 0, dim, return_fixed_size_list=True) - child_nulls = pc.is_null(arr.values) - parent_nulls = pc.is_null(arr) - if pa.types.is_list(arr.type): - original_child_nulls = pa.ListArray.from_arrays( - arr.offsets, child_nulls, mask=parent_nulls - ) - elif pa.types.is_large_list(arr.type): - original_child_nulls = pa.LargeListArray.from_arrays( - arr.offsets, child_nulls, mask=parent_nulls - ) - else: - original_child_nulls = pa.FixedSizeListArray.from_arrays( - child_nulls, arr.type.list_size, mask=parent_nulls - ) - sliced_child_nulls = pc.list_slice( - original_child_nulls, 0, dim, return_fixed_size_list=True - ) - needs_fill = pc.is_null(sliced_child_nulls.values) values = sliced.values + needs_fill = pc.is_null(values) if pa.types.is_floating(values.type): values_for_nan_check = ( values.cast(pa.float32()) if pa.types.is_float16(values.type) else values @@ -5395,6 +5379,7 @@ def _fill_bad_vector_values( def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray: + """Find vectors with NaN or null components, including pandas-converted NaNs.""" if isinstance(arr, pa.ChunkedArray): values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks]) else: @@ -5405,6 +5390,9 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray values_has_nan = pc.is_nan(values.cast(pa.float32())) else: values_has_nan = pc.is_nan(values) + # Pandas represents NaN components as Arrow nulls. is_nan returns null + # for these values, so count them as invalid before filtering parent indices. + values_has_nan = pc.fill_null(values_has_nan, True) values_indices = pc.list_parent_indices(arr) if arr.null_count: # Match flatten(), which excludes child spans under null parent lists. @@ -6318,7 +6306,7 @@ class AsyncTable: "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 6ccba606f..7d92b0e6d 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -2170,6 +2170,71 @@ def test_add_with_nans(mem_db: DBConnection): assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0])) +@pytest.mark.parametrize("operation", ["create", "append", "overwrite"]) +@pytest.mark.parametrize("input_type", ["pandas", "rows", "arrow"]) +@pytest.mark.parametrize("on_bad_vectors", [None, "error", "drop", "fill", "null"]) +def test_write_vectors_with_null_elements( + mem_db: DBConnection, operation, input_type, on_bad_vectors +): + schema = pa.schema({"vector": pa.list_(pa.float32(), 2), "id": pa.int64()}) + vectors = [[3.0, 4.0], [np.nan, 1.0], [np.nan, np.nan]] + if input_type == "pandas": + pd = pytest.importorskip("pandas") + + data = pd.DataFrame( + { + "vector": [np.array(v, dtype=np.float32) for v in vectors], + "id": [0, 1, 2], + } + ) + else: + # Pandas turns the NaNs above into null child elements. Explicit nulls + # in other input formats must follow the same bad-vector policy. + data = [ + {"vector": [3.0, 4.0], "id": 0}, + {"vector": [None, 1.0], "id": 1}, + {"vector": [None, None], "id": 2}, + ] + if input_type == "arrow": + data = pa.Table.from_pylist(data, schema=schema) + + options = {} if on_bad_vectors is None else {"on_bad_vectors": on_bad_vectors} + if on_bad_vectors == "fill": + options["fill_value"] = 42.0 + + table = None + if operation != "create": + table = mem_db.create_table("test", schema=schema) + + def write(): + if operation == "create": + return mem_db.create_table("test", data=data, schema=schema, **options) + table.add(data, mode=operation, **options) + return table + + if on_bad_vectors in (None, "error"): + # Create/overwrite can wrap errors from lazy Arrow readers in RuntimeError. + expected_error = ( + ValueError if operation == "append" else (ValueError, RuntimeError) + ) + with pytest.raises(expected_error, match="NaN|null"): + write() + if table is not None: + assert table.count_rows() == 0 + return + + table = write() + if on_bad_vectors == "drop": + expected = [{"vector": [3.0, 4.0], "id": 0}] + else: + expected = [ + {"vector": [3.0, 4.0], "id": 0}, + {"vector": [42.0, 1.0] if on_bad_vectors == "fill" else None, "id": 1}, + {"vector": [42.0, 42.0] if on_bad_vectors == "fill" else None, "id": 2}, + ] + assert table.to_arrow().to_pylist() == expected + + def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection): class Schema(LanceModel): text: str diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 4fb16bb31..e85afb5bc 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -406,7 +406,7 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): ).read_all() if on_bad_vectors == "drop": - expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]]) + expected = pa.array([[1.0, 2.0], [4.0, 5.0]]) elif on_bad_vectors == "fill": expected = pa.array( [ @@ -414,12 +414,12 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): [3.0, 42.0], [4.0, 5.0], [6.0, 7.0], - [None, 9.0], + [42.0, 9.0], [42.0, 42.0], ] ) elif on_bad_vectors == "null": - expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None]) + expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, None, None]) assert output["vector"].combine_chunks() == expected @@ -430,17 +430,17 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): ( pa.list_(pa.float64()), [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]], - [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]], + [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]], ), ( pa.large_list(pa.float64()), [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]], - [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]], + [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]], ), ( pa.list_(pa.float64(), 2), [[1.0, float("nan")], None, [None, 3.0]], - [[1.0, 42.0], [42.0, 42.0], [None, 3.0]], + [[1.0, 42.0], [42.0, 42.0], [42.0, 3.0]], ), ], ) @@ -467,7 +467,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors): ).read_all() output = exception_output(e) assert output == ( - "ValueError: Vector column 'vector' has NaNs. Set " + "ValueError: Vector column 'vector' has NaNs or null values. Set " "on_bad_vectors='drop' to remove them, set on_bad_vectors='fill' " "and fill_value= to replace them, or set on_bad_vectors='null' " "to replace them with null." @@ -490,6 +490,46 @@ def test_handle_bad_vectors_nan(on_bad_vectors): assert output["vector"].combine_chunks() == expected +@pytest.mark.parametrize("value_type", [pa.float16(), pa.float32(), pa.float64()]) +@pytest.mark.parametrize( + "list_type", [pa.list_, pa.large_list, lambda t: pa.list_(t, 2)] +) +@pytest.mark.parametrize("chunked", [False, True]) +def test_has_nan_values_with_null_elements(value_type, list_type, chunked): + arr = pa.array( + [[None, None], [1.0, 2.0], [None, 3.0], None, [4.0, float("nan")], [5.0, 6.0]], + type=list_type(value_type), + ).slice(1, 5) + if chunked: + arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2)]) + # Whole null vectors have no components to inspect. Their handling is + # separate from missing or NaN components in a valid vector. + assert has_nan_values(arr).to_pylist() == [False, True, False, True, False] + + +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +def test_sanitize_vectors_with_only_null_elements(on_bad_vectors): + data = pa.table({"vector": [[None, None]]}) + schema = pa.schema({"vector": pa.list_(pa.float32(), 2)}) + + def sanitize(): + return _sanitize_data( + data, schema, on_bad_vectors=on_bad_vectors, fill_value=42.0 + ).read_all() + + if on_bad_vectors == "error": + with pytest.raises(ValueError, match="null"): + sanitize() + return + + expected = { + "drop": [], + "fill": [{"vector": [42.0, 42.0]}], + "null": [{"vector": None}], + } + assert sanitize().to_pylist() == expected[on_bad_vectors] + + @pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()]) @pytest.mark.parametrize("large", [False, True]) @pytest.mark.parametrize("chunked", [False, True]) diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index 12404ec73..445f29a18 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -45,10 +45,12 @@ pub struct AddResult { #[derive(Debug, Default, Clone, Copy)] pub enum NaNVectorBehavior { - /// Reject any vectors containing NaN values (the default) + /// Reject any vectors containing NaN values or null elements (the default). + /// Whole null vectors are allowed. #[default] Error, - /// Allow NaN values to be added, but they will not be indexed for search + /// Allow NaN values and null elements to be added, but they will not be indexed + /// for search. Keep, } @@ -104,11 +106,12 @@ impl AddDataBuilder { self } - /// Configure how to handle NaN values in vector columns. + /// Configure how to handle NaN values and null elements in vector columns. /// - /// By default, any vectors containing NaN values will be rejected with an - /// error, since NaNs cannot be indexed for search. Setting this to `Keep` - /// will allow NaN values to be added to the table, but they will not be + /// By default, any vectors containing NaN values or null elements will be + /// rejected with an error, since these vectors cannot be indexed for search. + /// Whole null vectors are allowed. Setting this to `Keep` will allow NaN values + /// and null elements to be added to the table, but these vectors will not be /// indexed and will not be searchable. pub fn on_nan_vectors(mut self, behavior: NaNVectorBehavior) -> Self { self.on_nan_vectors = behavior; @@ -953,8 +956,11 @@ mod tests { ); } + #[rstest::rstest] + #[case::nan(Some(f32::NAN))] + #[case::null(None)] #[tokio::test] - async fn test_add_rejects_nan_vectors() { + async fn test_add_rejects_nan_or_null_vectors(#[case] bad_value: Option) { let schema = Arc::new(Schema::new(vec![Field::new( "embedding", DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), @@ -974,7 +980,12 @@ mod tests { FixedSizeListArray::try_new( Arc::new(Field::new("item", DataType::Float32, true)), 4, - Arc::new(Float32Array::from(vec![0.1, 0.2, f32::NAN, 0.4])), + Arc::new(Float32Array::from(vec![ + Some(0.1), + Some(0.2), + bad_value, + Some(0.4), + ])), None, ) .unwrap(), @@ -984,8 +995,8 @@ mod tests { let res = table.add(batch.clone()).execute().await; let err = res.unwrap_err(); assert!( - err.to_string().contains("NaN"), - "Expected error mentioning NaN values, but got: {err:?}" + err.to_string().contains("NaN or null"), + "Expected error mentioning NaN or null values, but got: {err:?}" ); table diff --git a/rust/lancedb/src/table/datafusion/reject_nan.rs b/rust/lancedb/src/table/datafusion/reject_nan.rs index 084c9722b..462216929 100644 --- a/rust/lancedb/src/table/datafusion/reject_nan.rs +++ b/rust/lancedb/src/table/datafusion/reject_nan.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! A DataFusion projection that rejects vectors containing NaN values. +//! A DataFusion projection that rejects vectors containing NaN or null elements. use std::sync::{Arc, LazyLock}; @@ -31,10 +31,12 @@ fn is_vector_field(field: &Field) -> bool { } } -/// Wraps the input plan with a projection that checks vector columns for NaN values. +/// Wraps the input plan with a projection that checks vector columns for NaN or +/// null elements. /// /// Non-vector columns pass through unchanged. Vector columns are wrapped with a -/// UDF that returns the column as-is if no NaNs are present, or errors otherwise. +/// UDF that returns the column as-is if no NaN or null elements are present, or +/// errors otherwise. Whole null vectors pass through unchanged. pub fn reject_nan_vectors(input: Arc) -> Result> { let schema = input.schema(); let config = Arc::new(ConfigOptions::default()); @@ -70,7 +72,7 @@ pub fn reject_nan_vectors(input: Arc) -> Result datafusion_common::Result<()> { ) })?; - // Only inspect elements that are both in a valid parent row and non-null - // themselves. Values backing null parent rows or null child elements may - // contain garbage (including NaN) per the Arrow spec. - let has_nan = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| { + // Values backing null parent rows may contain garbage per the Arrow spec. + // A null child in a valid row is an invalid vector component, however: + // pandas can represent NaN components as Arrow nulls. + let has_bad_values = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| { let row = fsl.value(i); + if row.null_count() > 0 { + return true; + } match row.data_type() { DataType::Float16 => row .as_any() @@ -150,10 +155,10 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> { } }); - if has_nan { + if has_bad_values { return Err(datafusion_common::DataFusionError::ArrowError( Box::new(arrow_schema::ArrowError::ComputeError( - "Vector column contains NaN values".to_string(), + "Vector column contains NaN or null values".to_string(), )), None, )); @@ -165,7 +170,7 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> { #[cfg(test)] mod tests { use super::*; - use arrow_array::Float32Array; + use arrow_array::{Float16Array, Float32Array, Float64Array}; #[test] fn test_passes_clean_vectors() { @@ -209,23 +214,43 @@ mod tests { } #[test] - fn test_skips_null_elements_within_valid_row() { - // A valid row with null child elements: the underlying buffer may hold - // NaN but the null bitmap says they're absent — should not reject. - let values = Float32Array::from(vec![ - Some(1.0), - None, // null element — buffer may contain NaN - Some(3.0), - None, // null element - ]); + fn test_rejects_null_elements_within_valid_row() { + let arrays: Vec> = vec![ + Arc::new(Float16Array::from(vec![Some(half::f16::ONE), None])), + Arc::new(Float32Array::from(vec![Some(1.0), None])), + Arc::new(Float64Array::from(vec![Some(1.0), None])), + ]; + for values in arrays { + let fsl = FixedSizeListArray::try_new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + 2, + values, + None, + ) + .unwrap(); + let err = check_no_nans(&fsl).unwrap_err(); + assert!(err.to_string().contains("null")); + } + } + + #[test] + fn test_skips_null_elements_in_null_rows_and_outside_slice() { let fsl = FixedSizeListArray::try_new( Arc::new(Field::new("item", DataType::Float32, true)), 2, - Arc::new(values), - None, // both rows are valid + Arc::new(Float32Array::from(vec![ + None, + Some(1.0), + Some(2.0), + Some(3.0), + None, + None, + ])), + Some(vec![true, true, false].into()), ) .unwrap(); - assert!(check_no_nans(&fsl).is_ok()); + assert!(check_no_nans(&fsl.slice(1, 2)).is_ok()); + assert!(check_no_nans(&fsl.slice(0, 2)).is_err()); } #[test] From 22fc2a2087c1ed78033eadb056b1aef63abd1b69 Mon Sep 17 00:00:00 2001 From: devjoinedthechat Date: Mon, 26 Oct 2026 11:17:14 +0100 Subject: [PATCH 31/36] fix(python)!: put OpenCLIP models in eval mode before embedding (#4536) `OpenClipEmbeddings` loads its model with `open_clip.create_model_and_transforms()`, which hands the model back in training mode (the open_clip README calls `model.eval()` right after it for that reason). `OpenClipEmbeddings` never switches it out of training mode, so layers that behave differently during training stay active while we embed: stochastic depth in the timm-based image towers (ViTamin, ConvNeXt, MobileCLIP2-L-14), dropout in the Hugging Face text towers (`xlm-roberta-*`, `roberta-ViT-B-32`, `nllb-clip-*`), and batch norm in the ResNet towers loaded with non-OpenAI weights. The `torch.no_grad()` blocks turn off autograd but don't change how these layers behave. The default `ViT-B-32` has none of those layers, which is probably why nobody has run into this. With dropout or stochastic depth active, the same input gives a different vector on every call, so the vector stored at ingestion and the one computed for an identical query don't match. With batch norm in training mode, an image's vector depends on the other images in its batch: ```python import numpy as np from PIL import Image from lancedb.embeddings import get_registry func = get_registry().get("open-clip").create(name="ViTamin-S", pretrained="datacomp1b") image = Image.new("RGB", (256, 256), "red") a = func.generate_image_embedding(image) b = func.generate_image_embedding(image) print(func._model.training, np.array_equal(a, b)) # before: True False # after: False True ``` The fix calls `model.eval()` once after loading, same as the SigLIP, ColPali and ImageBind embedding functions already do. The new test swaps in a stand-in for the `open_clip` module, so it runs without torch or a model download, and checks that the model handed back by `create_model_and_transforms` ends up in eval mode. BREAKING CHANGE: Embeddings change for OpenCLIP models with training-sensitive layers, including ViTamin, ConvNeXt, MobileCLIP2, Hugging Face RoBERTa/XLM-RoBERTa/NLLB text towers, and ResNet models loaded with non-OpenAI weights. For tables embedded with affected models, regenerate the stored embeddings in eval mode and rebuild their vector indexes before querying with the new embeddings. Rebuilding an index alone does not regenerate embeddings. The default ViT-B-32 model is unaffected. Co-authored-by: devjoinedthechat <235875317+devjoinedthechat@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/embeddings/open_clip.py | 3 ++ python/python/tests/test_embeddings.py | 33 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/python/python/lancedb/embeddings/open_clip.py b/python/python/lancedb/embeddings/open_clip.py index 8b3f0eb06..6199ef271 100644 --- a/python/python/lancedb/embeddings/open_clip.py +++ b/python/python/lancedb/embeddings/open_clip.py @@ -49,6 +49,9 @@ class OpenClipEmbeddings(EmbeddingFunction): self.name, pretrained=self.pretrained ) model.to(self.device) + # open_clip returns the model in training mode, which keeps dropout, + # stochastic depth and batch norm in their training behaviour. + model.eval() self._model, self._preprocess = model, preprocess self._tokenizer = open_clip.get_tokenizer(self.name) self._ndims = None diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 207a3d253..3c968a821 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -1039,6 +1039,39 @@ def test_open_clip_opens_percent_encoded_file_uri(tmp_path): assert image.size == (4, 4) +def test_open_clip_model_in_eval_mode(): + """ + open_clip.create_model_and_transforms returns the model in training mode, + which leaves dropout, stochastic depth and batch norm in their training + behaviour. OpenCLIP must switch the model to eval mode before embedding. + """ + from lancedb.embeddings.open_clip import OpenClipEmbeddings + + class FakeModel: + training = True + + def to(self, device): + return self + + def eval(self): + self.training = False + return self + + open_clip = MagicMock() + open_clip.create_model_and_transforms.return_value = (FakeModel(), None, None) + + with patch( + "lancedb.embeddings.open_clip.attempt_import_or_raise", + return_value=open_clip, + ): + func = OpenClipEmbeddings(name="ViTamin-B", pretrained="datacomp1b") + + open_clip.create_model_and_transforms.assert_called_once_with( + "ViTamin-B", pretrained="datacomp1b" + ) + assert func._model.training is False + + def test_jina_generate_image_input_dict_local_path(tmp_path): """ JinaEmbeddings._generate_image_input_dict must accept a local image path From 66357275bb3cbf55e565113e3d8966cd0f7a291c Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 18:57:25 +0800 Subject: [PATCH 32/36] fix: inherit Lance defaults for IVF RQ quantization (#4585) `IvfRq()` forced `num_bits=1` through Python's Rust binding, overriding Lance's five-bit default. The shared native index creation path also substituted one bit when the Rust or Node builder left the value unset. Make Python's bit count optional and forward only explicit values. Use `RQBuildParams::default()` for an unset native bit count, and omit unset values when dispatching to pylance for accelerator builds. Update the Python, Rust, Node, and vector-index documentation. New indexes with an unset bit count now inherit Lance's default, currently five bits, across Python, Rust, and Node. Five-bit vector codes use roughly five times the storage of one-bit codes, plus other index overhead. This also changes reader compatibility: indexes created with `num_bits >= 2` use a newer on-disk layout that some older LanceDB versions cannot read. Upgrade readers to versions that support multi-bit IVF_RQ, or explicitly set `num_bits=1` (for example, `IvfRq(num_bits=1)` in Python) when sharing indexes with readers limited to one-bit IVF_RQ. Validation: - Reproduced six regression failures before the fix: one-bit default metadata, rejected `None`, and incorrect remote/accelerator forwarding. - After merging `main` through `b762ec22`, rebuilt both native bindings against Lance 15.0.0-beta.1. Python validation passed with 106 tests and 15 optional Ollama skips. Coverage includes the full index suite, omitted and explicit bit counts, reading Lance-created IVF_RQ indexes, accelerator forwarding, CUDA-unavailability handling, timeout validation, remote requests, and table/namespace listing behavior. - Rust formatting and Python package Ruff lint/format checks pass. Node `pnpm lint`, `pnpm build` (including TypeScript checks), `pnpm run docs`, the full connection suite, and the existing IVF_RQ regression pass. Generated documentation adds no further changes. The merge preserves upstream's explicit RaBitQ heading anchor alongside the `IvfRq()` example and five-bit default/reader-compatibility guidance. Repository-wide `ruff check .` still reports 49 existing errors in unchanged documentation examples under `docs/web-tests`; root formatting also finds 11 unchanged documentation example files. The Python package checks pass. Fixes #4582 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- docs/src/js/interfaces/IvfRqOptions.md | 3 ++- docs/web/indexing/vector-index.mdx | 17 +++++++++++------ nodejs/lancedb/indices.ts | 3 ++- python/python/lancedb/index.py | 8 ++++++-- python/python/lancedb/table.py | 7 ++++++- python/python/tests/test_index.py | 14 ++++++++++++-- python/python/tests/test_remote_db.py | 11 +++++++++-- python/python/tests/test_table.py | 22 ++++++++++++++++++++++ python/src/index.rs | 8 +++++--- rust/lancedb/src/index/vector.rs | 3 +++ rust/lancedb/src/table/create_index.rs | 5 ++++- 11 files changed, 82 insertions(+), 19 deletions(-) diff --git a/docs/src/js/interfaces/IvfRqOptions.md b/docs/src/js/interfaces/IvfRqOptions.md index 6b48cf653..298565a62 100644 --- a/docs/src/js/interfaces/IvfRqOptions.md +++ b/docs/src/js/interfaces/IvfRqOptions.md @@ -58,7 +58,8 @@ Number of bits per dimension for residual quantization. This value controls how much each residual component is compressed. The more bits, the more accurate the index will be but the slower search. Typical values -are small integers; the default is 1 bit per dimension. +are small integers. If unset, uses Lance's default (currently 5 bits per +dimension). Set to 1 for maximum compression. *** diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index b26e0c295..037d75f8d 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -58,7 +58,7 @@ Lance with `index_type`. | If your priority is... | Use this index | Why | Approx. compression ratio | Python config class | | :--- | :--- | :--- | :--- | :--- | -| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size | `IvfRq` | +| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size with `num_bits=1` | `IvfRq` | | Higher accuracy at small dimensions (`dimension <= 256`) | `IVF_PQ` | IVF indexing with product quantization | Usually `1/64` to `1/16` of raw size (depends on `num_sub_vectors`) | `IvfPq` | | Best recall/latency trade-off | `IVF_HNSW_SQ` | IVF-HNSW indexing with scalar | Typically a little larger than `1/4` of raw size | `IvfHnswSq` | | Highest recall / no quantization | `IVF_HNSW_FLAT` | IVF-HNSW indexing with no quantization; embedded library only | Around raw vector size plus HNSW graph overhead | `IvfHnswFlat` | @@ -377,7 +377,7 @@ which has been compressed to $4$ chunks $\times 8$-bit integers $= 32$ bits of q RaBitQ is an advanced quantization technique that outperforms PQ in several ways. It needs no codebook to train, estimates distances very quickly at query-time, -and crucially, quantizes each vector in (with some small overhead) just **one bit per dimension!** +and can quantize each vector in (with some small overhead) as little as **one bit per dimension!** In practice, RaBitQ compresses a $1024$-dimensional `float32` vector into just a few thousand bits, while maintaining good recall. The inner workings of RaBitQ quantization are rather mathematically dense. It generates a quantization codebook by @@ -385,14 +385,19 @@ applying a uniformly random, approximately distance-preserving orthogonal transf where $d$ is the dimensionality of the dataset. We defer the details, and an elegant theoretical error bound, to [the original paper.](https://arxiv.org/pdf/2405.12497) ##### Using RaBitQ {#using-rabitq} -Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, call `create_index` with `index_type="IVF_RQ"`. +Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, import `IvfRq` from `lancedb.index` and call `table.create_index("vector", config=IvfRq())`. Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`. `num_bits` determines how many bits are used to quantize each dimension. -`1` is the standard RaBitQ setting. Increase to `2`, `4`, or `8` bits to achieve better recall for additional storage and query-time compute. +By default, Python, Rust, and Node leave this unset and inherit Lance's default, currently `5` bits. +Higher values improve recall at the cost of additional storage and query-time compute. +To retain the previous one-bit default and maximum compression, explicitly set `num_bits=1` (for example, `IvfRq(num_bits=1)` in Python). +Five-bit quantized vector codes use roughly five times the storage of one-bit codes; total index size also includes other overhead. -RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout, and cannot be read by some older LanceDB versions. +RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout and cannot be read by some older LanceDB versions. +The new five-bit default therefore affects reader compatibility as well as index size. +Upgrade readers to versions that support multi-bit IVF_RQ before creating indexes with the default settings, or explicitly set `num_bits=1` when sharing indexes with readers that support only one-bit IVF_RQ. See this [blog post](https://lancedb.com/blog/feature-rabitq-quantization/) for further discussion and benchmarking of LanceDB's RaBitQ implementation. @@ -410,7 +415,7 @@ among all integral multiples $\{j \times 10^5: -128 \leq j \leq 127 \}$. | Parameter | Description | | :--- | :--- | -| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (default `1`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | +| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (`IvfRq()` inherits Lance's default, currently `5`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | `max_iterations` and `sample_rate` also affect quantizer training, but since they apply to every IVF/HNSW index type (not just quantized ones), they're documented as general [Build-time Parameters](/indexing/vector-index#build-time-parameters) instead. All three are passed as keyword arguments to `create_index`, alongside `index_type`: diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index 43d1cf1b5..0e383e984 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -131,7 +131,8 @@ export interface IvfRqOptions { * * This value controls how much each residual component is compressed. The more * bits, the more accurate the index will be but the slower search. Typical values - * are small integers; the default is 1 bit per dimension. + * are small integers. If unset, uses Lance's default (currently 5 bits per + * dimension). Set to 1 for maximum compression. */ numBits?: number; diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index cf3bf8e4f..80179f4a7 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -882,8 +882,12 @@ class IvfRq: num_partitions: int, default sqrt(num_rows) Number of IVF partitions to create. - num_bits: int, default 1 + num_bits: int, optional Number of bits to encode each dimension in the RabitQ codebook. + If None, uses Lance's default (currently 5), matching Rust and Node. + Set to 1 to retain the previous Python default and maximum compression. + Five-bit vector codes use roughly five times the storage of one-bit codes, + in addition to other index overhead. max_iterations: int, default 50 Max iterations to train kmeans when computing IVF partitions. @@ -897,7 +901,7 @@ class IvfRq: distance_type: Literal["l2", "cosine", "dot"] = "l2" num_partitions: Optional[int] = None - num_bits: int = 1 + num_bits: Optional[int] = None max_iterations: int = 50 sample_rate: int = 256 target_partition_size: Optional[int] = None diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 2ad7f821d..cc455f746 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3754,6 +3754,11 @@ class LanceTable(Table): cfg_type = type(config).__name__ lance_index_type = index_type_map.get(cfg_type, "IVF_PQ") + quantization_kwargs = {} + num_bits = getattr(config, "num_bits", 8) + if num_bits is not None: + quantization_kwargs["num_bits"] = num_bits + dataset.create_index( column=column, index_type=lance_index_type, @@ -3762,12 +3767,12 @@ class LanceTable(Table): num_sub_vectors=getattr(config, "num_sub_vectors", None), replace=replace, accelerator=acc, - num_bits=getattr(config, "num_bits", 8), m=getattr(config, "m", 20), ef_construction=getattr(config, "ef_construction", 300), target_partition_size=getattr( config, "target_partition_size", None ), + **quantization_kwargs, ) self.checkout_latest() return diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 0252a09ad..2ec14f6f8 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -421,13 +421,23 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable): @pytest.mark.asyncio -async def test_create_ivfrq_index(some_table: AsyncTable): - await some_table.create_index("vector", config=IvfRq(num_bits=1)) +@pytest.mark.parametrize( + "config, expected_bits", + [ + pytest.param(IvfRq(), 5, id="default"), + pytest.param(IvfRq(num_bits=None), 5, id="none"), + pytest.param(IvfRq(num_bits=1), 1, id="one-bit"), + pytest.param(IvfRq(num_bits=4), 4, id="multi-bit"), + ], +) +async def test_create_ivfrq_index(some_table: AsyncTable, config, expected_bits): + await some_table.create_index("vector", config=config) indices = await some_table.list_indices() assert len(indices) == 1 assert indices[0].index_type == "IvfRq" assert indices[0].columns == ["vector"] assert indices[0].name == "vector_idx" + assert indices[0].index_details["compression"]["num_bits"] == expected_bits @pytest.mark.asyncio diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 95da3eba2..049a5fc58 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1593,7 +1593,10 @@ def test_remote_json_fts_config(server_version): assert request["block_size"] == 256 -def test_remote_create_index_new_api(caplog): +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +def test_remote_create_index_new_api(caplog, rq_kwargs): received_requests = [] def handler(request): @@ -1666,7 +1669,7 @@ def test_remote_create_index_new_api(caplog): table.create_index("category", config=BTree()) table.create_index("text", config=FTS(block_size=256)) # IvfRq via new API - table.create_index("vector", config=IvfRq(distance_type="l2")) + table.create_index("vector", config=IvfRq(distance_type="l2", **rq_kwargs)) table.create_index( "vector", config=IvfPq(distance_type="l2"), replace=False ) @@ -1691,6 +1694,10 @@ def test_remote_create_index_new_api(caplog): "vector", ] assert received_requests[2]["block_size"] == 256 + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in received_requests[3] + else: + assert received_requests[3]["num_bits"] == rq_kwargs["num_bits"] assert received_requests[4]["replace"] is False assert received_requests[5]["replace"] is False assert "replace is not supported" not in caplog.text diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 7d92b0e6d..3b3ccd122 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1751,6 +1751,28 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection): job.cancel() +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +@patch("lancedb.table.LanceTable.to_lance") +def test_create_ivfrq_index_accelerator(mock_to_lance, mem_db: DBConnection, rq_kwargs): + from lancedb.index import IvfRq + + table = mem_db.create_table("test", [{"vector": [1.0] * 8}]) + torch = MagicMock() + torch.cuda.is_available.return_value = True + with patch.dict(sys.modules, {"torch": torch}): + table.create_index("vector", config=IvfRq(accelerator="cuda", **rq_kwargs)) + + kwargs = mock_to_lance.return_value.create_index.call_args.kwargs + assert kwargs["index_type"] == "IVF_RQ" + assert kwargs["accelerator"] == "cuda" + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in kwargs + else: + assert kwargs["num_bits"] == rq_kwargs["num_bits"] + + @pytest.mark.parametrize("method", ["create_index", "create_index_async"]) @pytest.mark.parametrize("wait_timeout", ["invalid", 1, []]) def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout): diff --git a/python/src/index.rs b/python/src/index.rs index d4e16a67e..c6c3fe17d 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -136,8 +136,10 @@ pub fn extract_index_params(source: &Option>) -> PyResult, - num_bits: u32, + num_bits: Option, max_iterations: u32, sample_rate: u32, target_partition_size: Option, diff --git a/rust/lancedb/src/index/vector.rs b/rust/lancedb/src/index/vector.rs index bce77a610..3e6794b54 100644 --- a/rust/lancedb/src/index/vector.rs +++ b/rust/lancedb/src/index/vector.rs @@ -362,6 +362,9 @@ impl IvfRqIndexBuilder { impl_distance_type_setter!(); impl_ivf_params_setter!(); + /// Set the number of bits used to encode each vector dimension. + /// + /// If unset, uses Lance's default (currently 5). Set to 1 for maximum compression. pub fn num_bits(mut self, num_bits: u32) -> Self { self.num_bits = Some(num_bits); self diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 3ac759bdd..1795b778e 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -355,7 +355,10 @@ impl NativeTable { index.sample_rate, index.max_iterations, ); - let rq_params = RQBuildParams::new(index.num_bits.unwrap_or(1) as u8); + let rq_params = index + .num_bits + .map(|num_bits| RQBuildParams::new(num_bits as u8)) + .unwrap_or_default(); let lance_idx_params = VectorIndexParams::with_ivf_rq_params( index.distance_type.into(), ivf_params, From 3dd579785edf209f0e410b339be89ba7cc71757d Mon Sep 17 00:00:00 2001 From: KarthikS Date: Mon, 26 Oct 2026 07:28:44 -0400 Subject: [PATCH 33/36] fix(python)!: align synchronous table method signatures (#4498) Closes #4396. Calls written against `Table` can fail or bind different arguments when used with `RemoteTable`. Align the affected sync signatures, positional order, and defaults; preserve renamed remote keywords as deprecated aliases. Share legacy vector configuration and tokenizer inference, forward index tuning and timeout options, and accept `fast_search` consistently. Preserve current upstream conversion behavior: local `to_polars(batch_size=...)` returns a LazyFrame; remote tables accept the same argument signature but raise the documented unsupported-operation error and direct callers to `table.search().to_polars()`. The abstract `Table` method retains its unsupported-operation error, with the implementation confined to `LanceTable`. Regression coverage checks argument binding, deprecated aliases, option forwarding, native behavior through both wrappers, HTTP index requests, and remote conversion rejection. The merge preserves current tokenizer defaults, remote query defaults, HNSW-PQ support, and unsupported-operation handling. Runtime validation at `c4fb01f2` (Linux x86_64, uv-managed Python 3.12.15): - Focused Ruff lint and formatting passed for all four changed Python files. - Rebuilt and installed the current native extension successfully with the repository's `ci` Cargo profile. - Seven regression suites passed: `test_errors.py`, `test_table_signatures.py`, `test_remote_db.py`, `test_table.py`, `test_index.py`, `test_fts.py`, and `test_query.py`: **874 passed, 1 deselected**, with slow and S3 tests excluded. - CI caught an abstract-base conversion regression in the initial merge. The existing error-contract test reproduced it locally (1 failed, 7 passed); restoring the implementation to `LanceTable` fixes it while preserving signature parity. - Full-root Ruff reports 49 lint errors and 11 files needing formatting, all in unchanged upstream `docs/web-tests` files. The four changed Python files pass both checks. - Validated against upstream `197403e8`, using Lance 15.0.0-beta.1 and LanceDB 0.41.0-beta.14 bindings. Historical validation before this merge: 618 tests passed and one was deselected. That result does not establish correctness of the current merge. The original client-side reproduction produced nine argument-binding errors without contacting a server; the initial regression suite failed 32 cases before the fix and passed 39 after it. Documentation review follow-up: restored the `batch_size` parameter description on the abstract `Table.to_polars` method so it appears in the public API reference. This adds six docstring lines and changes no runtime behavior. After this edit, all 47 focused error-contract and signature tests pass; changed-file Ruff lint/format and `git diff --check` pass. Repository-wide Ruff still reports the same 49 findings and 11 formatting files in unchanged `docs/web-tests`. The full documentation site and seven-suite runtime run were not repeated for this documentation-only follow-up. The reported ARM64 production-server deployment has not been recreated. HTTP integration tests use the repository's local mock server. BREAKING CHANGE: Remote synchronous table methods now use the positional argument order and defaults documented by `Table`. Migrate calls using the former remote-only positional order to named arguments, particularly `create_index` and `search`; pass `index_type` by keyword to `create_scalar_index`. Remote scalar and vector index creation now default to `replace=True`; specify `replace=False` to retain rejection of an existing index. Renamed remote keywords remain supported as deprecated aliases, as documented above. --------- Co-authored-by: Xuanwo --- docs/src/python/python.md | 19 + python/python/lancedb/remote/table.py | 343 ++++++++------- python/python/lancedb/table.py | 433 ++++++++++--------- python/python/tests/test_remote_db.py | 31 +- python/python/tests/test_table_signatures.py | 253 +++++++++++ 5 files changed, 733 insertions(+), 346 deletions(-) create mode 100644 python/python/tests/test_table_signatures.py diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 02ccced2c..4d0dfc87a 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -297,6 +297,25 @@ listing a storage directory. ## Tables (Synchronous) +Local and remote tables use the same arguments for index creation, deletion, +index statistics, and search. For example, `table.delete(where="id = 1")`, +`table.create_index("cosine", 2)` (the deprecated vector API), and +`table.search([1.0, 1.0], fast_search=True)` work with either backend. +`fast_search=True` excludes rows that have not been indexed. + +`create_scalar_index` defaults to `replace=True` on both backends. Both scalar +and full-text index creation accept `wait_timeout`. Remote maintenance methods +accept the local options: cleanup and compaction warn that they are no-ops, +while `optimize` reports an unsupported operation. Removed Tantivy options +produce the same explanatory errors on both backends. + +The old remote keyword names `predicate` for `delete`, `index_name` for +`drop_index`, `index_uuid` for `index_stats`, and `column` for `create_fts_index` +remain available with deprecation warnings. Use `where`, `name`, `index_name`, +and `field_names`, respectively. For positional calls, use the argument order +documented by `Table` below; specify index options by keyword when migrating +code that used the former remote-only argument order. + ::: lancedb.table.Table ::: lancedb.table.FragmentStatistics diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index dd22dea48..faf15c987 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -4,7 +4,7 @@ from datetime import timedelta import deprecation import logging -from functools import cached_property +from functools import cached_property, wraps import os from typing import ( TYPE_CHECKING, @@ -42,19 +42,12 @@ from lancedb.index import ( FTS, BTree, Bitmap, - HnswFlat, - HnswPq, - HnswSq, - IvfFlat, - IvfPq, - IvfRq, - IvfSq, LabelList, ) from lancedb.job import Job from lancedb.functions import FunctionApplication, RefreshColumnResult from lancedb.remote.db import LOOP -from lancedb.table import IndexConfigType, KNOWN_METRICS +from lancedb.table import IndexConfigType, _maybe_add_fts_error_note import pyarrow as pa from lancedb.common import DATA, VEC, VECTOR_COLUMN_NAME @@ -69,7 +62,6 @@ from ..query import ( FullTextQuery, LanceQueryBuilder, LanceTakeQueryBuilder, - LanceVectorQueryBuilder, ) from ..table import ( AsyncTable, @@ -80,7 +72,7 @@ from ..table import ( Table, Tags, ) -from ..types import BaseTokenizerType +from ..types import BaseTokenizerType, ScalarIndexType from ..util import infer_vector_column_name if TYPE_CHECKING: @@ -88,6 +80,28 @@ if TYPE_CHECKING: import polars as pl +def _deprecated_keyword(old_name: str, new_name: str): + """Keep old remote-only keyword names without changing the Table signature.""" + + def decorate(method): + @wraps(method) + def wrapped(self, *args, **kwargs): + if old_name in kwargs: + if new_name in kwargs: + raise TypeError(f"Cannot specify both {old_name} and {new_name}") + warnings.warn( + f"{old_name} is deprecated; use {new_name} instead", + DeprecationWarning, + stacklevel=2, + ) + kwargs[new_name] = kwargs.pop(old_name) + return method(self, *args, **kwargs) + + return wrapped + + return decorate + + class RemoteTable(Table): def __init__( self, @@ -268,7 +282,7 @@ class RemoteTable(Table): """ return LOOP.run(self._table.to_lance(**kwargs)) - def to_polars(self, **kwargs) -> "pl.LazyFrame": + def to_polars(self, batch_size=None) -> "pl.LazyFrame": """Use ``table.search().to_polars()`` to read remote query results.""" raise NotImplementedError( "to_polars() is not supported for remote tables. " @@ -315,9 +329,10 @@ class RemoteTable(Table): self._table.tokenize(query, column=column, index_name=index_name) ) - def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]: + @_deprecated_keyword("index_uuid", "index_name") + def index_stats(self, index_name: str) -> Optional[IndexStatistics]: """List all the stats of a specified index""" - return LOOP.run(self._table.index_stats(index_uuid)) + return LOOP.run(self._table.index_stats(index_name)) @deprecation.deprecated( deprecated_in="0.25.0", @@ -327,43 +342,38 @@ class RemoteTable(Table): def create_scalar_index( self, column: str, - index_type: Literal["BTREE", "BITMAP", "LABEL_LIST", "scalar"] = "scalar", *, - replace: bool = False, + replace: bool = True, + index_type: ScalarIndexType = "BTREE", wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): - """Creates a scalar index. + """Create a scalar index on a column. .. deprecated:: 0.25.0 Use :meth:`create_index` with a BTree, Bitmap, or LabelList config instead. Example: ``table.create_index("column", config=BTree())`` - - Parameters - ---------- - column : str - The column to be indexed. Must be a boolean, integer, float, - or string column. - index_type : str - The index type of the scalar index. Must be "scalar" (BTREE), - "BTREE", "BITMAP", or "LABEL_LIST", - replace : bool - If True, replace the existing index with the new one. """ - if index_type == "scalar" or index_type == "BTREE": + if index_type == "scalar": + warnings.warn( + 'index_type="scalar" is deprecated; use "BTREE" instead', + DeprecationWarning, + stacklevel=2, + ) + index_type = "BTREE" + if index_type == "BTREE": config = BTree() elif index_type == "BITMAP": config = Bitmap() elif index_type == "LABEL_LIST": config = LabelList() else: - raise ValueError(f"Unknown index type: {index_type}") - - LOOP.run( + raise ValueError(f"Unknown index type {index_type}") + return LOOP.run( self._table.create_index( column, - config=config, replace=replace, + config=config, wait_timeout=wait_timeout, name=name, ) @@ -374,12 +384,16 @@ class RemoteTable(Table): current_version=__version__, details="Use create_index() with config=FTS() instead.", ) + @_deprecated_keyword("column", "field_names") def create_fts_index( self, - column: str, + field_names: Union[str, List[str]], *, + ordering_field_names: Optional[Union[str, List[str]]] = None, replace: bool = False, - wait_timeout: Optional[timedelta] = None, + writer_heap_size: Optional[int] = 1024 * 1024 * 1024, + use_tantivy: bool = False, + tokenizer_name: Optional[str] = None, with_position: bool = False, # tokenizer configs: base_tokenizer: BaseTokenizerType = "simple", @@ -395,6 +409,7 @@ class RemoteTable(Table): prefix_only: bool = False, block_size: int = 128, document_granularity: DocumentGranularity = DocumentGranularity.ROW, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a full-text search index on a column. @@ -403,31 +418,70 @@ class RemoteTable(Table): Use :meth:`create_index` with an FTS config instead. Example: ``table.create_index("text_column", config=FTS())`` """ + if use_tantivy: + raise ValueError( + "Tantivy-based FTS has been removed. " + "Remove use_tantivy and recreate the index with native FTS." + ) + if ordering_field_names is not None: + raise ValueError( + "ordering_field_names was only supported by the removed " + "Tantivy-based FTS implementation." + ) + if writer_heap_size != 1024 * 1024 * 1024: + raise ValueError( + "writer_heap_size was only supported by the removed " + "Tantivy-based FTS implementation." + ) + if not isinstance(field_names, str): + raise ValueError( + "Native FTS indexes can only be created on a single field " + "at a time. To search over multiple text fields, create a " + "separate FTS index for each field." + ) + + if tokenizer_name is None: + tokenizer_configs = { + "base_tokenizer": base_tokenizer, + "language": language, + "with_position": with_position, + "max_token_length": max_token_length, + "lower_case": lower_case, + "stem": stem, + "remove_stop_words": remove_stop_words, + "custom_stop_words": custom_stop_words, + "ascii_folding": ascii_folding, + "ngram_min_length": ngram_min_length, + "ngram_max_length": ngram_max_length, + "prefix_only": prefix_only, + } + else: + tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name) + tokenizer_configs["custom_stop_words"] = custom_stop_words + config = FTS( - with_position=with_position, - base_tokenizer=base_tokenizer, - language=language, - max_token_length=max_token_length, - lower_case=lower_case, - stem=stem, - remove_stop_words=remove_stop_words, - custom_stop_words=custom_stop_words, - ascii_folding=ascii_folding, - ngram_min_length=ngram_min_length, - ngram_max_length=ngram_max_length, - prefix_only=prefix_only, block_size=block_size, document_granularity=document_granularity, + **tokenizer_configs, ) - LOOP.run( - self._table.create_index( - column, - config=config, - replace=replace, - wait_timeout=wait_timeout, - name=name, + + try: + LOOP.run( + self._table.create_index( + field_names, + replace=replace, + config=config, + name=name, + wait_timeout=wait_timeout, + ) ) - ) + except (ValueError, RuntimeError) as e: + _maybe_add_fts_error_note( + e, + base_tokenizer=config.base_tokenizer, + language=config.language, + ) + raise e # New unified API overload @overload @@ -437,6 +491,7 @@ class RemoteTable(Table): /, *, config: IndexConfigType, + replace: bool = ..., wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., @@ -447,38 +502,56 @@ class RemoteTable(Table): def create_index( self, metric: Literal["l2", "cosine", "dot", "hamming"] = ..., - vector_column_name: str = ..., - index_cache_size: Optional[int] = ..., num_partitions: Optional[int] = ..., num_sub_vectors: Optional[int] = ..., - replace: Optional[bool] = ..., + vector_column_name: str = ..., + replace: bool = ..., accelerator: Optional[str] = ..., - index_type: Literal[ - "VECTOR", "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" - ] = ..., - wait_timeout: Optional[timedelta] = ..., - *, + index_cache_size: Optional[int] = ..., num_bits: int = ..., + index_type: Literal[ + "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" + ] = ..., + max_iterations: int = ..., + sample_rate: int = ..., + m: int = ..., + ef_construction: int = ..., + *, + wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., + target_partition_size: Optional[int] = ..., ) -> None: ... def create_index( self, metric: str = "l2", - vector_column_name: str = VECTOR_COLUMN_NAME, - index_cache_size: Optional[int] = None, num_partitions: Optional[int] = None, num_sub_vectors: Optional[int] = None, - replace: Optional[bool] = None, + vector_column_name: str = VECTOR_COLUMN_NAME, + replace: bool = True, accelerator: Optional[str] = None, - index_type="vector", - wait_timeout: Optional[timedelta] = None, - *, + index_cache_size: Optional[int] = None, num_bits: int = 8, + index_type: Literal[ + "IVF_FLAT", + "IVF_SQ", + "IVF_PQ", + "IVF_RQ", + "IVF_HNSW_SQ", + "IVF_HNSW_PQ", + "IVF_HNSW_FLAT", + ] = "IVF_PQ", + max_iterations: int = 50, + sample_rate: int = 256, + m: int = 20, + ef_construction: int = 300, + *, config: Optional[IndexConfigType] = None, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, train: bool = True, + target_partition_size: Optional[int] = None, ): """Create an index on a column. @@ -489,6 +562,24 @@ class RemoteTable(Table): argument plus separate ``vector_column_name`` / ``num_partitions`` / etc. parameters, and emits a ``DeprecationWarning``. + Parameters + ---------- + metric : str + For new API: the column name to index. + For legacy API: the distance metric ("l2", "cosine", "dot", "hamming"). + config : IndexConfigType, optional + The index configuration object. If provided, uses the new unified API. + Can be one of: IvfFlat, IvfPq, IvfSq, IvfRq, HnswPq, HnswSq, + BTree, Bitmap, LabelList, Fm, FTS. + replace : bool, default True + Whether to replace an existing index on this column. + wait_timeout : timedelta, optional + Timeout to wait for async indexing to complete. + name : str, optional + Custom name for the index. + train : bool, default True + Whether to train the index with existing data. + Examples -------- New API (recommended): @@ -514,7 +605,6 @@ class RemoteTable(Table): vector_column_name, accelerator, index_cache_size, - replace, ) if is_legacy: @@ -530,6 +620,7 @@ class RemoteTable(Table): stacklevel=2, ) + # Legacy API: first arg is the distance metric column = vector_column_name if accelerator is not None: @@ -538,45 +629,27 @@ class RemoteTable(Table): "and will be ignored." ) - idx_type = index_type.upper() - if idx_type == "VECTOR" or idx_type == "IVF_PQ": - config = IvfPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - ) - elif idx_type == "IVF_RQ": - config = IvfRq( - distance_type=metric, - num_partitions=num_partitions, - num_bits=num_bits, - ) - elif idx_type == "IVF_SQ": - config = IvfSq(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_HNSW_PQ": - config = HnswPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - ) - elif idx_type == "IVF_HNSW_SQ": - config = HnswSq(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_HNSW_FLAT": - config = HnswFlat(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_FLAT": - config = IvfFlat(distance_type=metric, num_partitions=num_partitions) - else: - raise ValueError( - f"Unknown vector index type: {idx_type}. Valid options are" - " 'IVF_FLAT', 'IVF_PQ', 'IVF_RQ', 'IVF_SQ'," - " 'IVF_HNSW_PQ', 'IVF_HNSW_SQ', 'IVF_HNSW_FLAT'" - ) + # Build config from legacy parameters + config = self._build_vector_config_from_legacy_params( + metric=metric, + index_type="IVF_PQ" + if index_type.upper() == "VECTOR" + else index_type.upper(), + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=None, + ) + else: column = metric - LOOP.run( + return LOOP.run( self._table.create_index( column, replace=replace, @@ -615,37 +688,6 @@ class RemoteTable(Table): ) ) - def _is_legacy_create_index_call( - self, - first_arg: str, - config: Optional[IndexConfigType], - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - vector_column_name: str, - accelerator: Optional[str], - index_cache_size: Optional[int], - replace: Optional[bool], - ) -> bool: - """Detect if this is a legacy create_index call.""" - if config is not None: - return False - if any( - x is not None - for x in ( - num_partitions, - num_sub_vectors, - accelerator, - index_cache_size, - replace, - ) - ): - return True - if vector_column_name != VECTOR_COLUMN_NAME: - return True - if first_arg.lower() in KNOWN_METRICS: - return True - return False - def add( self, data: DATA, @@ -717,9 +759,10 @@ class RemoteTable(Table): query: Union[VEC, str] = None, vector_column_name: Optional[str] = None, query_type="auto", + ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, fast_search: Optional[bool] = None, - ) -> LanceVectorQueryBuilder: + ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -810,6 +853,7 @@ class RemoteTable(Table): query, query_type, vector_column_name=vector_column_name, + ordering_field_name=ordering_field_name, fts_columns=fts_columns or [], fast_search=fast_search, ) @@ -869,7 +913,8 @@ class RemoteTable(Table): self._table._do_merge(merge, new_data, on_bad_vectors, fill_value) ) - def delete(self, predicate: str) -> DeleteResult: + @_deprecated_keyword("predicate", "where") + def delete(self, where: Union[str, Expr]) -> DeleteResult: """Delete rows from the table. This can be used to delete a single row, many rows, all rows, or @@ -877,7 +922,7 @@ class RemoteTable(Table): Parameters ---------- - predicate: str + where: Union[str, Expr] The SQL where clause to use when deleting rows. - For example, 'x = 2' or 'x IN (1, 2, 3)'. @@ -921,7 +966,7 @@ class RemoteTable(Table): x vector _distance # doctest: +SKIP 0 2 [3.0, 4.0] 85.0 # doctest: +SKIP """ - return LOOP.run(self._table.delete(predicate)) + return LOOP.run(self._table.delete(where)) def update( self, @@ -983,7 +1028,12 @@ class RemoteTable(Table): self._table.update(where=where, updates=values, updates_sql=values_sql) ) - def cleanup_old_versions(self, *_): + def cleanup_old_versions( + self, + older_than: Optional[timedelta] = None, + *, + delete_unverified: bool = False, + ): """ cleanup_old_versions() is a no-op for remote tables. """ @@ -993,7 +1043,7 @@ class RemoteTable(Table): ) pass - def compact_files(self, *_): + def compact_files(self, *args, **kwargs): """ compact_files() is a no-op for remote tables. """ @@ -1008,12 +1058,14 @@ class RemoteTable(Table): *, cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, + retrain: bool = False, ): """optimize() is not supported for remote tables.""" return LOOP.run( self._table.optimize( cleanup_older_than=cleanup_older_than, delete_unverified=delete_unverified, + retrain=retrain, ) ) @@ -1121,8 +1173,9 @@ class RemoteTable(Table): """No-op for remote tables (no local shard writers).""" return LOOP.run(self._table.close_rte_writers()) - def drop_index(self, index_name: str): - return LOOP.run(self._table.drop_index(index_name)) + @_deprecated_keyword("index_name", "name") + def drop_index(self, name: str) -> None: + return LOOP.run(self._table.drop_index(name)) def prewarm_index(self, name: str) -> None: """Prewarm an index in the table. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index cc455f746..d2ff02683 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -265,10 +265,8 @@ if TYPE_CHECKING: OnBadVectorsType, AddMode, CreateMode, - VectorIndexType, ScalarIndexType, BaseTokenizerType, - DistanceType, ) # Type alias for index configuration objects @@ -1489,12 +1487,18 @@ class Table(ABC): """ raise NotImplementedError("to_lance is not supported for this table type") - def to_polars(self, **kwargs) -> "pl.LazyFrame": + def to_polars(self, batch_size=None) -> "pl.LazyFrame": """Return the table as a Polars LazyFrame. For remote tables, use ``table.search().to_polars()`` to convert query results to a Polars DataFrame. + Parameters + ---------- + batch_size: int, optional + Passed to Polars as the maximum row count for scanned PyArrow + record batches. + Note ---- The Polars streaming engine is not supported because it does not currently @@ -1507,6 +1511,193 @@ class Table(ABC): """ raise NotImplementedError("to_polars is not supported for this table type") + def _is_legacy_create_index_call( + self, + first_arg: str, + config: Optional[IndexConfigType], + num_partitions: Optional[int], + num_sub_vectors: Optional[int], + vector_column_name: str, + accelerator: Optional[str], + index_cache_size: Optional[int], + ) -> bool: + """Detect if this is a legacy create_index call.""" + # If config is provided, it's definitely the new API + if config is not None: + return False + + # If old-style parameters were explicitly set, it's legacy + if any( + x is not None + for x in (num_partitions, num_sub_vectors, accelerator, index_cache_size) + ): + return True + + # If vector_column_name differs from default, it's legacy + if vector_column_name != VECTOR_COLUMN_NAME: + return True + + # If first arg is a known metric, assume legacy + if first_arg.lower() in KNOWN_METRICS: + return True + + # Otherwise assume new API + return False + + def _build_vector_config_from_legacy_params( + self, + metric: str, + index_type: str, + num_partitions: Optional[int], + num_sub_vectors: Optional[int], + num_bits: int, + max_iterations: int, + sample_rate: int, + m: int, + ef_construction: int, + target_partition_size: Optional[int], + accelerator: Optional[str], + ) -> IndexConfigType: + """Build an index config object from legacy parameters.""" + if index_type == "IVF_FLAT": + return IvfFlat( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_SQ": + return IvfSq( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_PQ": + return IvfPq( + distance_type=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_RQ": + return IvfRq( + distance_type=metric, + num_partitions=num_partitions, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_PQ": + return HnswPq( + distance_type=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_SQ": + return HnswSq( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_FLAT": + return HnswFlat( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + ) + else: + raise ValueError(f"Unknown index type {index_type}") + + @staticmethod + def infer_tokenizer_configs(tokenizer_name: str) -> dict: + if tokenizer_name == "default": + return { + "base_tokenizer": "simple", + "language": "English", + "max_token_length": 40, + "lower_case": True, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + elif tokenizer_name == "raw": + return { + "base_tokenizer": "raw", + "language": "English", + "max_token_length": None, + "lower_case": False, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + elif tokenizer_name == "whitespace": + return { + "base_tokenizer": "whitespace", + "language": "English", + "max_token_length": None, + "lower_case": False, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + + # or it's with language stemming with pattern like "en_stem" + if len(tokenizer_name) != 7: + raise ValueError(f"Invalid tokenizer name {tokenizer_name}") + lang = tokenizer_name[:2] + if tokenizer_name[-5:] != "_stem": + raise ValueError(f"Invalid tokenizer name {tokenizer_name}") + if lang not in lang_mapping: + raise ValueError(f"Invalid language code {lang}") + return { + "base_tokenizer": "simple", + "language": lang_mapping[lang], + "max_token_length": 40, + "lower_case": True, + "stem": True, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + # New unified API overload @overload def create_index( @@ -1532,14 +1723,16 @@ class Table(ABC): replace: bool = ..., accelerator: Optional[str] = ..., index_cache_size: Optional[int] = ..., - *, - index_type: VectorIndexType = ..., - wait_timeout: Optional[timedelta] = ..., num_bits: int = ..., + index_type: Literal[ + "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" + ] = ..., max_iterations: int = ..., sample_rate: int = ..., m: int = ..., ef_construction: int = ..., + *, + wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., target_partition_size: Optional[int] = ..., @@ -1547,22 +1740,30 @@ class Table(ABC): def create_index( self, - metric: DistanceType = "l2", + metric: str = "l2", num_partitions: Optional[int] = None, num_sub_vectors: Optional[int] = None, vector_column_name: str = VECTOR_COLUMN_NAME, replace: bool = True, accelerator: Optional[str] = None, index_cache_size: Optional[int] = None, - *, - index_type: VectorIndexType = "IVF_PQ", - wait_timeout: Optional[timedelta] = None, num_bits: int = 8, + index_type: Literal[ + "IVF_FLAT", + "IVF_SQ", + "IVF_PQ", + "IVF_RQ", + "IVF_HNSW_SQ", + "IVF_HNSW_PQ", + "IVF_HNSW_FLAT", + ] = "IVF_PQ", max_iterations: int = 50, sample_rate: int = 256, m: int = 20, ef_construction: int = 300, + *, config: Optional[IndexConfigType] = None, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, train: bool = True, target_partition_size: Optional[int] = None, @@ -2019,6 +2220,7 @@ class Table(ABC): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -3816,130 +4018,6 @@ class LanceTable(Table): ) ) - def _is_legacy_create_index_call( - self, - first_arg: str, - config: Optional[IndexConfigType], - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - vector_column_name: str, - accelerator: Optional[str], - index_cache_size: Optional[int], - ) -> bool: - """Detect if this is a legacy create_index call.""" - # If config is provided, it's definitely the new API - if config is not None: - return False - - # If old-style parameters were explicitly set, it's legacy - if any( - x is not None - for x in (num_partitions, num_sub_vectors, accelerator, index_cache_size) - ): - return True - - # If vector_column_name differs from default, it's legacy - if vector_column_name != VECTOR_COLUMN_NAME: - return True - - # If first arg is a known metric, assume legacy - if first_arg.lower() in KNOWN_METRICS: - return True - - # Otherwise assume new API - return False - - def _build_vector_config_from_legacy_params( - self, - metric: str, - index_type: str, - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - num_bits: int, - max_iterations: int, - sample_rate: int, - m: int, - ef_construction: int, - target_partition_size: Optional[int], - accelerator: Optional[str], - ) -> IndexConfigType: - """Build an index config object from legacy parameters.""" - if index_type == "IVF_FLAT": - return IvfFlat( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_SQ": - return IvfSq( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_PQ": - return IvfPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_RQ": - return IvfRq( - distance_type=metric, - num_partitions=num_partitions, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_PQ": - return HnswPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_SQ": - return HnswSq( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_FLAT": - return HnswFlat( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - ) - else: - raise ValueError(f"Unknown index type {index_type}") - def drop_index(self, name: str) -> None: """ Drops an index from the table @@ -4050,6 +4128,7 @@ class LanceTable(Table): *, replace: bool = True, index_type: ScalarIndexType = "BTREE", + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a scalar index on a column. @@ -4058,6 +4137,13 @@ class LanceTable(Table): Use :meth:`create_index` with a BTree, Bitmap, or LabelList config instead. Example: ``table.create_index("column", config=BTree())`` """ + if index_type == "scalar": + warnings.warn( + 'index_type="scalar" is deprecated; use "BTREE" instead', + DeprecationWarning, + stacklevel=2, + ) + index_type = "BTREE" if index_type == "BTREE": config = BTree() elif index_type == "BITMAP": @@ -4067,7 +4153,13 @@ class LanceTable(Table): else: raise ValueError(f"Unknown index type {index_type}") return LOOP.run( - self._table.create_index(column, replace=replace, config=config, name=name) + self._table.create_index( + column, + replace=replace, + config=config, + wait_timeout=wait_timeout, + name=name, + ) ) @deprecation.deprecated( @@ -4099,6 +4191,7 @@ class LanceTable(Table): prefix_only: bool = False, block_size: int = 128, document_granularity: DocumentGranularity = DocumentGranularity.ROW, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a full-text search index on a column. @@ -4163,6 +4256,7 @@ class LanceTable(Table): replace=replace, config=config, name=name, + wait_timeout=wait_timeout, ) ) except (ValueError, RuntimeError) as e: @@ -4173,69 +4267,6 @@ class LanceTable(Table): ) raise e - @staticmethod - def infer_tokenizer_configs(tokenizer_name: str) -> dict: - if tokenizer_name == "default": - return { - "base_tokenizer": "simple", - "language": "English", - "max_token_length": 40, - "lower_case": True, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - elif tokenizer_name == "raw": - return { - "base_tokenizer": "raw", - "language": "English", - "max_token_length": None, - "lower_case": False, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - elif tokenizer_name == "whitespace": - return { - "base_tokenizer": "whitespace", - "language": "English", - "max_token_length": None, - "lower_case": False, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - - # or it's with language stemming with pattern like "en_stem" - if len(tokenizer_name) != 7: - raise ValueError(f"Invalid tokenizer name {tokenizer_name}") - lang = tokenizer_name[:2] - if tokenizer_name[-5:] != "_stem": - raise ValueError(f"Invalid tokenizer name {tokenizer_name}") - if lang not in lang_mapping: - raise ValueError(f"Invalid language code {lang}") - return { - "base_tokenizer": "simple", - "language": lang_mapping[lang], - "max_token_length": 40, - "lower_case": True, - "stem": True, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - def add( self, data: DATA, @@ -4386,6 +4417,7 @@ class LanceTable(Table): query_type: Literal["vector"] = "vector", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceVectorQueryBuilder: ... @overload @@ -4396,6 +4428,7 @@ class LanceTable(Table): query_type: Literal["fts"] = "fts", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceFtsQueryBuilder: ... @overload @@ -4408,6 +4441,7 @@ class LanceTable(Table): query_type: Literal["hybrid"] = "hybrid", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceHybridQueryBuilder: ... @overload @@ -4418,6 +4452,7 @@ class LanceTable(Table): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceEmptyQueryBuilder: ... def search( @@ -4429,6 +4464,7 @@ class LanceTable(Table): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -4513,6 +4549,7 @@ class LanceTable(Table): vector_column_name=vector_column_name, ordering_field_name=ordering_field_name, fts_columns=fts_columns or [], + fast_search=fast_search, ) @classmethod diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 049a5fc58..db3394735 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1155,8 +1155,9 @@ def test_table_create_indices(): # Test create_fts_index with custom name (legacy method) with pytest.warns(DeprecationWarning, match="create_fts_index"): table.create_fts_index( - "text", + field_names="text", wait_timeout=timedelta(seconds=2), + tokenizer_name="en_stem", block_size=256, custom_stop_words=["cloud"], name="custom_fts_idx", @@ -1165,6 +1166,8 @@ def test_table_create_indices(): # Test create_index with custom name (legacy form: vector_column_name kwarg) with pytest.warns(DeprecationWarning, match="create_index"): table.create_index( + "cosine", + 2, vector_column_name="vector", wait_timeout=timedelta(seconds=10), name="custom_vector_idx", @@ -1177,7 +1180,8 @@ def test_table_create_indices(): scalar_req = received_requests[0] assert "name" in scalar_req assert scalar_req["name"] == "custom_scalar_idx" - assert scalar_req["replace"] is False + # The wire protocol omits replace=True (the server default). + assert "replace" not in scalar_req # Check FTS index request has custom name fts_req = received_requests[1] @@ -1186,18 +1190,39 @@ def test_table_create_indices(): assert fts_req["replace"] is False assert fts_req["block_size"] == 256 assert fts_req["custom_stop_words"] == ["cloud"] + assert fts_req["language"] == "English" + assert fts_req["stem"] is True # Check vector index request has custom name vector_req = received_requests[2] assert "name" in vector_req assert vector_req["name"] == "custom_vector_idx" assert "replace" not in vector_req + assert vector_req["metric_type"] == "cosine" + assert vector_req["num_partitions"] == 2 + + with pytest.warns(DeprecationWarning, match="create_index"): + table.create_index( + index_type="IVF_HNSW_SQ", + num_partitions=1, + m=10, + ef_construction=50, + max_iterations=12, + sample_rate=64, + ) + tuned_req = received_requests[-1] + assert tuned_req["index_type"] == "IVF_HNSW_SQ" + assert tuned_req["num_partitions"] == 1 + assert tuned_req["m"] == 10 + assert tuned_req["ef_construction"] == 50 + assert tuned_req["max_iterations"] == 12 + assert tuned_req["sample_rate"] == 64 table.wait_for_index(["custom_scalar_idx"], timedelta(seconds=2)) table.wait_for_index( ["custom_fts_idx", "custom_vector_idx"], timedelta(seconds=2) ) - table.drop_index("custom_vector_idx") + table.drop_index(name="custom_vector_idx") table.drop_index("custom_scalar_idx") table.drop_index("custom_fts_idx") diff --git a/python/python/tests/test_table_signatures.py b/python/python/tests/test_table_signatures.py new file mode 100644 index 000000000..346935b2b --- /dev/null +++ b/python/python/tests/test_table_signatures.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import timedelta +import inspect +from unittest.mock import AsyncMock, MagicMock + +import lancedb +from lancedb.expr import col +from lancedb.index import BTree, FTS, HnswSq, IvfPq +from lancedb.remote.table import RemoteTable +from lancedb.table import LanceTable, Table +import pyarrow as pa +import pytest + + +@pytest.mark.parametrize( + "method", + [ + "delete", + "drop_index", + "index_stats", + "create_index", + "create_scalar_index", + "create_fts_index", + "search", + "cleanup_old_versions", + "compact_files", + "optimize", + "to_polars", + ], +) +def test_sync_table_signatures(method): + """A Table call must bind the same arguments on either backend.""" + + def parameters(cls): + return [ + (p.name, p.kind, p.default) + for p in inspect.signature(getattr(cls, method)).parameters.values() + ] + + assert parameters(Table) == parameters(LanceTable) == parameters(RemoteTable) + + +@pytest.fixture(params=[LanceTable, RemoteTable], ids=["local", "remote"]) +def sync_table(request, monkeypatch): + """Exercise real sync wrappers with a recording async backend.""" + handle = MagicMock() + handle.name = "test" + handle.create_index = AsyncMock() + handle.delete = AsyncMock() + handle.drop_index = AsyncMock() + handle.index_stats = AsyncMock() + if request.param is RemoteTable: + table = RemoteTable(handle, "test_db") + else: + table = LanceTable.__new__(LanceTable) + table._table = handle + monkeypatch.setattr(table, "_ensure_no_legacy_fts_index", lambda: None) + return table, handle + + +def test_sync_keyword_names(sync_table): + table, handle = sync_table + predicate = col("id") == 1 + table.delete(where=predicate) + handle.delete.assert_awaited_once() + expected = predicate._inner if isinstance(table, LanceTable) else predicate + assert handle.delete.await_args.args[0] is expected + table.drop_index(name="vector_idx") + handle.drop_index.assert_awaited_once_with("vector_idx") + table.index_stats(index_name="vector_idx") + handle.index_stats.assert_awaited_once_with("vector_idx") + + +def test_legacy_vector_positional_arguments(sync_table): + table, handle = sync_table + with pytest.warns(DeprecationWarning): + table.create_index("cosine", 2) + args, kwargs = handle.create_index.await_args + assert args == ("vector",) + assert isinstance(kwargs["config"], IvfPq) + assert kwargs["config"].distance_type == "cosine" + assert kwargs["config"].num_partitions == 2 + assert kwargs["replace"] is True + + +def test_vector_tuning_options(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_index( + index_type="IVF_HNSW_SQ", + num_partitions=1, + m=10, + ef_construction=50, + max_iterations=12, + sample_rate=64, + target_partition_size=1000, + wait_timeout=timeout, + name="custom_idx", + replace=False, + train=False, + ) + kwargs = handle.create_index.await_args.kwargs + config = kwargs.pop("config") + assert isinstance(config, HnswSq) + for name, value in { + "num_partitions": 1, + "m": 10, + "ef_construction": 50, + "max_iterations": 12, + "sample_rate": 64, + "target_partition_size": 1000, + }.items(): + assert getattr(config, name) == value + assert kwargs == dict( + wait_timeout=timeout, name="custom_idx", replace=False, train=False + ) + + +def test_scalar_defaults_and_wait_timeout(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_scalar_index("id", wait_timeout=timeout) + kwargs = handle.create_index.await_args.kwargs + assert isinstance(kwargs["config"], BTree) + assert kwargs["replace"] is True + assert kwargs["wait_timeout"] == timeout + + +def test_fts_tokenizer_and_wait_timeout(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_fts_index( + field_names="text", + tokenizer_name="en_stem", + replace=True, + wait_timeout=timeout, + custom_stop_words=["example"], + ) + args, kwargs = handle.create_index.await_args + assert args == ("text",) + assert kwargs["replace"] is True + assert kwargs["wait_timeout"] == timeout + config = kwargs["config"] + assert isinstance(config, FTS) + assert config.language == "English" + assert config.stem is True + assert config.custom_stop_words == ["example"] + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"use_tantivy": True}, "Tantivy-based FTS has been removed"), + ({"ordering_field_names": "id"}, "ordering_field_names"), + ({"writer_heap_size": 123}, "writer_heap_size"), + ({"tokenizer_name": "invalid"}, "Invalid tokenizer name"), + ], +) +def test_fts_unsupported_options(sync_table, kwargs, message): + table, handle = sync_table + with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match=message): + table.create_fts_index("text", **kwargs) + handle.create_index.assert_not_awaited() + + +@pytest.mark.parametrize( + "method, kwargs", + [ + ( + "cleanup_old_versions", + {"older_than": timedelta(days=1), "delete_unverified": True}, + ), + ("compact_files", {"target_rows_per_fragment": 1000}), + ], +) +def test_remote_maintenance_options_warn(method, kwargs): + table = RemoteTable(MagicMock(name="test"), "test_db") + with pytest.warns(UserWarning, match="no-op"): + getattr(table, method)(**kwargs) + + +def test_remote_optimize_options(): + handle = MagicMock() + handle.optimize = AsyncMock(side_effect=NotImplementedError("remote optimize")) + table = RemoteTable(handle, "test_db") + age = timedelta(days=1) + with pytest.raises(NotImplementedError, match="remote optimize"): + table.optimize(cleanup_older_than=age, delete_unverified=True, retrain=True) + handle.optimize.assert_awaited_once_with( + cleanup_older_than=age, delete_unverified=True, retrain=True + ) + + +@pytest.mark.parametrize( + "method, old, new, value", + [ + ("delete", "predicate", "where", "id = 1"), + ("drop_index", "index_name", "name", "vector_idx"), + ("index_stats", "index_uuid", "index_name", "vector_idx"), + ("create_fts_index", "column", "field_names", "text"), + ], +) +def test_deprecated_remote_keywords(method, old, new, value): + handle = MagicMock() + handle.name = "test" + target = "create_index" if method == "create_fts_index" else method + setattr(handle, target, AsyncMock()) + table = RemoteTable(handle, "test_db") + with pytest.warns(DeprecationWarning, match=f"{old} is deprecated"): + getattr(table, method)(**{old: value}) + assert getattr(handle, target).await_args.args == (value,) + with pytest.raises(TypeError, match="Cannot specify both"): + getattr(table, method)(**{old: value, new: value}) + + +@pytest.mark.parametrize("remote_wrapper", [False, True], ids=["local", "remote"]) +def test_sync_search_and_polars(tmp_path, remote_wrapper): + """Use the native engine through both sync wrappers, without a cloud service.""" + local = lancedb.connect(tmp_path).create_table( + "test", + [{"id": 1, "text": "quick fox", "vector": [1.0, 1.0]}], + ) + table = RemoteTable(local._table, "test_db") if remote_wrapper else local + table.create_index("text", config=FTS()) + assert table.search("quick", query_type="fts").to_arrow()["id"].to_pylist() == [1] + # Both backends must accept the deprecated ordering argument in the same slot. + with pytest.warns(DeprecationWarning, match="ordering_field_name"): + table.search("quick", None, "fts", "id", ["text"]) + fast = table.search([1.0, 1.0], fast_search=True).to_arrow() + assert len(fast) == 0 # No vector index: fast_search excludes unindexed rows. + assert table.search([1.0, 1.0]).to_arrow()["id"].to_pylist() == [1] + if remote_wrapper: + with pytest.raises(NotImplementedError, match=r"table.search\(\).to_polars"): + table.to_polars(batch_size=1) + assert table.search().to_polars()["id"].to_list() == [1] + else: + assert table.to_polars(batch_size=1).collect()["id"].to_list() == [1] + table.delete(where="id = 1") + assert table.count_rows() == 0 + + +def test_index_replacement_defaults(tmp_path): + local = lancedb.connect(tmp_path).create_table("test", pa.table({"id": [1]})) + table = RemoteTable(local._table, "test_db") + for sync in [local, table]: + with pytest.warns(DeprecationWarning): + sync.create_scalar_index("id") + assert len(list(table.list_indices())) == 1 From 982b2a1229ecbce435ed10a24bd0e54d1228f7ca Mon Sep 17 00:00:00 2001 From: KarthikS Date: Mon, 26 Oct 2026 08:12:53 -0400 Subject: [PATCH 34/36] fix(python): restore sync remote connection and table API parity (#4473) Fixes #4400. Sync remote connections fail on `uri`, `len(db)`, and `name in db`, and remote tables lack `namespace`, `id`, and `replace_field_metadata`. Share the paginated connection helpers and table identifier contract through the base classes, and delegate remote metadata replacement to the existing async implementation. Unsupported remote session and read-consistency properties raise explicit `NotImplementedError`s. Reconciled with main at `3dd57978`, including the synchronous method signatures from #4498. The local session fix has already landed in #4433, so this PR preserves its supplied-session-or-None contract and adds no Rust/PyO3 session access. Namespace connections retain their writable properties. Existing custom `Table` subclasses remain instantiable: `namespace` defaults to a fresh empty list, and deprecated `replace_field_metadata` delegates to `update_field_metadata` with `replace=True` while retaining its `None` return. Native and remote overrides are preserved. Regression tests instantiate an old-style subclass through normal ABC class creation and verify identifiers, independent namespace lists, deprecation, and exact metadata delegation. Historical validation at `03c43f7b` before the latest main integration and subclass compatibility repair, Ubuntu x86-64, Python 3.12.13: - Rebuilt the native extension with the repository `ci` profile and locked Python dependencies. - 273 tests passed across `test_sync_api_parity.py`, `test_db.py`, `test_remote_db.py`, and `test_session.py`. - 84 namespace and catalog compatibility tests passed (357 tests total). - Python SDK lint and formatting pass (158 files); changed-file checks and `git diff --check` pass. - Repository-wide lint/format checks report existing issues in unchanged `docs/web-tests` files (49 lint findings and 11 formatting files). These are outside this patch. Remote tests use the existing HTTP fixtures through the rebuilt native client; the original aarch64 server deployment was not recreated. Validation after merging main `3dd57978` and the compatibility repair (macOS ARM64, uv-managed Python 3.12.12, rebuilt editable native extension): - 60 focused parity, signature, and error-contract tests passed. - 381 database, remote, session, namespace, catalog, and sync parity tests passed; 2 platform-dependent tests skipped. - The two new subclass regressions fail against the pre-repair implementation with the expected abstract-member instantiation error, then pass with the repair. - Python SDK Ruff lint/format passes (167 files); repository-wide findings remain confined to unchanged docs/web-tests (49 lint findings, 11 formatting files). --------- Co-authored-by: Xuanwo --- python/python/lancedb/db.py | 33 ++-- python/python/lancedb/remote/db.py | 19 ++- python/python/lancedb/remote/table.py | 12 ++ python/python/lancedb/table.py | 40 ++++- python/python/tests/test_sync_api_parity.py | 170 ++++++++++++++++++++ 5 files changed, 251 insertions(+), 23 deletions(-) create mode 100644 python/python/tests/test_sync_api_parity.py diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 9ed772e1c..b6dfe8fb8 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -823,8 +823,26 @@ class DBConnection(EnforceOverrides): @property def uri(self) -> str: + """The URI of this database.""" return self._uri + def _all_table_names(self) -> Generator[str, None, None]: + page_token = None + while True: + response = self.list_tables(page_token=page_token) + yield from response.tables + page_token = response.page_token + if not page_token: + return + + def __len__(self) -> int: + """Return the number of tables in the root namespace across all pages.""" + return sum(1 for _ in self._all_table_names()) + + def __contains__(self, name: str) -> bool: + """Return whether the root namespace contains the named table.""" + return name in self._all_table_names() + def namespace_client(self) -> LanceNamespace: """Get the equivalent namespace client for this connection. @@ -1609,21 +1627,6 @@ class LanceDBConnection(DBConnection): ) ) - def _all_table_names(self) -> Generator[str, None, None]: - page_token = None - while True: - response = self.list_tables(page_token=page_token) - yield from response.tables - page_token = response.page_token - if not page_token: - return - - def __len__(self) -> int: - return sum(1 for _ in self._all_table_names()) - - def __contains__(self, name: str) -> bool: - return name in self._all_table_names() - @override def create_table( self, diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index c11e01ff4..4ffd3ac4a 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -47,7 +47,7 @@ from ..secrets import EnvVarSecret, SecretInfo from ..view import ViewDescription if TYPE_CHECKING: - from .._lancedb import JobInfo + from .._lancedb import JobInfo, Session from ..embeddings import EmbeddingFunctionConfig from lance_namespace import ( LanceNamespace, @@ -240,6 +240,23 @@ class RemoteDBConnection(DBConnection): def __repr__(self) -> str: return f"RemoteConnect(name={self.db_name})" + @property + @override + def uri(self) -> str: + return self._conn.uri + + @property + def session(self) -> Optional["Session"]: + """Remote connections do not expose a local cache session.""" + raise NotImplementedError("session is not supported by remote connections") + + @property + def read_consistency_interval(self) -> Optional[timedelta]: + """Remote connections do not expose a local read consistency interval.""" + raise NotImplementedError( + "read_consistency_interval is not supported by remote connections" + ) + @override def serialize(self) -> str: if getattr(self, "_catalog_oauth", False): diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index faf15c987..6b7835470 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -192,6 +192,18 @@ class RemoteTable(Table): """The name of the table""" return self._name + @property + def namespace(self) -> List[str]: + return self._namespace_path + + @deprecation.deprecated( + deprecated_in="0.33.1", + current_version=__version__, + details="Use update_field_metadata() instead.", + ) + def replace_field_metadata(self, field_name: str, new_metadata: Dict[str, str]): + LOOP.run(self._table.replace_field_metadata(field_name, new_metadata)) + def __repr__(self) -> str: return f"RemoteTable({self.db_name}.{self.name})" diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index d2ff02683..f1e538c14 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1360,6 +1360,16 @@ class Table(ABC): """The name of this Table""" raise NotImplementedError("name is not supported for this table type") + @property + def namespace(self) -> List[str]: + """Return the namespace path of the table (empty for the root namespace).""" + return [] + + @property + def id(self) -> str: + """Return the full identifier of the table (namespace$name).""" + return "$".join(self.namespace + [self.name]) + @property @abstractmethod def version(self) -> int: @@ -3106,6 +3116,29 @@ class Table(ABC): version: the new table version after the update. """ + @deprecation.deprecated( + deprecated_in="0.33.1", + current_version=__version__, + details="Use update_field_metadata() instead.", + ) + def replace_field_metadata(self, field_name: str, new_metadata: Dict[str, str]): + """Replace all metadata of a field in the schema. + + .. deprecated:: 0.33.1 + Use [update_field_metadata][lancedb.table.Table.update_field_metadata] + with ``replace=True`` instead. + + Parameters + ---------- + field_name: str + The name of the field whose metadata is replaced. + new_metadata: dict + The new metadata, replacing all existing keys. + """ + self.update_field_metadata( + {"path": field_name, "metadata": new_metadata, "replace": True} + ) + @abstractmethod def drop_columns(self, columns: Iterable[str]) -> DropColumnsResult: """ @@ -3315,13 +3348,6 @@ class LanceTable(Table): """Return the namespace path of the table.""" return self._namespace_path - @property - def id(self) -> str: - """Return the full identifier of the table (namespace$name).""" - if self._namespace_path: - return "$".join(self._namespace_path + [self.name]) - return self.name - @classmethod async def from_inner(cls, tbl: LanceDBTable): from .db import AsyncConnection, LanceDBConnection diff --git a/python/python/tests/test_sync_api_parity.py b/python/python/tests/test_sync_api_parity.py new file mode 100644 index 000000000..a12efd0c9 --- /dev/null +++ b/python/python/tests/test_sync_api_parity.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import timedelta +import json +from urllib.parse import parse_qs, urlparse + +import lancedb +from lancedb.table import Table +import pytest + +from test_remote_db import mock_lancedb_connection + + +@pytest.fixture +def legacy_table(): + """Implement the pre-existing abstract API without the two new members.""" + + def unsupported(self, *args, **kwargs): + raise NotImplementedError + + def update_field_metadata(self, updates): + self.metadata_updates.append(updates) + return object() # The deprecated wrapper must not expose this result. + + # Supply concrete implementations through normal ABC class creation. Leave + # namespace and replace_field_metadata inherited, as an old subclass would. + implementations = { + name: unsupported + for name in Table.__abstractmethods__ + if name not in {"namespace", "replace_field_metadata"} + } + implementations.update( + name=property(lambda self: "legacy"), + update_field_metadata=update_field_metadata, + ) + legacy_class = type("LegacyTable", (Table,), implementations) + table = legacy_class() + table.metadata_updates = [] + return table + + +def test_legacy_table_namespace_and_id(legacy_table): + assert legacy_table.namespace == [] + assert legacy_table.id == "legacy" + namespace = legacy_table.namespace + namespace.append("changed") + assert legacy_table.namespace == [] + assert legacy_table.id == "legacy" + + +def test_legacy_table_replace_field_metadata(legacy_table): + with pytest.warns(DeprecationWarning, match="update_field_metadata"): + result = legacy_table.replace_field_metadata("nested.field", {"a": "b"}) + assert result is None + assert legacy_table.metadata_updates == [ + {"path": "nested.field", "metadata": {"a": "b"}, "replace": True} + ] + + +def respond(request, body): + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(body).encode()) + + +def test_remote_uri(): + with mock_lancedb_connection(lambda request: respond(request, {})) as db: + assert db.uri == "db://dev" + + +@pytest.mark.parametrize("attribute", ["session", "read_consistency_interval"]) +def test_remote_unsupported_properties(attribute): + with mock_lancedb_connection(lambda request: respond(request, {})) as db: + with pytest.raises(NotImplementedError, match=attribute): + getattr(db, attribute) + + +@pytest.mark.parametrize("tables", [[], ["first", "last"]]) +def test_remote_len_and_contains_follow_pages(tables): + requests = [] + + def handler(request): + requests.append(request.path) + query = parse_qs(urlparse(request.path).query) + if tables and "page_token" not in query: + respond(request, {"tables": [tables[0]], "page_token": "opaque-next"}) + else: + if tables: + assert query["page_token"] == ["opaque-next"] + respond(request, {"tables": tables[1:]}) + + with mock_lancedb_connection(handler) as db: + assert len(db) == len(tables) + assert ("last" in db) == bool(tables) + assert "missing" not in db + requests.clear() + assert ("first" in db) == bool(tables) + assert len(requests) == 1 + + +@pytest.mark.parametrize("namespace", [[], ["parent", "child"]]) +def test_remote_table_identifiers(namespace): + def handler(request): + respond(request, {"version": 1, "schema": {"fields": []}}) + + with mock_lancedb_connection(handler) as db: + table = db.open_table("rfm", namespace_path=namespace) + assert table.namespace == namespace + assert table.id == "$".join(namespace + ["rfm"]) + + +def test_remote_replace_field_metadata(): + updates = [] + + def handler(request): + if request.path.endswith("/describe/"): + respond(request, {"version": 1, "schema": {"fields": []}}) + else: + assert request.path == "/v1/table/rfm/update_field_metadata/" + body = request.rfile.read(int(request.headers["Content-Length"])) + updates.append(json.loads(body)) + respond(request, {"version": 2}) + + with mock_lancedb_connection(handler) as db: + table = db.open_table("rfm") + with pytest.warns(DeprecationWarning, match="update_field_metadata"): + assert table.replace_field_metadata("id", {"a": "b"}) is None + assert updates == [ + {"updates": [{"path": "id", "metadata": {"a": "b"}, "replace": True}]} + ] + + +@pytest.mark.parametrize("custom_session", [False, True]) +def test_local_properties_and_metadata(tmp_path, custom_session): + supplied = lancedb.Session() if custom_session else None + db = lancedb.connect( + tmp_path, session=supplied, read_consistency_interval=timedelta(seconds=5) + ) + session = db.session + assert session is supplied + before = session.size_bytes if session is not None else 0 + table = db.create_table("rfm", [{"id": i} for i in range(3)]) + assert table.to_arrow().num_rows == 3 + if supplied is not None: + assert supplied.size_bytes > before + assert db.read_consistency_interval == timedelta(seconds=5) + assert table.namespace == [] + assert table.id == "rfm" + with pytest.warns(DeprecationWarning): + table.replace_field_metadata("id", {"old": "value"}) + with pytest.warns(DeprecationWarning): + table.replace_field_metadata("id", {"a": "b"}) + assert table.schema.field("id").metadata == {b"a": b"b"} + + +def test_namespace_connection_properties_remain_writable(tmp_path): + session = lancedb.Session() + interval = timedelta(seconds=5) + with lancedb.connect_namespace( + "dir", + {"root": str(tmp_path)}, + session=session, + read_consistency_interval=interval, + ) as db: + assert db.session is session + assert db.read_consistency_interval == interval + db.read_consistency_interval = timedelta(0) + assert db.read_consistency_interval == timedelta(0) From 180eaaf68eae595e0948f3bb2f02974d4cf0535e Mon Sep 17 00:00:00 2001 From: KarthikS Date: Mon, 26 Oct 2026 09:31:13 -0400 Subject: [PATCH 35/36] fix!: reject unsupported remote connection and table options (#4486) Fixes #4399. Remote connections currently discard local-only settings, including table write options that determine storage format. The synchronous Python client also binds the fourth positional `create_table` argument to `on_bad_vectors`, so an intended overwrite is sent as a create. This change aligns the remote Python signature with `DBConnection.create_table`, preserves the legacy remote timeout keywords, and raises `ValueError` for other unknown connection keywords. The shared Rust client rejects unsupported connection, create, and open options before sending HTTP requests; open validation also runs before a table-cache hit. Python reports these as `NotImplementedError`, including explicit `enable_v2_manifest_paths=False`. Supported Enterprise connection-level storage options remain available. This implements option 2 from the issue: explicit rejection. It deliberately does not add unverified server write parameters. Calls that previously supplied ignored options must omit them to use server-managed storage and caching. Database URI prefixes are already preserved on current main and are unchanged here. Regression tests cover Python sync/async, TypeScript create/open calls, and the Rust remote client, including empty-table creation, overwrite mode, and cached table handles. Python API documentation describes the behavior and corrected positional order. Reconciled with main at `0ea932da`, preserving upstream blob and error-handling tests and updating the Python test that previously expected open options to be ignored. Validation on the merged source: - Rebuilt the Python native extension; 304 tests passed across remote options, database, remote database, and session suites. - Rust remote database suite: 135 passed. Rebuilt the Node native extension; TypeScript remote suite: 62 passed. TypeScript compilation and documentation generation passed. - Full Python SDK lint and formatting pass (159 files); full TypeScript lint and formatting pass (66 files). Rust formatting and `git diff --check` pass. - Repository-wide Python lint/format reports existing issues in unchanged `docs/web-tests` files (49 lint findings and 11 formatting files). Environment: Ubuntu under WSL2, Linux x86-64, CPython 3.12.13, Rust 1.97.0, and Node.js 22.16.0. Locked dependencies and the repository `ci` build profile. Remote tests use HTTP fixtures; this does not claim validation against the original ARM64 single-node deployment. Documentation review follow-up: added a draft Enterprise migration note in `docs/web/basics/sdk.mdx`. It explains that `db://` connections reject local session/cache and per-table storage/write options, while supported connection-level Enterprise storage options remain available. Following the repository's release conventions, the note stays in an MDX comment below the migration list until the behavior is released. Vale 3.24.0 with the pinned LanceDB style package reports no findings; `git diff --check` passes. This follow-up changes documentation only. BREAKING CHANGE: Remote connections now reject unsupported local session/cache and per-table storage/write options instead of silently ignoring them. Callers must omit these options and use server-managed storage and caching; supported Enterprise connection-level storage options remain available. Integration with current main preserves its synchronous API signatures, including the remote `create_table(mode=None)` default, and the table compatibility defaults. The positional-order repair is already on main via #4608; this PR retains the distinct unsupported-option rejection behavior. CI follow-up: the inherited IVF search-effort test occasionally measured recall 0.375 against its unchanged 0.5 threshold after random index training. A matching local reproducer observed this once in 30 builds. The fixture now places the eight exact neighbors together, preserving the 16-partition setup, two writes, all probe/precedence assertions, and the recall threshold. The adjusted fixture returned recall 1.0 in 100 builds; the full query and remote suites passed (102 tests). This stabilizes test data and does not change search behavior or promise production recall. --------- Co-authored-by: Xuanwo --- docs/src/python/python.md | 11 ++ docs/web/basics/sdk.mdx | 7 + nodejs/__test__/query.test.ts | 16 +- nodejs/__test__/remote.test.ts | 49 ++++++ python/python/lancedb/__init__.py | 15 ++ python/python/lancedb/remote/db.py | 37 +++-- python/python/tests/test_remote_db.py | 10 +- python/python/tests/test_remote_options.py | 154 ++++++++++++++++++ rust/lancedb/src/connection.rs | 19 +++ rust/lancedb/src/connection/create_table.rs | 5 + rust/lancedb/src/remote/db.rs | 164 +++++++++++++++++++- 11 files changed, 460 insertions(+), 27 deletions(-) create mode 100644 python/python/tests/test_remote_options.py diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 4d0dfc87a..356b49ae7 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -20,6 +20,17 @@ pip install lancedb The following methods describe the synchronous API client. There is also an [asynchronous API client](#connections-asynchronous). +Connections to `db://` databases reject local-only options with +`NotImplementedError`. These include `session` and `manifest_enabled=True` on +connect, per-table storage and write options on create, and `index_cache_size`, +`storage_options`, or `location` on open (where the method accepts them). +Omit these options to use server-managed storage and caching. Connection-level +`storage_options` supported by an Enterprise deployment remain available. + +The synchronous remote `create_table` uses the same positional order as the +local API: `name, data, schema, mode, exist_ok, on_bad_vectors, fill_value, +embedding_functions`. Prefer keyword arguments for optional parameters. + ## Connections (Synchronous) ::: lancedb.connect diff --git a/docs/web/basics/sdk.mdx b/docs/web/basics/sdk.mdx index fc386bc7a..b20fc1662 100644 --- a/docs/web/basics/sdk.mdx +++ b/docs/web/basics/sdk.mdx @@ -92,6 +92,13 @@ manages background maintenance. [background maintenance](/maintenance/background-jobs) rather than calling the embedded table's `optimize()` method. +{/* Draft for the remote option validation release. +- **Remove local storage and cache settings.** `db://` connections reject local + sessions, table cache settings, and per-table storage and write options. Omit + these options so Enterprise manages storage and caching. Connection-level + storage options supported by your Enterprise deployment remain available. +*/} + For SQL connection details, see [SQL](/basics/sql). For the operating model and deployment options, see [LanceDB Enterprise](/enterprise). diff --git a/nodejs/__test__/query.test.ts b/nodejs/__test__/query.test.ts index 786f04015..068e55c2f 100644 --- a/nodejs/__test__/query.test.ts +++ b/nodejs/__test__/query.test.ts @@ -538,11 +538,17 @@ describe("IVF search effort", () => { const db = await connect(directory.name); const rows = Array.from({ length: 512 }, (_, id) => ({ id, - vector: Array.from( - { length: 16 }, - (_, j) => - (Math.floor(id / 32) === j ? 1 : 0) + Math.sin(id * 17 + j) * 0.001, - ), + // Keep all eight exact neighbors together even if random IVF training + // splits the first cluster. This tests effort, not training recall. + vector: + id < 8 + ? [...vector] + : Array.from( + { length: 16 }, + (_, j) => + (Math.floor(id / 32) === j ? 1 : 0) + + Math.sin(id * 17 + j) * 0.001, + ), })); table = await db.createTable("effort", rows.slice(0, 256)); await table.add(rows.slice(256)); diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 7e9479f94..d443537b0 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -98,6 +98,55 @@ async function withMockDatabase( } describe("remote connection", () => { + it.each([ + { dataStorageVersion: "2.0" }, + { enableV2ManifestPaths: false }, + { storageOptions: { bogusOptionXyz: "1" } }, + ])("rejects local create options %j before sending data", async (options) => { + const requests: string[] = []; + await withMockDatabase( + (req, res) => { + requests.push(req.url ?? ""); + res.writeHead(200).end(); + }, + async (db) => { + await expect( + db.createTable("test", [{ value: 1 }], { + ...options, + mode: "overwrite", + }), + ).rejects.toThrow(/write_options.*remote/); + await expect( + db.createEmptyTable( + "test", + new Schema([new Field("value", new Int64(), true)]), + options, + ), + ).rejects.toThrow(/write_options.*remote/); + }, + ); + expect(requests).toEqual([]); + }); + + it.each([{ indexCacheSize: 5 }, { storageOptions: { x: "y" } }])( + "rejects local open options %j before sending a request", + async (options) => { + const requests: string[] = []; + await withMockDatabase( + (req, res) => { + requests.push(req.url ?? ""); + res.writeHead(200).end(); + }, + async (db) => { + await expect( + db.openTable("test", undefined, options), + ).rejects.toThrow(/not supported for remote/); + }, + ); + expect(requests).toEqual([]); + }, + ); + it.each([false, true])( "normalizes remote blob writes (range target: %s)", async (wide) => { diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index a3aa55209..5e879f61d 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -281,6 +281,21 @@ def connect( "uri is required when not connecting through a namespace client" ) if isinstance(uri, str) and uri.startswith("db://"): + # Preserve the deprecated remote timeout options, but validate all other + # unknown keywords just as the local connection does. + unknown_kwargs = { + key: value + for key, value in kwargs.items() + if key not in ("connection_timeout", "read_timeout") + } + if unknown_kwargs: + raise ValueError(f"Unknown keyword arguments: {unknown_kwargs}") + if session is not None: + raise NotImplementedError("session is not supported for remote connections") + if manifest_enabled: + raise NotImplementedError( + "manifest_enabled is not supported for remote connections" + ) if api_key is None: api_key = os.environ.get("LANCEDB_API_KEY") if api_key is None: diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 4ffd3ac4a..291cc52d6 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -7,7 +7,6 @@ from __future__ import annotations from dataclasses import replace from datetime import timedelta import json -import logging from concurrent.futures import ThreadPoolExecutor import sys from typing import ( @@ -502,18 +501,14 @@ class RemoteDBConnection(DBConnection): if namespace_path is None: namespace_path = [] - if storage_options is not None: - logging.info( - "storage_options is ignored for remote tables" - " (storage is managed; set storage_options on connect() instead)" + table = LOOP.run( + self._conn.open_table( + name, + namespace_path=namespace_path, + storage_options=storage_options, + index_cache_size=index_cache_size, ) - if index_cache_size is not None: - logging.info( - "index_cache_size is ignored for remote tables" - " (there is no local cache to configure)" - ) - - table = LOOP.run(self._conn.open_table(name, namespace_path=namespace_path)) + ) tbl = RemoteTable( table, self.db_name, @@ -593,6 +588,9 @@ class RemoteDBConnection(DBConnection): embedding_functions: Optional[List[EmbeddingFunctionConfig]] = None, *, namespace_path: Optional[List[str]] = None, + storage_options: Optional[Dict[str, str]] = None, + data_storage_version: Optional[str] = None, + enable_v2_manifest_paths: Optional[bool] = None, ) -> Table: """Create a [Table][lancedb.table.Table] in the database. @@ -632,6 +630,12 @@ class RemoteDBConnection(DBConnection): embedding_functions: list of EmbeddingFunctionConfig, optional The embedding functions to store in the table schema and use to generate vectors when creating the table or adding data. + storage_options: dict, optional + Not supported for remote table creation. Raises NotImplementedError. + data_storage_version: str, optional + Not supported for remote table creation. Raises NotImplementedError. + enable_v2_manifest_paths: bool, optional + Not supported for remote table creation. Raises NotImplementedError. Returns ------- @@ -698,6 +702,14 @@ class RemoteDBConnection(DBConnection): LanceTable(table4) """ + for option, value in ( + ("data_storage_version", data_storage_version), + ("enable_v2_manifest_paths", enable_v2_manifest_paths), + ): + if value is not None: + raise NotImplementedError( + f"{option} is not supported for remote table creation" + ) if exist_ok: if mode == "create": mode = "exist_ok" @@ -719,6 +731,7 @@ class RemoteDBConnection(DBConnection): on_bad_vectors=on_bad_vectors, fill_value=fill_value, embedding_functions=embedding_functions, + storage_options=storage_options, ) ) return RemoteTable( diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index db3394735..32f999d18 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1823,12 +1823,12 @@ def test_remote_unsupported_legacy_hnsw_pq(): ) -def test_remote_python_only_messages(caplog): +def test_remote_python_only_messages(): with mock_lancedb_connection(remote_unsupported_handler) as db: - with caplog.at_level("INFO"): - table = db.open_table("test", storage_options={}, index_cache_size=1) - assert "storage_options is ignored for remote tables" in caplog.text - assert "index_cache_size is ignored for remote tables" in caplog.text + for options in [{"storage_options": {}}, {"index_cache_size": 1}]: + with pytest.raises(NotImplementedError, match=next(iter(options))): + db.open_table("test", **options) + table = db.open_table("test") for method in ["to_arrow", "to_pandas", "to_polars"]: with pytest.raises(NotImplementedError) as error: getattr(table, method)() diff --git a/python/python/tests/test_remote_options.py b/python/python/tests/test_remote_options.py new file mode 100644 index 000000000..2022aa29d --- /dev/null +++ b/python/python/tests/test_remote_options.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import inspect +from urllib.parse import parse_qs, urlparse + +import lancedb +import pyarrow as pa +import pytest + +from test_remote_db import mock_lancedb_connection, mock_lancedb_connection_async + + +@pytest.mark.parametrize("uri", ["memory://", "db://dev"]) +def test_unknown_connect_options(uri): + with pytest.raises(ValueError, match="Unknown keyword arguments.*bogus"): + lancedb.connect(uri, api_key="fake", bogus=1) + + +@pytest.mark.parametrize("option", ["session", "manifest_enabled"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_remote_connect_rejects_local_options(option, asynchronous): + value = lancedb.Session() if option == "session" else True + with pytest.raises(NotImplementedError, match=option): + if asynchronous: + await lancedb.connect_async("db://dev", api_key="fake", **{option: value}) + else: + lancedb.connect("db://dev", api_key="fake", **{option: value}) + + +def test_remote_create_signature(): + from lancedb.db import DBConnection + from lancedb.remote.db import RemoteDBConnection + + expected = inspect.signature(DBConnection.create_table).parameters + actual = inspect.signature(RemoteDBConnection.create_table).parameters + assert list(actual) == list(expected) + for name in expected: + assert actual[name].kind == expected[name].kind + if name == "mode": + # Remote connections retain their existing server-default mode. + assert actual[name].default is None + else: + assert actual[name].default == expected[name].default + + +def test_remote_create_positional_mode(): + modes = [] + + def handler(request): + assert urlparse(request.path).path == "/v1/table/ct_pos/create/" + modes.append(parse_qs(urlparse(request.path).query)["mode"]) + request.send_response(200) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + data = [{"value": 1}] + db.create_table("ct_pos", data, None, "overwrite") + db.create_table("ct_pos", data, None, "overwrite") + db.create_table("ct_pos", data, None, "create", True) + with pytest.raises(ValueError, match="Invalid mode"): + db.create_table("ct_pos", data, None, "drop") + assert modes == [["overwrite"], ["overwrite"], ["exist_ok"]] + + +CREATE_OPTIONS = [ + {"new_table_data_storage_version": "2.0"}, + {"new_table_enable_v2_manifest_paths": "false"}, + {"bogus_option_xyz": "1"}, +] + + +@pytest.mark.parametrize("storage_options", CREATE_OPTIONS) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.asyncio +async def test_remote_create_rejects_write_options(storage_options, empty): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + + kwargs = {"schema": pa.schema([("value", pa.int64())])} + if not empty: + kwargs["data"] = [{"value": 1}] + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError, match="write_options.*remote"): + await db.create_table("t", storage_options=storage_options, **kwargs) + assert requests == [] + + +@pytest.mark.parametrize( + "options", + [ + {"data_storage_version": "2.0"}, + {"enable_v2_manifest_paths": False}, + {"storage_options": {"x": "1"}}, + ], +) +def test_sync_remote_create_rejects_write_options(options): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError, match="remote"): + db.create_table("t", [{"value": 1}], **options) + assert requests == [] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "options", [{"index_cache_size": 10}, {"storage_options": {"x": "y"}}] +) +@pytest.mark.asyncio +async def test_remote_open_rejects_local_options(options, asynchronous): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + if asynchronous: + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError, match="remote"): + await db.open_table("t", **options) + else: + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError, match="remote"): + db.open_table("t", **options) + assert requests == [] + + +@pytest.mark.asyncio +async def test_remote_open_rejects_location(): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError, match="location.*remote"): + await db.open_table("t", location="/tmp/zzz") + assert requests == [] diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 26afef5e8..bff97bf31 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -174,6 +174,8 @@ impl OpenTableBuilder { /// Advanced parameters that can be used to customize table reads /// + /// Remote connections reject these parameters with [`Error::NotSupported`]. + /// /// If set, these will take precedence over any overlapping `OpenTableOptions` options pub fn lance_read_params(mut self, params: ReadParams) -> Self { self.request.lance_read_params = Some(params); @@ -1530,6 +1532,7 @@ impl ConnectBuilder { /// Enable or disable manifest-backed directory namespace mode for local /// native connections. + /// Remote connections reject enabling this option with [`Error::NotSupported`]. /// /// When enabled, the connection uses the directory namespace database /// directly for all table operations and forces @@ -1565,6 +1568,7 @@ impl ConnectBuilder { } /// Set a custom session for object stores and caching. + /// Remote connections reject this option with [`Error::NotSupported`]. /// /// By default, a new session with default configuration will be created. /// This method allows you to provide a custom session with your own @@ -1596,6 +1600,21 @@ impl ConnectBuilder { fn execute_remote(self) -> Result { use crate::remote::db::RemoteDatabaseOptions; + for (option, provided) in [ + ("session", self.request.session.is_some()), + ("manifest_enabled", self.request.manifest_enabled), + ( + "namespace_client_properties", + !self.request.namespace_client_properties.is_empty(), + ), + ] { + if provided { + return Err(Error::NotSupported { + message: format!("{option} is not supported for remote connections"), + }); + } + } + let mut merged_options = self.request.options.clone(); Self::apply_env_defaults(&ENV_VARS_TO_STORAGE_OPTS, &mut merged_options); let sql_host_override = merged_options.get(OPT_REMOTE_SQL_HOST_OVERRIDE).cloned(); diff --git a/rust/lancedb/src/connection/create_table.rs b/rust/lancedb/src/connection/create_table.rs index 39cc82ec0..9d4408ba7 100644 --- a/rust/lancedb/src/connection/create_table.rs +++ b/rust/lancedb/src/connection/create_table.rs @@ -45,6 +45,9 @@ impl CreateTableBuilder { } /// Apply the given write options when writing the initial data + /// + /// Remote connections reject explicit Lance write parameters with + /// [`Error::NotSupported`] instead of silently using server defaults. pub fn write_options(mut self, write_options: WriteOptions) -> Self { self.request.write_options = write_options; self @@ -70,6 +73,8 @@ impl CreateTableBuilder { /// Set multiple options for the storage layer. /// + /// Table-level storage options are not supported by remote connections. + /// /// Options already set on the connection will be inherited by the table, /// but can be overridden here. /// diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index d33565653..5c299dc5a 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1481,6 +1481,23 @@ impl Database for RemoteDatabase { } async fn create_table(&self, mut request: CreateTableRequest) -> Result> { + // The remote create endpoint does not transmit Lance write parameters. + // Reject them before consuming data or issuing a request, especially + // for overwrite: silently using the server's format can lose intent. + for (option, provided) in [ + ( + "write_options (including storage_options)", + request.write_options.lance_write_params.is_some(), + ), + ("location", request.location.is_some()), + ("namespace_client", request.namespace_client.is_some()), + ] { + if provided { + return Err(Error::NotSupported { + message: format!("{option} is not supported for remote table creation"), + }); + } + } let data_schema = request.data.schema(); let body = stream_as_body(request.data.scan_as_stream())?; @@ -1605,6 +1622,23 @@ impl Database for RemoteDatabase { } async fn open_table(&self, request: OpenTableRequest) -> Result> { + // Validate before looking in the cache as well as before sending HTTP. + for (option, provided) in [ + ("index_cache_size", request.index_cache_size.is_some()), + ( + "lance_read_params (including storage_options)", + request.lance_read_params.is_some(), + ), + ("location", request.location.is_some()), + ("namespace_client", request.namespace_client.is_some()), + ("managed_versioning", request.managed_versioning.is_some()), + ] { + if provided { + return Err(Error::NotSupported { + message: format!("{option} is not supported for remote table access"), + }); + } + } let identifier = build_table_identifier(&request.name, &request.namespace_path)?; let cache_key = build_cache_key(&request.name, &request.namespace_path); @@ -2504,14 +2538,14 @@ mod tests { let table = conn.open_table("table1").execute().await.unwrap(); assert_eq!(table.name(), "table1"); - // Storage options should be ignored. - let table = conn + // Per-table storage options must not be silently ignored. + let result = conn .open_table("table1") .storage_option("key", "value") .execute() - .await - .unwrap(); - assert_eq!(table.name(), "table1"); + .await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains("storage_options"))); } #[tokio::test] @@ -2735,6 +2769,126 @@ mod tests { ); } + #[rstest] + #[case("new_table_data_storage_version", "2.0")] + #[case("new_table_enable_v2_manifest_paths", "false")] + #[case("bogus_option_xyz", "1")] + #[tokio::test] + async fn test_create_rejects_local_storage_options(#[case] key: &str, #[case] value: &str) { + let conn = Connection::new_with_handler(|_| -> http::Response<&'static str> { + panic!("must not send a create request") + }); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + for builder in [ + conn.create_table("table1", data), + conn.create_empty_table("table1", schema), + ] { + let result = builder + .mode(CreateTableMode::Overwrite) + .storage_option(key, value) + .execute() + .await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains("write_options") && message.contains("remote"))); + } + } + + #[tokio::test] + async fn test_create_rejects_local_write_params_and_location() { + let conn = Connection::new_with_handler(|_| -> http::Response<&'static str> { + panic!("must not send a create request") + }); + let schema = Arc::new(Schema::empty()); + for builder in [ + conn.create_empty_table("table1", schema.clone()) + .write_options(crate::table::WriteOptions { + lance_write_params: Some(lance::dataset::WriteParams::default()), + }), + conn.create_empty_table("table1", schema) + .location("/tmp/elsewhere"), + ] { + assert!(matches!( + builder.execute().await, + Err(Error::NotSupported { .. }) + )); + } + } + + #[rstest] + #[tokio::test] + async fn test_open_rejects_local_options_even_when_cached(#[values(false, true)] cached: bool) { + let mut db = super::RemoteDatabase::new_mock(|request| { + assert_eq!(request.url().path(), "/v1/table/table1/create/"); + http::Response::builder().status(200).body("").unwrap() + }); + db.table_cache = moka::future::Cache::new(10); + let table_cache = db.table_cache.clone(); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + if cached { + conn.create_empty_table("table1", Arc::new(Schema::empty())) + .execute() + .await + .unwrap(); + } + assert_eq!( + table_cache + .get(&build_cache_key("table1", &[])) + .await + .is_some(), + cached + ); + for (option, builder) in [ + ( + "index_cache_size", + conn.open_table("table1").index_cache_size(10), + ), + ( + "storage_options", + conn.open_table("table1").storage_option("x", "y"), + ), + ("location", conn.open_table("table1").location("/tmp/zzz")), + ( + "managed_versioning", + conn.open_table("table1").managed_versioning(false), + ), + ( + "lance_read_params", + conn.open_table("table1") + .lance_read_params(Default::default()), + ), + ] { + let result = builder.execute().await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains(option) && message.contains("remote"))); + } + } + + #[tokio::test] + async fn test_connect_rejects_local_options() { + for (option, builder) in [ + ( + "manifest_enabled", + ConnectBuilder::new("db://dev").manifest_enabled(true), + ), + ( + "namespace_client_properties", + ConnectBuilder::new("db://dev").namespace_client_properties([("root", "/tmp/db")]), + ), + ] { + let result = builder.api_key("fake").region("us-east-1").execute().await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains(option) && message.contains("remote"))); + } + } + #[rstest] #[case::matching(Field::new("a", DataType::Int32, false), true)] #[case::different_name(Field::new("x", DataType::Int32, false), false)] From 7952ee2a71d5bb921d097f3cd81f8039bf014818 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:31:19 +0800 Subject: [PATCH 36/36] fix(python): cover FP16 GPU index queries (#3819) Keep FP16 vectors through IVF-PQ index creation and search in the synchronous and asynchronous quickstarts. Add native and conditional CUDA regression coverage using a float64 query against an explicitly typed FP16 column. This is test and documentation-example coverage for dtype preservation already present in Lance; it does not change production code. Related #2120. Validation on `bbad20940632b21a9ee85bdcb46b4cc19dbd65e3` with Python 3.12.12, Lance 15.0.0-beta.1, PyArrow 24.0.0, and Torch 2.14.1: - The native FP16 parameter, Arrow FP16 ingestion test, and both quickstarts passed: 4 passed, 1 CUDA-only skip. - The two affected test files passed: 386 passed, 1 CUDA-only skip (`-m "not slow and not s3_test"`). - Ruff 0.9.9 lint and formatting checks passed for the Python package. Real CUDA execution has not been performed. Torch is installed, but `torch.cuda.is_available()` is false; `test_create_f16_table[cuda]` skips with `CUDA not available`. The native CPU result does not establish that the original GPU failure is resolved. Issue #2120 remains open for verification on an actual CUDA device. Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/tests/docs/test_basic.py | 16 ++-------------- python/python/tests/test_table.py | 26 +++++++++++++++++++++----- 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/python/python/tests/docs/test_basic.py b/python/python/tests/docs/test_basic.py index 35d7aac10..ef68114aa 100644 --- a/python/python/tests/docs/test_basic.py +++ b/python/python/tests/docs/test_basic.py @@ -91,19 +91,13 @@ def test_quickstart(tmp_path): } ) # --8<-- [end:alter_columns_vector] - # Change it back since we can get a panic with fp16 - tbl.alter_columns( - { - "path": "vector", - "data_type": pa.list_(pa.float32(), list_size=2), - } - ) # --8<-- [start:drop_columns] tbl.drop_columns(["dbl_price"]) # --8<-- [end:drop_columns] # --8<-- [start:create_index] tbl.create_index(num_sub_vectors=1) # --8<-- [end:create_index] + tbl.search([100, 100]).limit(2).to_pandas() # --8<-- [start:delete_rows] tbl.delete("item = 'fizz'") # --8<-- [end:delete_rows] @@ -185,13 +179,6 @@ async def test_quickstart_async(tmp_path): } ) # --8<-- [end:alter_columns_async_vector] - # Change it back since we can get a panic with fp16 - await tbl.alter_columns( - { - "path": "vector", - "data_type": pa.list_(pa.float32(), list_size=2), - } - ) # --8<-- [start:drop_columns_async] await tbl.drop_columns(["dbl_price"]) # --8<-- [end:drop_columns_async] @@ -200,6 +187,7 @@ async def test_quickstart_async(tmp_path): # --8<-- [start:create_index_async] await tbl.create_index("vector") # --8<-- [end:create_index_async] + await tbl.vector_search([100, 100]).limit(2).to_pandas() # --8<-- [start:delete_rows_async] await tbl.delete("item = 'fizz'") # --8<-- [end:delete_rows_async] diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 3b3ccd122..97f65b180 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -3965,26 +3965,42 @@ def test_create_f16_table_from_arrow_data(mem_db: DBConnection): assert "s-2" in expected["text"].to_pylist() -def test_create_f16_table(mem_db: DBConnection): +@pytest.mark.parametrize("accelerator", [None, "cuda"]) +def test_create_f16_table(tmp_path, accelerator): + if accelerator == "cuda": + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + class MyTable(LanceModel): text: str vector: Vector(32, value_type=pa.float16()) rng = np.random.default_rng(42) + vectors = rng.standard_normal((512, 32)).astype(np.float16) df = pa.table( { "text": [f"s-{i}" for i in range(512)], - "vector": [rng.standard_normal(32).astype(np.float16) for _ in range(512)], + "vector": list(vectors), } ) - table = mem_db.create_table( + db = lancedb.connect(tmp_path) + table = db.create_table( "f16_tbl", schema=MyTable, ) table.add(df) - table.create_index(num_partitions=2, num_sub_vectors=2) + table.create_index( + "vector", + config=IvfPq( + num_partitions=2, + num_sub_vectors=2, + accelerator=accelerator, + ), + ) - query = df["vector"][2].as_py() + # Match the issue's float64 query against an explicitly typed float16 column. + query = vectors[2].astype(np.float64) expected = table.search(query).limit(2).to_arrow() assert "s-2" in expected["text"].to_pylist()