diff --git a/docs/src/js/classes/AutoQuery.md b/docs/src/js/classes/AutoQuery.md index 79d666124..12f517dbd 100644 --- a/docs/src/js/classes/AutoQuery.md +++ b/docs/src/js/classes/AutoQuery.md @@ -417,6 +417,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/docs/src/js/classes/Query.md b/docs/src/js/classes/Query.md index 0abf49755..3038bf2b3 100644 --- a/docs/src/js/classes/Query.md +++ b/docs/src/js/classes/Query.md @@ -491,6 +491,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/docs/src/js/classes/QueryBase.md b/docs/src/js/classes/QueryBase.md index 664337bad..0ef4fa286 100644 --- a/docs/src/js/classes/QueryBase.md +++ b/docs/src/js/classes/QueryBase.md @@ -252,6 +252,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + *** ### withRowId() diff --git a/docs/src/js/classes/TakeQuery.md b/docs/src/js/classes/TakeQuery.md index 79a53f753..f3d59a2a3 100644 --- a/docs/src/js/classes/TakeQuery.md +++ b/docs/src/js/classes/TakeQuery.md @@ -267,6 +267,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from [`QueryBase`](QueryBase.md).[`toArrow`](QueryBase.md#toarrow) diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index 29749c28f..f49963d60 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -740,6 +740,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/nodejs/__test__/query.test.ts b/nodejs/__test__/query.test.ts index 5f3e68b16..2cfecf823 100644 --- a/nodejs/__test__/query.test.ts +++ b/nodejs/__test__/query.test.ts @@ -110,6 +110,129 @@ describe("Query outputSchema", () => { }); }); +describe("Query empty results", () => { + let tmpDir: tmp.DirResult; + let table: Table; + + beforeEach(async () => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + const schema = new Schema( + [ + new Field("id", new Int64(), false), + new Field( + "text", + new Utf8(), + true, + new Map([["description", "searchable text"]]), + ), + new Field( + "vector", + new FixedSizeList(2, new Field("item", new Float32())), + true, + ), + ], + new Map([["source", "empty-query-regression"]]), + ); + table = await db.createTable( + "test", + makeArrowTable([{ id: 1n, text: "hello", vector: [1, 2] }], { schema }), + ); + }); + + afterEach(() => { + tmpDir.removeCallback(); + }); + + it.each([ + { + name: "filtered scan", + query: (table: Table) => table.query().where("id < 0"), + fields: ["id", "text", "vector"], + }, + { + name: "zero limit", + query: (table: Table) => table.query().limit(0), + fields: ["id", "text", "vector"], + }, + { + name: "column projection", + query: (table: Table) => + table.query().where("id < 0").select(["text", "id"]), + fields: ["text", "id"], + }, + { + name: "dynamic projection", + query: (table: Table) => + table.query().where("id < 0").select({ doubled: "id * 2" }), + fields: ["doubled"], + }, + { + name: "filtered vector search", + query: (table: Table) => + table.search([0.5, 0.5]).where("id < 0").limit(3), + fields: ["id", "text", "vector", "_distance"], + }, + { + name: "vector projection with row id", + query: (table: Table) => + table + .vectorSearch([0.5, 0.5]) + .where("id < 0") + .select(["id"]) + .withRowId(), + fields: ["id", "_distance", "_rowid"], + }, + { + name: "fast search without an index", + query: (table: Table) => table.search([0.5, 0.5]).fastSearch(), + fields: ["id", "text", "vector", "_rowid", "_distance"], + }, + { + name: "take with no offsets", + query: (table: Table) => table.takeOffsets([]).select(["text"]), + fields: ["text"], + }, + ])("preserves the output schema for $name", async ({ query, fields }) => { + const result = await query(table).toArrow({ maxBatchLength: 1 }); + + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual(fields); + for (const name of fields) { + expect(result.getChild(name)).not.toBeNull(); + expect(result.getChild(name)?.length).toBe(0); + } + }); + + it("preserves types, nullability, and metadata", async () => { + const result = await table.query().where("id < 0").toArrow(); + const schema = await table.schema(); + + expect(result.schema).toEqual(schema); + expect(result.getChild("id")?.type.toString()).toBe("Int64"); + expect(result.schema.fields[0].nullable).toBe(false); + expect(result.schema.fields[1].metadata.get("description")).toBe( + "searchable text", + ); + expect(result.schema.metadata.get("source")).toBe("empty-query-regression"); + }); + + it("preserves the score column for full-text search with no matches", async () => { + await table.createIndex("text", { config: Index.fts() }); + const result = await table + .search("missing", "fts") + .select(["id"]) + .toArrow(); + + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual([ + "id", + "_score", + ]); + expect(result.getChild("_score")?.type.toString()).toBe("Float32"); + }); +}); + describe("Search pagination", () => { let tmpDir: tmp.DirResult; let table: Table; diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 898ba533e..5422945b6 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -14,6 +14,14 @@ import { TlsConfig, connect, } from "../lancedb"; +import { + Table as ArrowTable, + Field, + Int64, + RecordBatch, + Schema, + tableToIPC, +} from "../lancedb/arrow"; import { HeaderProvider, OAuthHeaderProvider, @@ -85,6 +93,66 @@ async function withMockDatabase( } describe("remote connection", () => { + it.each([false, true])( + "preserves an empty query's schema with an empty batch: %s", + async (withEmptyBatch) => { + const schema = new Schema( + [new Field("doubled", new Int64(), false)], + new Map([["source", "query-output"]]), + ); + const result = new ArrowTable( + schema, + withEmptyBatch ? [new RecordBatch(schema, undefined)] : [], + ); + const response = Buffer.from(tableToIPC(result, "stream")); + let queryRequests = 0; + + await withMockDatabase( + (req, res) => { + if (req.url?.endsWith("/describe/")) { + res.writeHead(200, { "Content-Type": "application/json" }).end( + JSON.stringify({ + name: "items", + version: 1, + schema: { + fields: [ + { name: "id", type: { type: "int64" }, nullable: false }, + ], + }, + }), + ); + } else if (req.url?.endsWith("/query/")) { + queryRequests++; + req.resume(); + req.on("end", () => { + res + .writeHead(200, { + "Content-Type": "application/vnd.apache.arrow.stream", + }) + .end(response); + }); + } else { + res.writeHead(404).end(); + } + }, + async (db) => { + const table = await db.openTable("items"); + const result = await table + .query() + .where("id < 0") + .select({ doubled: "id * 2" }) + .toArrow(); + + expect(result.numRows).toBe(0); + expect(result.schema).toEqual(schema); + expect(result.getChild("doubled")?.length).toBe(0); + }, + ); + + expect(queryRequests).toBe(1); + }, + ); + it("rejects the external blob opt-in without blocking regular adds", async () => { let insertRequests = 0; const describeRequests: string[] = []; diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 75a787fcd..7c0dc9221 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -41,26 +41,6 @@ export async function* RecordBatchIterator( } } -class RecordBatchIterable< - NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery, -> implements AsyncIterable -{ - private inner: NativeQueryType; - private options?: QueryExecutionOptions; - - constructor(inner: NativeQueryType, options?: QueryExecutionOptions) { - this.inner = inner; - this.options = options; - } - - // biome-ignore lint/suspicious/noExplicitAny: skip - [Symbol.asyncIterator](): AsyncIterator, any, undefined> { - return RecordBatchIterator( - this.inner.execute(this.options?.maxBatchLength, this.options?.timeoutMs), - ); - } -} - /** * Options that control the behavior of a particular query execution */ @@ -271,13 +251,20 @@ export class QueryBase< return RecordBatchIterator(this.nativeExecute()); } - /** Collect the results as an Arrow @see {@link ArrowTable}. */ + /** + * Collect the results as an Arrow @see {@link ArrowTable}. + * + * Empty results retain the query's output schema. + */ async toArrow(options?: Partial): Promise { const batches = []; - const inner = await this.getInner(); - for await (const batch of new RecordBatchIterable(inner, options)) { + const iterator = this.nativeExecute(options); + for await (const batch of RecordBatchIterator(iterator)) { batches.push(batch); } + if (batches.length === 0) { + return tableFromIPC((await iterator).schema()); + } return new ArrowTable(batches); } diff --git a/nodejs/src/iterator.rs b/nodejs/src/iterator.rs index 8d65dfc3f..96d08c765 100644 --- a/nodejs/src/iterator.rs +++ b/nodejs/src/iterator.rs @@ -7,6 +7,8 @@ use lancedb::ipc::batches_to_ipc_file; use napi::bindgen_prelude::*; use napi_derive::napi; +use crate::util::schema_to_buffer; + /** Typescript-style Async Iterator over RecordBatches */ #[napi] pub struct RecordBatchIterator { @@ -19,6 +21,12 @@ impl RecordBatchIterator { Self { inner } } + /// Return the executed stream's output schema as an Arrow IPC file. + #[napi] + pub fn schema(&self) -> napi::Result { + schema_to_buffer(&self.inner.schema()) + } + #[napi(catch_unwind)] pub async unsafe fn next(&mut self) -> napi::Result> { if let Some(rst) = self.inner.next().await {