diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index 8d3749f38..c0b33dcf0 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -20,12 +20,13 @@ 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 top-level blob column holds the blob bytes as - `LargeBinary`. 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 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. -Blobs nested in a struct or a list keep their descriptors. +Blobs nested in a list keep their descriptors. 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__/table.test.ts b/nodejs/__test__/table.test.ts index 84340c72c..92e3df456 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3385,6 +3385,106 @@ describe("when dealing with blob columns", () => { expect(bytes[0]!.equals(payload)).toBe(true); }); + it("returns bytes for blobs nested in structs", async () => { + const { table, alpha, empty } = await openNestedBlobTable(); + const rows = await table.query().toArray({ blobMode: "bytes" }); + const byId = new Map(rows.map((row) => [Number(row.id), row])); + expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha); + expect(byId.get(1)!.info.caption).toBe("one"); + expect(byId.get(2)!.info.image).toBeNull(); + expect(byId.get(2)!.info.caption).toBe("two"); + expect(byId.get(3)!.info).toBeNull(); + expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty); + }); + + it("lines up nested rows and bytes across single-row batches", async () => { + const { table, alpha, empty } = await openNestedBlobTable(); + const result = await table + .query() + .toArrow({ blobMode: "bytes", maxBatchLength: 1 }); + expect(result.batches.length).toBe(4); + const byId = new Map(result.toArray().map((row) => [Number(row.id), row])); + expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha); + expect(byId.get(1)!.info.caption).toBe("one"); + expect(byId.get(2)!.info.image).toBeNull(); + expect(byId.get(2)!.info.caption).toBe("two"); + expect(byId.get(3)!.info).toBeNull(); + expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty); + }); + + it("handles struct fields named like Object prototypes", async () => { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + new Field( + "info", + new Struct([ + new Field("constructor", new Struct([blob("image")]), true), + ]), + true, + ), + ]); + const payload = Buffer.from("proto"); + const table = await db.createTable( + "prototype_named_blobs", + [{ id: 1n, info: { constructor: { image: payload } } }], + { schema }, + ); + expect(await table.blobColumns()).toEqual(["info.constructor.image"]); + const [row] = await table.query().toArray({ blobMode: "bytes" }); + // Arrow's StructRow shadows the `constructor` property, so read the + // field through the JSON form instead of property access. + const info = JSON.parse(JSON.stringify(row.info)); + expect(Buffer.from(Object.values(info.constructor.image))).toEqual( + payload, + ); + }); + + it("returns nested blob bytes for a struct selection", async () => { + const { table, alpha } = await openNestedBlobTable(); + const [row] = await table + .query() + .where("id = 1") + .select(["info"]) + .toArray({ blobMode: "bytes" }); + expect(Object.keys(row)).toEqual(["info"]); + expect(Buffer.from(row.info.image)).toEqual(alpha); + expect(row.info.caption).toBe("one"); + }); + + it("returns nested blob bytes through aliases", async () => { + const { table, alpha } = await openNestedBlobTable(); + const [row] = await table + .query() + .where("id = 1") + .select( + new Map([ + ["picture", "info.image"], + ["details", "info"], + ]), + ) + .toArray({ blobMode: "bytes" }); + expect(Buffer.from(row.picture)).toEqual(alpha); + expect(Buffer.from(row.details.image)).toEqual(alpha); + expect(row.details.caption).toBe("one"); + }); + + it("returns a nested bytes schema for an empty result", async () => { + const { table } = await openNestedBlobTable(); + const result = await table + .query() + .where("id > 100") + .toArrow({ blobMode: "bytes" }); + expect(result.numRows).toBe(0); + const field = result.schema.fields.find((f) => f.name === "info")!; + expect(arrow.DataType.isStruct(field.type)).toBe(true); + const image = (field.type as Struct).children.find( + (f) => f.name === "image", + )!; + expect(arrow.DataType.isLargeBinary(image.type)).toBe(true); + expect(isBlobField(image)).toBe(false); + }); + it("creates list blobs without advertising unfetchable elements", async () => { const db = await connect(tmpDir.name); const schema = new Schema([ @@ -3593,6 +3693,31 @@ describe("when dealing with blob columns", () => { return { table, rowIds, alpha, beta }; } + async function openNestedBlobTable() { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + new Field( + "info", + new Struct([blob("image"), new Field("caption", new Utf8(), true)]), + true, + ), + ]); + const alpha = Buffer.from("alpha"); + const empty = Buffer.alloc(0); + const table = await db.createTable( + "nested_blobs", + [{ id: 1n, info: { image: alpha, caption: "one" } }], + { schema }, + ); + await table.add([ + { id: 2n, info: { image: null, caption: "two" } }, + { id: 3n, info: null }, + { id: 4n, info: { image: empty, caption: "four" } }, + ]); + return { table, alpha, empty }; + } + async function openExternalBlobTable() { const db = await connect(path.join(tmpDir.name, "db")); const payload = Buffer.from("external-payload"); diff --git a/nodejs/lancedb/arrow.ts b/nodejs/lancedb/arrow.ts index 68b35e256..56c8e1897 100644 --- a/nodejs/lancedb/arrow.ts +++ b/nodejs/lancedb/arrow.ts @@ -61,6 +61,7 @@ import { import { inferSchema } from "./schema"; export * from "apache-arrow"; +export type { Data as ArrowData } from "apache-arrow"; export type SchemaLike = | Schema | { diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 2abe6db97..823239ef9 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors import { + type ArrowData, Table as ArrowTable, + DataType, Field, type IntoVector, LargeBinary, @@ -70,16 +72,80 @@ function blobSources( blobColumns: string[], projection: [string, string][] | undefined, ): Map { - const columns = new Set(blobColumns); + const holds = (source: string) => + blobColumns.some( + (column) => column === source || column.startsWith(`${source}.`), + ); const pairs = projection ?? - blobColumns.map((column): [string, string] => [column, column]); - return new Map(pairs.filter(([, source]) => columns.has(source))); + 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; + } + 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, + }); } async function readBlobBytes( table: ArrowTable, sources: Map, + blobColumns: string[], snapshot: NativeTable, dropRowId: boolean, ): Promise { @@ -91,7 +157,11 @@ async function readBlobBytes( rowIdColumn === null ? [] : Array.from(rowIdColumn.toArray() as BigUint64Array); - const replacements = new Map(); + const fetch = async (column: string) => + (rowIds.length === 0 ? [] : await snapshot.fetchBlobs(column, rowIds)).map( + (value) => value ?? null, + ); + const replacements = new Map(); for (const [output, source] of sources) { const index = table.schema.fields.findIndex( (field) => field.name === output, @@ -99,12 +169,30 @@ async function readBlobBytes( if (index < 0) { continue; } - const values = - rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds); - replacements.set( - index, - values.map((value) => value ?? null), - ); + 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) { @@ -121,26 +209,37 @@ async function readBlobBytes( // 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) => - replacements.has(index) + const fields = table.schema.fields.map((field, index) => { + const bytes = replacements.get(index); + if (bytes === undefined) { + return field; + } + return isLeaf(bytes) ? new Field(field.name, new LargeBinary(), true) - : field, - ); + : new Field( + field.name, + nestedFieldType(field.type as Struct, bytes), + field.nullable, + ); + }); 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 values = replacements.get(index); - if (values === undefined) { + const bytes = replacements.get(index); + if (bytes === undefined) { return child; } + if (!isLeaf(bytes)) { + return nestedData(child, bytes, offset); + } return vectorFromArray( - values.slice(offset, offset + batch.numRows), + bytes.slice(offset, offset + batch.numRows), new LargeBinary(), ).data[0]; }); @@ -175,12 +274,13 @@ export interface QueryExecutionOptions { * * - `"descriptions"` (default): each blob is a descriptor. Read the bytes * with {@link Table.fetchBlobs} or {@link Table.fetchBlobFiles}. - * - `"bytes"`: each top-level blob column holds the blob bytes as - * `LargeBinary`. 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 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. * - * Blobs nested in a struct or a list keep their descriptors. + * Blobs nested in a list keep their descriptors. * * 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. @@ -429,17 +529,15 @@ 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 sources = blobSources( - await snapshot.blobColumns(), - context.projection, - ); + const blobColumns = await snapshot.blobColumns(); + const sources = blobSources(blobColumns, context.projection); if (sources.size === 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, snapshot, addRowId); + return readBlobBytes(table, sources, blobColumns, snapshot, addRowId); } if (attempt === BLOB_BYTES_ATTEMPTS) { throw new Error(