From a32866fb5879b1339955ed259f69002585cf3e52 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 20:51:07 +0000 Subject: [PATCH 1/2] feat(bigquery): support destination tables in createQueryArrowStream Read results of jobs.insert queries with an explicit destination table through BigQueryReadClient.createReadSession (dataFormat: ARROW, maxStreamCount: 1) once the job completes, instead of the job's default stream, which only serves anonymous destination tables. An empty destination table returns no streams, so emit one 0-row RecordBatch with the session's Arrow schema. --- handwritten/bigquery/src/arrow.ts | 121 +++++++++--- handwritten/bigquery/system-test/arrow.ts | 49 +++++ handwritten/bigquery/test/arrow.ts | 215 +++++++++++++++++++++- 3 files changed, 350 insertions(+), 35 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 9358e853e86..4c475d651ed 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -370,6 +370,32 @@ async function streamDefaultJobStream( pushBatches(stream, decodeArrowPayload(arrow, schema)); } +/** + * Creates an Arrow read session with a single stream on `table`, billed to + * `parentProjectId`. + * + * @private + */ +async function createArrowReadSession( + ctx: ArrowStreamContext, + parentProjectId: string | undefined, + table: bigquery.ITableReference, +): Promise { + const [session] = await getStorageReadClient( + ctx.bq, + ctx.storage, + ).createReadSession({ + parent: `projects/${parentProjectId}`, + readSession: { + table: `projects/${table.projectId}/datasets/${table.datasetId}/tables/${table.tableId}`, + dataFormat: 'ARROW', + }, + // One stream keeps rows in the query's order (e.g., ORDER BY). + maxStreamCount: 1, + }); + return session ?? undefined; +} + /** * Reads the Arrow schema of a finished query job's results from its * destination table. The default job stream sends no `ReadRowsResponse`, and @@ -383,7 +409,6 @@ async function readDestinationSchema( ctx: ArrowStreamContext, job: Job, ): Promise { - const {bq, storage} = ctx; // A script's destination table is only set once the job finishes, so read // the job's current metadata instead of the response that created it. const [metadata] = await job.getMetadata(); @@ -391,18 +416,15 @@ async function readDestinationSchema( if (!destRef) { return undefined; } - const [session] = await getStorageReadClient(bq, storage).createReadSession({ - parent: `projects/${resolveProjectId( - bq, + const session = await createArrowReadSession( + ctx, + resolveProjectId( + ctx.bq, metadata?.jobReference?.projectId, destRef.projectId, - )}`, - readSession: { - table: `projects/${destRef.projectId}/datasets/${destRef.datasetId}/tables/${destRef.tableId}`, - dataFormat: 'ARROW', - }, - maxStreamCount: 1, - }); + ), + destRef, + ); return session?.arrowSchema?.serializedSchema || undefined; } @@ -517,6 +539,45 @@ function streamReadRows( }); } +/** + * Creates a Storage Read API session for an explicit destination table and + * streams Arrow `RecordBatch`es from its single stream. + * + * @private + */ +async function streamDestinationTable( + ctx: ArrowStreamContext, + jobRef: bigquery.IJobReference | undefined, + table: bigquery.ITableReference, +): Promise { + const {arrow, stream} = ctx; + // The session is billed to the job's project; the table can be in another. + const session = await createArrowReadSession( + ctx, + resolveProjectId(ctx.bq, jobRef?.projectId), + table, + ); + const sessionSchema = session?.arrowSchema?.serializedSchema || undefined; + const streams = session?.streams ?? []; + const readStream = streams[0]?.name; + if (streams.length === 0 || !readStream) { + // An empty destination table returns 0 streams; emit a single 0-row + // RecordBatch carrying the session's Arrow schema. + pushBatches(stream, decodeArrowPayload(arrow, sessionSchema)); + return; + } + + const {emittedBatch, serializedSchema} = await streamReadRows(ctx, { + readStream, + offset: 0, + initialSchema: sessionSchema, + }); + // A stream that ends without rows still yields a 0-row RecordBatch. + if (!emittedBatch) { + pushBatches(stream, decodeArrowPayload(arrow, serializedSchema)); + } +} + /** * Waits for an unfinished BigQuery job to complete using `job.getQueryResults` * without fetching any `STRUCT_ENCODING` rows (`maxResults: 0`). Returns early @@ -549,10 +610,11 @@ async function waitForJobCompletion( } /** - * Waits for `job` to finish, then streams all of its rows from the default job - * stream. An unfinished jobs.query and jobs.insert return no rows or schema, - * so this reads every row from offset 0 and passes the job to look up the - * schema of an empty result. + * Waits for `job` to finish, then streams all of its rows: from + * `destinationTable` through a read session when the caller set one, or from + * the default job stream otherwise. An unfinished jobs.query and jobs.insert + * return no rows or schema, so the default job stream is read from offset 0, + * passing the job to look up the schema of an empty result. * * @private */ @@ -561,11 +623,16 @@ async function streamJobWhenComplete( job: Job, jobRef: bigquery.IJobReference | undefined, timeoutMs?: number, + destinationTable?: bigquery.ITableReference, ): Promise { await waitForJobCompletion(job, ctx.stream, timeoutMs); if (ctx.stream.destroyed) { return; } + if (destinationTable) { + await streamDestinationTable(ctx, jobRef, destinationTable); + return; + } await streamDefaultJobStream(ctx, jobRef, 0, undefined, job); } @@ -601,7 +668,18 @@ async function streamQueryJob( return; } - await streamJobWhenComplete(ctx, job, resp?.jobReference, timeoutMs); + // BigQuery echoes the caller's destination table, with its resolved + // project, in the job's configuration. + const destinationTable = queryObj.destination + ? (resp?.configuration?.query?.destinationTable ?? undefined) + : undefined; + await streamJobWhenComplete( + ctx, + job, + resp?.jobReference, + timeoutMs, + destinationTable, + ); } /** @@ -641,13 +719,6 @@ export function createQueryArrowStream_( ); } - if (queryObj.destination) { - // TODO: Support explicit destination tables via CreateReadSession. - throw new Error( - 'Explicit destination tables are not yet supported in createQueryArrowStream.', - ); - } - const ctx: ArrowStreamContext = {bq, arrow, storage, stream}; readCtx = ctx; @@ -657,8 +728,8 @@ export function createQueryArrowStream_( }); if (!queryReq) { // buildQueryRequest_ returns undefined when the query sets an option - // that only jobs.insert accepts (e.g., jobId, dryRun, or jobTimeoutMs), - // or when FAST_QUERY_PATH=DISABLED. + // that only jobs.insert accepts (e.g., destination, jobId, dryRun, or + // jobTimeoutMs), or when FAST_QUERY_PATH=DISABLED. await streamQueryJob(ctx, queryObj); return; } diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts index e801ee4c039..2a081a5bae5 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -632,4 +632,53 @@ describe('Arrow Query Results', () => { assert.ok(arrow.DataType.isInt(batches[0].schema.fields[0].type)); assert.ok(arrow.DataType.isUtf8(batches[0].schema.fields[1].type)); }); + + it('should stream Arrow results from an explicit destination table via CreateReadSession with default writeDisposition and WRITE_TRUNCATE', async () => { + const destTable = dataset.table('dest_results'); + + // 1. Default writeDisposition (writes to a new destination table). + const initialBatches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({ + query: + 'SELECT x AS id, CONCAT("row_", CAST(x AS STRING)) AS label FROM UNNEST(GENERATE_ARRAY(1, 100)) AS x ORDER BY x', + destination: destTable, + }) + .toArray(); + + const initialTable = new arrow.Table(initialBatches); + assert.strictEqual(initialTable.numRows, 100); + assert.strictEqual(initialTable.getChild('id')!.get(0), 1n); + assert.strictEqual(initialTable.getChild('id')!.get(99), 100n); + + // 2. WRITE_TRUNCATE overwrites the destination table and reads only the new rows. + const truncatedBatches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({ + query: + 'SELECT x AS id, CONCAT("new_", CAST(x AS STRING)) AS label FROM UNNEST(GENERATE_ARRAY(1, 10)) AS x ORDER BY x', + destination: destTable, + writeDisposition: 'WRITE_TRUNCATE', + }) + .toArray(); + + const truncatedTable = new arrow.Table(truncatedBatches); + assert.strictEqual(truncatedTable.numRows, 10); + assert.strictEqual(truncatedTable.getChild('id')!.get(0), 1n); + assert.strictEqual(truncatedTable.getChild('label')!.get(0), 'new_1'); + + // 3. 0-row query into an explicit destination table preserves the schema. + const emptyBatches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({ + query: 'SELECT 1 AS id, "empty" AS label FROM UNNEST([1]) WHERE FALSE', + destination: destTable, + writeDisposition: 'WRITE_TRUNCATE', + }) + .toArray(); + + assert.strictEqual(emptyBatches.length, 1); + assert.strictEqual(emptyBatches[0].numRows, 0); + assert.deepStrictEqual( + emptyBatches[0].schema.fields.map(f => f.name), + ['id', 'label'], + ); + }); }); diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 260d63cefe0..fc2416e4689 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -525,16 +525,6 @@ describe('Arrow Support', () => { ); }); - it('should fail when an explicit destination table is specified', async () => { - const destination = bq.dataset('ds').table('tbl'); - - // TODO: Replace with explicit destination table (CreateReadSession) tests once added. - await assert.rejects( - bq.createQueryArrowStream({query: 'SELECT 1', destination}).toArray(), - /Explicit destination tables are not yet supported in createQueryArrowStream\./, - ); - }); - it('should forward errors thrown while building the jobs.query request', async () => { // A null parameter without an explicit type makes buildQueryRequest_ throw. await assert.rejects( @@ -1625,4 +1615,209 @@ describe('Arrow Support', () => { ); }); }); + + describe('BigQuery#createQueryArrowStream (destination table)', () => { + let bq: BigQuery; + + beforeEach(() => { + bq = new BigQuery({projectId: 'test-project'}); + }); + + function stubDestinationJob( + tableId: string, + jobId: string, + opts?: { + destinationTable?: bigquery.ITableReference; + onGetQueryResults?: (reqOpts: QueryResultsOptions) => void; + }, + ) { + const destination = bq.dataset('dest_ds').table(tableId); + let capturedJobOpts: JobOptions | undefined; + + const fakeJob = bq.job(jobId, {location: 'US'}); + fakeJob.metadata = { + jobReference: { + projectId: 'test-project', + jobId, + location: 'US', + }, + // jobs.insert echoes the destination table in the job configuration. + configuration: { + query: { + destinationTable: opts?.destinationTable ?? { + projectId: 'test-project', + datasetId: 'dest_ds', + tableId, + }, + }, + }, + }; + stubJobGetQueryResults(fakeJob, (reqOpts, cb) => { + opts?.onGetQueryResults?.(reqOpts); + process.nextTick(() => { + cb(null, [], null, {jobComplete: true}); + }); + }); + stubCreateJob(bq, (reqOpts, cb) => { + capturedJobOpts = reqOpts; + process.nextTick(() => { + cb(null, fakeJob, fakeJob.metadata); + }); + }); + + return { + destination, + getCapturedJobOpts: () => capturedJobOpts, + }; + } + + it('should preserve the Table instance through createQueryJob, call createReadSession with maxStreamCount = 1 and dataFormat = ARROW, and stream from streams[0].name', async () => { + const capturedSessionReqs: CreateReadSessionRequest[] = []; + const capturedReadRowsReqs: Array<{ + readStream?: string; + offset?: number; + }> = []; + + const {destination, getCapturedJobOpts} = stubDestinationJob( + 'dest_tbl', + 'dest-job-1', + { + destinationTable: { + projectId: 'dest-project', + datasetId: 'dest_ds', + tableId: 'dest_tbl', + }, + onGetQueryResults: opts => { + // Verify createReadSession has not been called before the destination job completes. + assert.strictEqual(capturedSessionReqs.length, 0); + assert.strictEqual(opts.maxResults, 0); + }, + }, + ); + + stubStorageReadClient(bq, { + createReadSession: async req => { + capturedSessionReqs.push(req); + return [ + { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + streams: [ + { + name: 'projects/test-project/locations/US/sessions/sess-1/streams/stream-0', + }, + ], + }, + ]; + }, + readRows: (req: {readStream?: string; offset?: number}) => { + capturedReadRowsReqs.push(req); + // Omit arrowSchema on the ReadRows message to verify session.arrowSchema is used as initialSchema. + return fakeReadRowsStream([ + {arrowRecordBatch: SAMPLE_READ_ROWS_RESPONSE.arrowRecordBatch}, + ]); + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream({ + query: 'SELECT id, name FROM tbl', + destination, + writeDisposition: 'WRITE_TRUNCATE', + }) + .toArray(); + + // Verify createQueryJob translated destination into configuration.query.destinationTable. + const queryConfig = getCapturedJobOpts()?.configuration?.query; + assert.deepStrictEqual(queryConfig?.destinationTable, { + datasetId: 'dest_ds', + projectId: 'test-project', + tableId: 'dest_tbl', + }); + assert.strictEqual(queryConfig?.writeDisposition, 'WRITE_TRUNCATE'); + + // Verify createReadSession was called with maxStreamCount: 1 and dataFormat: 'ARROW'. + assert.deepStrictEqual(capturedSessionReqs, [ + { + parent: 'projects/test-project', + readSession: { + table: 'projects/dest-project/datasets/dest_ds/tables/dest_tbl', + dataFormat: 'ARROW', + }, + maxStreamCount: 1, + }, + ]); + + // Verify readRows was called on session.streams[0].name at offset = 0. + assert.deepStrictEqual(capturedReadRowsReqs, [ + { + readStream: + 'projects/test-project/locations/US/sessions/sess-1/streams/stream-0', + offset: 0, + }, + ]); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + + it('should emit a 0-row RecordBatch preserving session.arrowSchema when createReadSession returns no streams', async () => { + const {destination} = stubDestinationJob('empty_tbl', 'dest-empty-1'); + + let readRowsCalled = false; + stubStorageReadClient(bq, { + createReadSession: async () => [ + { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + streams: [], + }, + ], + readRows: () => { + readRowsCalled = true; + return new EventEmitter(); + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream({ + query: 'SELECT id, name FROM tbl WHERE FALSE', + destination, + }) + .toArray(); + + // Verify an empty destination table emits a single 0-row RecordBatch with the session's schema without calling readRows. + assert.strictEqual(readRowsCalled, false); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.deepStrictEqual( + batches[0].schema.fields.map(f => f.name), + ['id', 'name'], + ); + }); + + it('should forward createReadSession errors to the stream', async () => { + const {destination} = stubDestinationJob('err_tbl', 'dest-err-1'); + + stubStorageReadClient(bq, { + createReadSession: async () => { + throw new Error('bigquery.readsessions.create denied'); + }, + readRows: () => new EventEmitter(), + }); + + // Verify errors from createReadSession reject stream consumption. + await assert.rejects( + bq.createQueryArrowStream({query: 'SELECT 1', destination}).toArray(), + /bigquery\.readsessions\.create denied/, + ); + }); + }); }); From 7881f36490b9d653cd8b1b97e951a9183f3e48e9 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 21 Oct 2026 00:52:02 +0000 Subject: [PATCH 2/2] fix(bigquery): guard Arrow read session project; document destinations Throw a clear error instead of creating a read session with parent "projects/undefined" when no project ID can be resolved. Note in the createQueryArrowStream JSDoc that an explicit destination table is read in full through a read session after the query finishes. --- handwritten/bigquery/src/arrow.ts | 5 ++++ handwritten/bigquery/src/bigquery.ts | 3 ++ handwritten/bigquery/test/arrow.ts | 45 ++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index c9045d6d18e..7dec8fee610 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -378,6 +378,11 @@ async function createArrowReadSession( parentProjectId: string | undefined, table: bigquery.ITableReference, ): Promise { + if (!parentProjectId) { + throw new Error( + 'Cannot create an Arrow read session without a project ID.', + ); + } const [session] = await getStorageReadClient( ctx.bq, ctx.storage, diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index da52ab6e0ea..57a47762fbe 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -401,6 +401,9 @@ export class BigQuery extends Service { * ends, call `stream.destroy()` to cancel the read; breaking out of * `for await` does this for you. * + * With an explicit `destination` table, the stream reads that whole table + * through a Storage Read API session once the query finishes. + * * @param {string|ArrowQuery} [options] A string SQL query or configuration * object. See {@link BigQuery#query} for a complete list of options. * @param {number} [options.timeoutMs] How long to wait for the query to diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index cc23c4397d6..0069967a658 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -1933,5 +1933,50 @@ describe('Arrow Support', () => { /bigquery\.readsessions\.create denied/, ); }); + + it('should fail with a clear error instead of creating a read session for projects/undefined', async () => { + // Neither the job nor the client has a project ID. + bq.projectId = '{{projectId}}'; + const fakeJob = bq.job('dest-no-project', {location: 'US'}); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true}, + ]); + stubCreateQueryJob(bq, async () => [ + fakeJob, + { + jobReference: {jobId: 'dest-no-project', location: 'US'}, + configuration: { + query: { + destinationTable: { + projectId: 'other-project', + datasetId: 'dest_ds', + tableId: 'tbl', + }, + }, + }, + }, + ]); + let readSessions = 0; + stubStorageReadClient(bq, { + createReadSession: async () => { + readSessions++; + return [{}]; + }, + readRows: () => new EventEmitter(), + }); + + await assert.rejects( + bq + .createQueryArrowStream({ + query: 'SELECT 1', + destination: bq.dataset('dest_ds').table('tbl'), + }) + .toArray(), + /Cannot create an Arrow read session without a project ID\./, + ); + assert.strictEqual(readSessions, 0); + }); }); });