feat(node): add blobMode to return blob bytes from queries (#4353)

nodejs/__test__/table.test.ts

@@ -10,6 +10,7 @@ import * as tmp from "tmp";
 import {
   AutoQuery,
   type BlobInput,
+  type BlobMode,
   Connection,
   MatchQuery,
   PhraseQuery,
@@ -18,6 +19,7 @@ import {
   VectorQuery,
   blob,
   connect,
+  isBlobField,
   tokenize,
 } from "../lancedb";
 import {
@@ -3016,6 +3018,161 @@ describe("when dealing with blob columns", () => {
     expect(await handle!.readRange(0n, 4n)).toEqual(payload.subarray(0, 4));
   });
 
+  it("toArrow returns blob bytes with blobMode bytes", async () => {
+    const { table, alpha, beta } = await openBlobTable();
+    const result = await table.query().toArrow({ blobMode: "bytes" });
+    const field = result.schema.fields.find((f) => f.name === "image")!;
+    expect(arrow.DataType.isLargeBinary(field.type)).toBe(true);
+    expect(isBlobField(field)).toBe(false);
+    expect(result.schema.fields.map((f) => f.name)).toEqual(["id", "image"]);
+    const byId = new Map(
+      result.toArray().map((row) => [Number(row.id), row.image]),
+    );
+    expect(Buffer.from(byId.get(1)!)).toEqual(alpha);
+    expect(Buffer.from(byId.get(2)!)).toEqual(beta);
+    expect(byId.get(3)).toBeNull();
+  });
+
+  it("toArray returns blob bytes and keeps _rowid only when asked", async () => {
+    const { table, rowIds, alpha } = await openBlobTable();
+    const [plain] = await table
+      .query()
+      .where("id = 1")
+      .select(["id", "image"])
+      .toArray({ blobMode: "bytes" });
+    expect(Object.keys(plain)).toEqual(["id", "image"]);
+    expect(Buffer.from(plain.image)).toEqual(alpha);
+
+    const [withRowId] = await table
+      .query()
+      .where("id = 1")
+      .withRowId()
+      .toArray({ blobMode: "bytes" });
+    expect(withRowId._rowid).toBe(rowIds[0]);
+    expect(Buffer.from(withRowId.image)).toEqual(alpha);
+  });
+
+  it("fills a renamed blob column with bytes", async () => {
+    const { table, beta } = await openBlobTable();
+    const [row] = await table
+      .query()
+      .where("id = 2")
+      .select(
+        new Map([
+          ["key", "id"],
+          ["picture", "image"],
+        ]),
+      )
+      .toArray({ blobMode: "bytes" });
+    expect(Number(row.key)).toBe(2);
+    expect(Buffer.from(row.picture)).toEqual(beta);
+  });
+
+  it("returns blob bytes from a take query", async () => {
+    const { table, rowIds, alpha, beta } = await openBlobTable();
+    const rows = await table
+      .takeRowIds([rowIds[1], rowIds[0]])
+      .toArray({ blobMode: "bytes" });
+    const byId = new Map(rows.map((row) => [Number(row.id), row]));
+    expect(Buffer.from(byId.get(1)!.image)).toEqual(alpha);
+    expect(Buffer.from(byId.get(2)!.image)).toEqual(beta);
+    expect(Object.keys(rows[0])).not.toContain("_rowid");
+  });
+
+  it("does not change later executions of the same query", async () => {
+    const { table } = await openBlobTable();
+    const query = table.query().where("id = 1");
+    await query.toArray({ blobMode: "bytes" });
+    const [row] = await query.toArray();
+    expect(Object.keys(row)).not.toContain("_rowid");
+    expect(Number(row.image.size)).toBe(5);
+  });
+
+  it("skips row ids when no blob column is selected", async () => {
+    const { table } = await openBlobTable();
+    const result = await table
+      .query()
+      .select(["id"])
+      .toArrow({ blobMode: "bytes" });
+    expect(result.schema.fields.map((f) => f.name)).toEqual(["id"]);
+    expect(result.numRows).toBe(3);
+  });
+
+  it("returns an empty bytes column for an empty result", async () => {
+    const { table } = await openBlobTable();
+    const result = await table
+      .query()
+      .where("id > 100")
+      .toArrow({ blobMode: "bytes" });
+    expect(result.numRows).toBe(0);
+    expect(result.schema.fields.map((f) => f.name)).toEqual(["id", "image"]);
+    const field = result.schema.fields.find((f) => f.name === "image")!;
+    expect(arrow.DataType.isLargeBinary(field.type)).toBe(true);
+    expect(isBlobField(field)).toBe(false);
+  });
+
+  it("keeps the descriptor schema for an empty result", async () => {
+    const { table } = await openBlobTable();
+    const result = await table.query().where("id > 100").toArrow();
+    expect(result.numRows).toBe(0);
+    expect(result.schema.fields.map((f) => f.name)).toEqual(["id", "image"]);
+    const field = result.schema.fields.find((f) => f.name === "image")!;
+    expect(arrow.DataType.isStruct(field.type)).toBe(true);
+  });
+
+  it("rejects an unknown blobMode", async () => {
+    const { table } = await openBlobTable();
+    await expect(
+      table.query().toArray({ blobMode: "lazy" as BlobMode }),
+    ).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("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("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("rejects fetchBlobs on a non-blob column", async () => {
     const { table, rowIds } = await openBlobTable();
     await expect(table.fetchBlobs("id", rowIds)).rejects.toThrow(/blob/i);
@@ -3198,6 +3355,32 @@ describe("when dealing with blob columns", () => {
     ).map((value) => (value == null ? null : Number(value.size)));
   }
 
+  async function openVersionedBlobTable() {
+    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"),
+    ]);
+    const table = await db.createTable(
+      "versions",
+      [{ id: 1n, generation: 1n, image: Buffer.from("old") }],
+      { schema },
+    );
+    const native = (
+      table as unknown as { inner: { version: () => Promise<number> } }
+    ).inner;
+    const update = async (generation: bigint, image: string) => {
+      await table
+        .mergeInsert("id")
+        .whenMatchedUpdateAll()
+        .execute([{ id: 1n, generation, image: Buffer.from(image) }]);
+    };
+    return { table, native, update };
+  }
+
   async function openBlobTable() {
     const db = await connect(tmpDir.name);
     const schema = new Schema([

nodejs/lancedb/index.ts

@@ -122,6 +122,7 @@ export {
   VectorQuery,
   TakeQuery,
   AnalyzePlanDistributedMetrics,
+  BlobMode,
   QueryExecutionOptions,
   ColumnOrdering,
   FullTextSearchOptions,

nodejs/lancedb/query.ts

@@ -3,12 +3,18 @@
 
 import {
   Table as ArrowTable,
+  Field,
   type IntoVector,
+  LargeBinary,
   RecordBatch,
+  Schema,
+  Struct,
   extractVectorBuffer,
   fromBufferToRecordBatch,
   fromRecordBatchToBuffer,
+  makeData,
   tableFromIPC,
+  vectorFromArray,
 } from "./arrow";
 import { type IvfPqOptions } from "./indices";
 import {
@@ -41,6 +47,112 @@ export async function* RecordBatchIterator(
   }
 }
 
+/**
+ * How a query returns blob v2 columns. See {@link QueryExecutionOptions.blobMode}.
+ */
+export type BlobMode = "descriptions" | "bytes";
+
+/**
+ * The table a query reads, and the parts of the query that decide which blob
+ * columns `blobMode: "bytes"` fills in.
+ *
+ * @hidden
+ */
+export interface QueryBlobContext {
+  table: NativeTable;
+  withRowId: boolean;
+  projection?: [string, string][];
+}
+
+const BLOB_BYTES_ATTEMPTS = 3;
+
+function blobSources(
+  blobColumns: string[],
+  projection: [string, string][] | undefined,
+): Map<string, string> {
+  const columns = new Set(blobColumns);
+  const pairs =
+    projection ??
+    blobColumns.map((column): [string, string] => [column, column]);
+  return new Map(pairs.filter(([, source]) => columns.has(source)));
+}
+
+async function readBlobBytes(
+  table: ArrowTable,
+  sources: Map<string, string>,
+  snapshot: NativeTable,
+  dropRowId: boolean,
+): Promise<ArrowTable> {
+  const rowIdColumn = table.getChild("_rowid");
+  if (table.numRows > 0 && rowIdColumn === null) {
+    throw new Error('blobMode "bytes" needs _rowid in the query result');
+  }
+  const rowIds =
+    rowIdColumn === null
+      ? []
+      : Array.from(rowIdColumn.toArray() as BigUint64Array);
+  const replacements = new Map<number, (Buffer | null)[]>();
+  for (const [output, source] of sources) {
+    const index = table.schema.fields.findIndex(
+      (field) => field.name === output,
+    );
+    if (index < 0) {
+      continue;
+    }
+    const values =
+      rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds);
+    replacements.set(
+      index,
+      values.map((value) => value ?? null),
+    );
+  }
+  const result = withBinaryColumns(table, replacements);
+  if (!dropRowId) {
+    return result;
+  }
+  return result.select(
+    result.schema.fields
+      .map((field) => field.name)
+      .filter((name) => name !== "_rowid"),
+  );
+}
+
+// A fresh field, not a clone: the descriptor field's metadata must not follow
+// the bytes.
+function withBinaryColumns(
+  table: ArrowTable,
+  replacements: Map<number, (Buffer | null)[]>,
+): ArrowTable {
+  if (replacements.size === 0) {
+    return table;
+  }
+  const fields = table.schema.fields.map((field, index) =>
+    replacements.has(index)
+      ? new Field(field.name, new LargeBinary(), true)
+      : field,
+  );
+  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) {
+        return child;
+      }
+      return vectorFromArray(
+        values.slice(offset, offset + batch.numRows),
+        new LargeBinary(),
+      ).data[0];
+    });
+    offset += batch.numRows;
+    return new RecordBatch(
+      schema,
+      makeData({ type: new Struct(fields), length: batch.numRows, children }),
+    );
+  });
+  return new ArrowTable(schema, batches);
+}
+
 /**
  * Options that control the behavior of a particular query execution
  */
@@ -57,6 +169,32 @@ export interface QueryExecutionOptions {
    * Timeout for query execution in milliseconds
    */
   timeoutMs?: number;
+
+  /**
+   * How to return blob v2 columns. Applies to `toArrow()` and `toArray()`.
+   *
+   * - `"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.
+   *
+   * Blobs nested in a struct or 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.
+   *
+   * @example
+   * ```ts
+   * const rows = await table
+   *   .query()
+   *   .select(["id", "image"])
+   *   .toArray({ blobMode: "bytes" });
+   * const image: Uint8Array | null = rows[0].image;
+   * ```
+   */
+  blobMode?: BlobMode;
 }
 
 export type AnalyzePlanDistributedMetrics = "aggregate" | "per_worker" | "full";
@@ -119,10 +257,19 @@ export class QueryBase<
   /**
    * @hidden
    */
-  protected constructor(inner?: NativeQueryType | Promise<NativeQueryType>) {
+  protected blobContext?: QueryBlobContext;
+
+  /**
+   * @hidden
+   */
+  protected constructor(
+    inner?: NativeQueryType | Promise<NativeQueryType>,
+    blobContext?: QueryBlobContext,
+  ) {
     if (inner !== undefined) {
       this.inner = inner;
     }
+    this.blobContext = blobContext;
   }
 
   // call a function on the inner (either a promise or the actual object)
@@ -193,14 +340,22 @@ export class QueryBase<
       });
     };
 
+    let projection: [string, string][];
     if (typeof columns === "string") {
       selectColumns([columns]);
+      projection = [[columns, columns]];
     } else if (Array.isArray(columns)) {
       selectColumns(columns);
+      projection = columns.map((column) => [column, column]);
     } else if (columns instanceof Map) {
-      selectMapping(Array.from(columns.entries()));
+      projection = Array.from(columns.entries());
+      selectMapping(projection);
     } else {
-      selectMapping(Object.entries(columns));
+      projection = Object.entries(columns);
+      selectMapping(projection);
+    }
+    if (this.blobContext !== undefined) {
+      this.blobContext = { ...this.blobContext, projection };
     }
 
     return this;
@@ -215,6 +370,9 @@ export class QueryBase<
    */
   withRowId(): this {
     this.doCall((inner: NativeQueryType) => inner.withRowId());
+    if (this.blobContext !== undefined) {
+      this.blobContext = { ...this.blobContext, withRowId: true };
+    }
     return this;
   }
 
@@ -223,9 +381,14 @@ export class QueryBase<
    */
   protected async nativeExecute(
     options?: Partial<QueryExecutionOptions>,
+    withRowId = false,
   ): Promise<NativeBatchIterator> {
     const inner = await this.getInner();
-    return inner.execute(options?.maxBatchLength, options?.timeoutMs);
+    return inner.execute(
+      options?.maxBatchLength,
+      options?.timeoutMs,
+      withRowId,
+    );
   }
 
   /**
@@ -257,8 +420,41 @@ export class QueryBase<
    * Empty results retain the query's output schema.
    */
   async toArrow(options?: Partial<QueryExecutionOptions>): Promise<ArrowTable> {
+    const context = this.blobBytesContext(options?.blobMode);
+    if (context === undefined) {
+      return this.collect(options, false);
+    }
+    for (let attempt = 1; ; attempt++) {
+      // The snapshot pins the bytes. The query runs on the live table, so an
+      // 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,
+      );
+      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);
+      }
+      if (attempt === BLOB_BYTES_ATTEMPTS) {
+        throw new Error(
+          `blobMode "bytes": the table changed during each of ${BLOB_BYTES_ATTEMPTS} attempts to read it`,
+        );
+      }
+    }
+  }
+
+  private async collect(
+    options: Partial<QueryExecutionOptions> | undefined,
+    withRowId: boolean,
+  ): Promise<ArrowTable> {
     const batches = [];
-    const iterator = this.nativeExecute(options);
+    const iterator = this.nativeExecute(options, withRowId);
     for await (const batch of RecordBatchIterator(iterator)) {
       batches.push(batch);
     }
@@ -268,6 +464,23 @@ export class QueryBase<
     return new ArrowTable(batches);
   }
 
+  private blobBytesContext(
+    mode: BlobMode | undefined,
+  ): QueryBlobContext | undefined {
+    if (mode === undefined || mode === "descriptions") {
+      return undefined;
+    }
+    if (mode !== "bytes") {
+      throw new Error(
+        `blobMode must be "descriptions" or "bytes", got "${String(mode)}"`,
+      );
+    }
+    if (this.blobContext === undefined) {
+      throw new Error('blobMode "bytes" is not supported for this query');
+    }
+    return this.blobContext;
+  }
+
   /** Collect the results as an array of objects. */
   // biome-ignore lint/suspicious/noExplicitAny: arrow.toArrow() returns any[]
   async toArray(options?: Partial<QueryExecutionOptions>): Promise<any[]> {
@@ -355,8 +568,11 @@ export class StandardQueryBase<
   extends QueryBase<NativeQueryType>
   implements ExecutableQuery
 {
-  constructor(inner?: NativeQueryType | Promise<NativeQueryType>) {
-    super(inner);
+  constructor(
+    inner?: NativeQueryType | Promise<NativeQueryType>,
+    blobContext?: QueryBlobContext,
+  ) {
+    super(inner, blobContext);
   }
 
   /**
@@ -505,8 +721,11 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
   /**
    * @hidden
    */
-  constructor(inner: NativeVectorQuery | Promise<NativeVectorQuery>) {
-    super(inner);
+  constructor(
+    inner: NativeVectorQuery | Promise<NativeVectorQuery>,
+    blobContext?: QueryBlobContext,
+  ) {
+    super(inner, blobContext);
   }
 
   /**
@@ -760,7 +979,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
         addQueryVectorToNative(inner, outcome.value);
         return inner;
       })();
-      return new VectorQuery(res);
+      return new VectorQuery(res, this.blobContext);
     } else {
       this.doVectorCall((inner) => addQueryVectorToNative(inner, vector));
       return this;
@@ -849,7 +1068,7 @@ export function createAutoQuery(
     return nearestToNative(route.table.query(), vector);
   };
 
-  return new AutoQuery(createInner);
+  return new AutoQuery(createInner, { table, withRowId: false });
 }
 
 /**
@@ -858,8 +1077,8 @@ export function createAutoQuery(
  * @hideconstructor
  */
 export class TakeQuery extends QueryBase<NativeTakeQuery> {
-  constructor(inner: NativeTakeQuery) {
-    super(inner);
+  constructor(inner: NativeTakeQuery, blobContext?: QueryBlobContext) {
+    super(inner, blobContext);
   }
 
   /**
@@ -898,8 +1117,9 @@ export class AutoQuery extends StandardQueryBase<
     private readonly createInner: () => Promise<
       NativeQuery | NativeVectorQuery
     >,
+    blobContext?: QueryBlobContext,
   ) {
-    super();
+    super(undefined, blobContext);
   }
 
   /** @hidden */
@@ -933,7 +1153,7 @@ export class Query extends StandardQueryBase<NativeQuery> {
    * @hidden
    */
   constructor(tbl: NativeTable) {
-    super(tbl.query());
+    super(tbl.query(), { table: tbl, withRowId: false });
   }
 
   /**
@@ -979,14 +1199,15 @@ export class Query extends StandardQueryBase<NativeQuery> {
       const nativeQuery = inner.then(async (resolvedInner) =>
         nearestToNative(resolvedInner, await vector),
       );
-      return new VectorQuery(nativeQuery);
+      return new VectorQuery(nativeQuery, this.blobContext);
     }
     if (vector instanceof Promise) {
       return new VectorQuery(
         vector.then((resolvedVector) => nearestToNative(inner, resolvedVector)),
+        this.blobContext,
       );
     }
-    return new VectorQuery(nearestToNative(inner, vector));
+    return new VectorQuery(nearestToNative(inner, vector), this.blobContext);
   }
 
   nearestToText(query: string | FullTextQuery, columns?: string[]): Query {

nodejs/lancedb/table.ts

@@ -1240,11 +1240,17 @@ export class LocalTable extends Table {
   }
 
   takeOffsets(offsets: number[]): TakeQuery {
-    return new TakeQuery(this.inner.takeOffsets(offsets));
+    return new TakeQuery(this.inner.takeOffsets(offsets), {
+      table: this.inner,
+      withRowId: false,
+    });
   }
 
   takeRowIds(rowIds: readonly (bigint | number)[]): TakeQuery {
-    return new TakeQuery(this.inner.takeRowIds(rowIdsToBigInts(rowIds)));
+    return new TakeQuery(this.inner.takeRowIds(rowIdsToBigInts(rowIds)), {
+      table: this.inner,
+      withRowId: false,
+    });
   }
 
   blobColumns(): Promise<string[]> {

nodejs/src/query.rs

@@ -196,6 +196,7 @@ impl Query {
         &self,
         max_batch_length: Option<u32>,
         timeout_ms: Option<u32>,
+        with_row_id: Option<bool>,
     ) -> napi::Result<RecordBatchIterator> {
         let mut execution_opts = QueryExecutionOptions::default();
         if let Some(max_batch_length) = max_batch_length {
@@ -204,8 +205,12 @@ impl Query {
         if let Some(timeout_ms) = timeout_ms {
             execution_opts.timeout = Some(std::time::Duration::from_millis(timeout_ms as u64))
         }
-        let inner_stream = self
-            .inner
+        let query = if with_row_id == Some(true) {
+            self.inner.clone().with_row_id()
+        } else {
+            self.inner.clone()
+        };
+        let inner_stream = query
             .execute_with_options(execution_opts)
             .await
             .map_err(|e| {
@@ -432,6 +437,7 @@ impl VectorQuery {
         &self,
         max_batch_length: Option<u32>,
         timeout_ms: Option<u32>,
+        with_row_id: Option<bool>,
     ) -> napi::Result<RecordBatchIterator> {
         let mut execution_opts = QueryExecutionOptions::default();
         if let Some(max_batch_length) = max_batch_length {
@@ -440,8 +446,12 @@ impl VectorQuery {
         if let Some(timeout_ms) = timeout_ms {
             execution_opts.timeout = Some(std::time::Duration::from_millis(timeout_ms as u64))
         }
-        let inner_stream = self
-            .inner
+        let query = if with_row_id == Some(true) {
+            self.inner.clone().with_row_id()
+        } else {
+            self.inner.clone()
+        };
+        let inner_stream = query
             .execute_with_options(execution_opts)
             .await
             .map_err(|e| {
@@ -520,6 +530,7 @@ impl TakeQuery {
         &self,
         max_batch_length: Option<u32>,
         timeout_ms: Option<u32>,
+        with_row_id: Option<bool>,
     ) -> napi::Result<RecordBatchIterator> {
         let mut execution_opts = QueryExecutionOptions::default();
         if let Some(max_batch_length) = max_batch_length {
@@ -528,8 +539,12 @@ impl TakeQuery {
         if let Some(timeout_ms) = timeout_ms {
             execution_opts.timeout = Some(std::time::Duration::from_millis(timeout_ms as u64))
         }
-        let inner_stream = self
-            .inner
+        let query = if with_row_id == Some(true) {
+            self.inner.clone().with_row_id()
+        } else {
+            self.inner.clone()
+        };
+        let inner_stream = query
             .execute_with_options(execution_opts)
             .await
             .map_err(|e| {