From b301a47aa8b4ed0e1bad6ae7706029b376f9a107 Mon Sep 17 00:00:00 2001 From: Bilal Atique <74999164+BilalAtique@users.noreply.github.com> Date: Sat, 17 Oct 2026 00:50:38 +0500 Subject: [PATCH] feat(node): add blobMode to return blob bytes from queries (#4353) --- docs/src/js/globals.md | 1 + .../js/interfaces/QueryExecutionOptions.md | 32 +++ docs/src/js/type-aliases/BlobMode.md | 13 + nodejs/__test__/table.test.ts | 183 +++++++++++++ nodejs/lancedb/index.ts | 1 + nodejs/lancedb/query.ts | 255 ++++++++++++++++-- nodejs/lancedb/table.ts | 10 +- nodejs/src/query.rs | 27 +- 8 files changed, 497 insertions(+), 25 deletions(-) create mode 100644 docs/src/js/type-aliases/BlobMode.md diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 739b9033a..7310f3f90 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -161,6 +161,7 @@ - [BaseTokenizer](type-aliases/BaseTokenizer.md) - [BlobData](type-aliases/BlobData.md) - [BlobInput](type-aliases/BlobInput.md) +- [BlobMode](type-aliases/BlobMode.md) - [BlobOptions](type-aliases/BlobOptions.md) - [BlobRange](type-aliases/BlobRange.md) - [BlobReadOptions](type-aliases/BlobReadOptions.md) diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index 46b06ce7e..8d3749f38 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -10,6 +10,38 @@ Options that control the behavior of a particular query execution ## Properties +### blobMode? + +```ts +optional blobMode: BlobMode; +``` + +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. + +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; +``` + +*** + ### maxBatchLength? ```ts diff --git a/docs/src/js/type-aliases/BlobMode.md b/docs/src/js/type-aliases/BlobMode.md new file mode 100644 index 000000000..67d7c47ad --- /dev/null +++ b/docs/src/js/type-aliases/BlobMode.md @@ -0,0 +1,13 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BlobMode + +# Type Alias: BlobMode + +```ts +type BlobMode: "descriptions" | "bytes"; +``` + +How a query returns blob v2 columns. See [QueryExecutionOptions.blobMode](../interfaces/QueryExecutionOptions.md#blobmode). diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index f53f859b8..b30121725 100644 --- a/nodejs/__test__/table.test.ts +++ b/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 } } + ).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([ diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 3254a0742..3096cbe6b 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -122,6 +122,7 @@ export { VectorQuery, TakeQuery, AnalyzePlanDistributedMetrics, + BlobMode, QueryExecutionOptions, ColumnOrdering, FullTextSearchOptions, diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 3583e5992..ef7443bbb 100644 --- a/nodejs/lancedb/query.ts +++ b/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 { + 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, + snapshot: NativeTable, + dropRowId: boolean, +): Promise { + 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(); + 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, +): 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) { + protected blobContext?: QueryBlobContext; + + /** + * @hidden + */ + protected constructor( + inner?: NativeQueryType | Promise, + 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, + withRowId = false, ): Promise { 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): Promise { + 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 | undefined, + withRowId: boolean, + ): Promise { 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): Promise { @@ -355,8 +568,11 @@ export class StandardQueryBase< extends QueryBase implements ExecutableQuery { - constructor(inner?: NativeQueryType | Promise) { - super(inner); + constructor( + inner?: NativeQueryType | Promise, + blobContext?: QueryBlobContext, + ) { + super(inner, blobContext); } /** @@ -505,8 +721,11 @@ export class VectorQuery extends StandardQueryBase { /** * @hidden */ - constructor(inner: NativeVectorQuery | Promise) { - super(inner); + constructor( + inner: NativeVectorQuery | Promise, + blobContext?: QueryBlobContext, + ) { + super(inner, blobContext); } /** @@ -760,7 +979,7 @@ export class VectorQuery extends StandardQueryBase { 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 { - 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 { * @hidden */ constructor(tbl: NativeTable) { - super(tbl.query()); + super(tbl.query(), { table: tbl, withRowId: false }); } /** @@ -979,14 +1199,15 @@ export class Query extends StandardQueryBase { 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 { diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 4d6500a5f..2b5319bec 100644 --- a/nodejs/lancedb/table.ts +++ b/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 { diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index 68db6f810..dba8f665a 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -196,6 +196,7 @@ impl Query { &self, max_batch_length: Option, timeout_ms: Option, + with_row_id: Option, ) -> napi::Result { 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, timeout_ms: Option, + with_row_id: Option, ) -> napi::Result { 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, timeout_ms: Option, + with_row_id: Option, ) -> napi::Result { 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| {