From d5bdcaf3488ff3617cac3ab3122e415539369fc2 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 02:33:08 +0000 Subject: [PATCH 1/4] feat(bigquery): stream multi-page Arrow results from default job stream Implements the BigQuery Storage Read API fallback for inline jobs.query calls when rows remain after the first response page (pageRowCount < totalRows, or a pageToken is set): - Lazily initializes and reuses a single BigQueryReadClient on the BigQuery instance with explicit authClient, projectId, universeDomain, and 60s/60s gRPC keepalive settings. - Emits inline RecordBatch results first, then calls ReadRows on the anonymous default job stream (projects/{p}/locations/{l}/jobs/{j}/streams/_default) at offset = Number(pageRowCount). - Cancels only the active ReadRows gRPC stream if the caller terminates iteration early without closing the shared BigQueryReadClient. --- handwritten/bigquery/src/arrow.ts | 307 ++++++++++++-- handwritten/bigquery/src/bigquery.ts | 6 + handwritten/bigquery/system-test/arrow.ts | 219 ++++++++++ handwritten/bigquery/test/arrow.ts | 463 +++++++++++++++++++++- 4 files changed, 949 insertions(+), 46 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 223c6af54a9..7ba0dec4218 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -21,6 +21,31 @@ import type bigquery from './types'; type ArrowModule = typeof import('apache-arrow'); +interface IReadRowsResponse { + arrowSchema?: { + serializedSchema?: Uint8Array | string | null; + } | null; + arrowRecordBatch?: { + serializedRecordBatch?: Uint8Array | string | null; + } | null; +} + +interface ReadRowsStream extends NodeJS.EventEmitter { + cancel(): void; + pause(): void; + resume(): void; +} + +export interface BigQueryReadClient { + readRows(request: {readStream: string; offset?: number}): ReadRowsStream; +} + +export interface StorageModule { + BigQueryReadClient: new ( + options?: Record, + ) => BigQueryReadClient; +} + /** * [Beta] A stream of Apache Arrow `RecordBatch` objects returned by * {@link BigQuery#createQueryArrowStream}. @@ -37,6 +62,40 @@ export type ArrowRecordBatchStream = ResourceStream & { jobReference?: bigquery.IJobReference; }; +/** + * State shared by the helpers that stream rows for one + * {@link BigQuery#createQueryArrowStream} call. + */ +interface ArrowStreamContext { + bq: BigQuery; + arrow: ArrowModule; + storage: StorageModule; + stream: ArrowRecordBatchStream; + /** Called when the consumer wants more batches. */ + onRead?: () => void; +} + +const readClients = new WeakMap(); + +/** + * Returns the first project ID from `candidates`, falling back to the + * `BigQuery` client's project and ignoring the uninitialized `'{{projectId}}'` + * placeholder. + * + * @private + */ +function resolveProjectId( + bq: BigQuery, + ...candidates: Array +): string | undefined { + for (const candidate of [...candidates, bq.projectId]) { + if (candidate && candidate !== '{{projectId}}') { + return candidate; + } + } + return undefined; +} + /** * Requires a dependency and verifies its minimum major.minor version. * @@ -81,17 +140,18 @@ function requireWithMinVersion( } /** - * Loads `apache-arrow` and verifies that the optional - * `@google-cloud/bigquery-storage` (>= 6.1.0) package is installed. + * Loads `apache-arrow` and the optional `@google-cloud/bigquery-storage` + * (>= 6.1.0) package, verifying the latter's version. * * `apache-arrow` is a regular dependency, so its version is guaranteed by * `package.json`. It is loaded lazily so that only Arrow callers pay for it. * * @private */ -export function checkArrowDependencies( - requireFn: NodeRequire = require, -): ArrowModule { +export function checkArrowDependencies(requireFn: NodeRequire = require): { + arrow: ArrowModule; + storage: StorageModule; +} { const arrow: ArrowModule = requireFn('apache-arrow'); // Package managers only warn when an installed version is outside an // optional peer range, and an older version's protos silently drop newer @@ -99,8 +159,40 @@ export function checkArrowDependencies( // TODO: Bump minimum version (here and in package.json) once // @google-cloud/bigquery-storage is released with // @google-cloud/bigquery-storage-api >= 0.5.0. - requireWithMinVersion(requireFn, '@google-cloud/bigquery-storage', 6, 1); - return arrow; + const storage = requireWithMinVersion( + requireFn, + '@google-cloud/bigquery-storage', + 6, + 1, + ); + return {arrow, storage}; +} + +/** + * Returns the shared `BigQueryReadClient` for a `BigQuery` instance, creating + * it lazily on first use with the `BigQuery` client's credentials, universe + * domain, and gRPC keepalive settings. + * + * @private + */ +export function getStorageReadClient( + bq: BigQuery, + storage: StorageModule, +): BigQueryReadClient { + let client = readClients.get(bq); + if (!client) { + client = new storage.BigQueryReadClient({ + auth: bq.authClient, + projectId: resolveProjectId(bq), + universeDomain: bq.universeDomain, + // Ping while ReadRows is paused for a slow consumer, so a connection + // silently dropped by a proxy or NAT errors instead of hanging. + 'grpc.keepalive_time_ms': 60_000, + 'grpc.keepalive_timeout_ms': 60_000, + }); + readClients.set(bq, client); + } + return client; } /** @@ -163,6 +255,138 @@ export function decodeArrowPayload( : [createEmptyRecordBatch(arrow, reader.schema)]; } +/** + * Streams Arrow `RecordBatch`es from a job's sessionless default stream + * (`projects/{projectId}/locations/{location}/jobs/{jobId}/streams/_default`) + * starting at `offset`. + * + * @private + */ +function streamDefaultJobStream( + ctx: ArrowStreamContext, + jobRef: bigquery.IJobReference | undefined, + offset: number, + initialSchema?: Uint8Array | string, +): void { + const {bq, stream} = ctx; + const projectId = resolveProjectId(bq, jobRef?.projectId); + const jobId = jobRef?.jobId; + const location = jobRef?.location || bq.location; + + if (!projectId || !jobId || !location) { + stream.destroy( + new Error('Cannot stream remaining Arrow rows without a jobReference.'), + ); + return; + } + + const readStream = `projects/${projectId}/locations/${location}/jobs/${jobId}/streams/_default`; + streamReadRows(ctx, {readStream, offset, initialSchema}); +} + +/** + * Streams Arrow `RecordBatch`es from a Storage Read API stream starting at + * `offset`, pausing `ReadRows` while the consumer is behind. + * + * @private + */ +function streamReadRows( + ctx: ArrowStreamContext, + options: { + readStream: string; + offset: number; + initialSchema?: Uint8Array | string; + }, +): void { + const {bq, arrow, storage, stream} = ctx; + const {readStream, offset, initialSchema} = options; + + let readClient: BigQueryReadClient; + try { + readClient = getStorageReadClient(bq, storage); + } catch (err) { + stream.destroy(err as Error); + return; + } + + const grpcStream = readClient.readRows({ + readStream, + offset, + }); + + // Cancel only the active ReadRows gRPC stream if the caller ends iteration + // early, keeping the shared BigQueryReadClient open for subsequent queries. + let finished = false; + const cancelGrpcStream = () => { + if (!finished) { + finished = true; + grpcStream.cancel(); + } + }; + stream.once('close', cancelGrpcStream); + + // ReadRows is paused when push() reports a full buffer; resume it when the + // consumer asks for more so slow consumers don't buffer the whole result. + ctx.onRead = () => grpcStream.resume(); + + let serializedSchema: Uint8Array | undefined = initialSchema + ? toUint8Array(initialSchema) + : undefined; + + grpcStream.on('data', (resp: IReadRowsResponse) => { + if (stream._ended || stream.destroyed) { + cancelGrpcStream(); + return; + } + // The Storage Read API sends arrowSchema on the first ReadRowsResponse and + // may also emit responses without record batch bytes (e.g., heartbeats). + if (resp.arrowSchema?.serializedSchema) { + serializedSchema = toUint8Array(resp.arrowSchema.serializedSchema); + } + const rawBatch = resp.arrowRecordBatch?.serializedRecordBatch; + const hasBatchBytes = + rawBatch && + (typeof rawBatch === 'string' + ? rawBatch.length > 0 + : rawBatch.byteLength > 0); + if (!hasBatchBytes) { + return; + } + try { + const batches = decodeArrowPayload(arrow, serializedSchema, rawBatch); + for (const batch of batches) { + if (stream._ended || stream.destroyed) { + cancelGrpcStream(); + return; + } + if (!stream.push(batch)) { + grpcStream.pause(); + } + } + } catch (parseErr) { + cancelGrpcStream(); + stream.destroy(parseErr as Error); + } + }); + + grpcStream.on('error', (err: Error) => { + stream.removeListener('close', cancelGrpcStream); + if (finished || stream._ended || stream.destroyed) { + return; + } + finished = true; + stream.destroy(err); + }); + + grpcStream.on('end', () => { + stream.removeListener('close', cancelGrpcStream); + finished = true; + if (!stream._ended && !stream.destroyed) { + stream.end(); + } + }); +} + /** * Executes a query and returns a stream of `apache-arrow` `RecordBatch` objects. * @@ -178,15 +402,19 @@ export function createQueryArrowStream_( ) as ArrowRecordBatchStream; let started = false; + // Set once the query starts, so later reads can resume the active ReadRows. + let readCtx: ArrowStreamContext | undefined; stream._read = () => { if (started) { + readCtx?.onRead?.(); return; } started = true; let arrow: ArrowModule; + let storage: StorageModule; try { - arrow = checkArrowDependencies(); + ({arrow, storage} = checkArrowDependencies()); } catch (err) { stream.destroy(err as Error); return; @@ -219,6 +447,9 @@ export function createQueryArrowStream_( return; } + const ctx: ArrowStreamContext = {bq, arrow, storage, stream}; + readCtx = ctx; + let queryReq: bigquery.IQueryRequest | undefined; try { queryReq = bq.buildQueryRequest_(queryObj, {}); @@ -264,39 +495,51 @@ export function createQueryArrowStream_( } // pageRowCount and totalRows are returned as uint64 strings in REST JSON. + const inlineRowsCount = BigInt(res?.pageRowCount ?? 0); + const totalRowsCount = BigInt(res?.totalRows ?? 0); // Also check pageToken: on cached multi-page results, jobs.query can // return totalRows "0" and omit pageRowCount while setting pageToken. - if ( - Boolean(res?.pageToken) || - BigInt(res?.pageRowCount ?? 0) < BigInt(res?.totalRows ?? 0) - ) { - // TODO: Stream remaining rows via @google-cloud/bigquery-storage (streams/_default). - stream.destroy( - new Error( - 'Query results were not returned inline in the jobs.query response.', - ), - ); - return; - } + const hasRemainingRows = + Boolean(res?.pageToken) || inlineRowsCount < totalRowsCount; - try { - const batches = decodeArrowPayload( - arrow, - res?.arrowSchema?.serializedSchema, - res?.arrowRecordBatch?.serializedRecordBatch, - ); - for (const batch of batches) { - if (stream._ended || stream.destroyed) { - return; + // Emit the inline first page (or 0-row schema-preserving batch when totalRows is 0) + // before resuming remaining rows from streams/_default at offset = pageRowCount. + if (!hasRemainingRows || inlineRowsCount > 0n) { + try { + const batches = decodeArrowPayload( + arrow, + res?.arrowSchema?.serializedSchema, + res?.arrowRecordBatch?.serializedRecordBatch, + ); + for (const batch of batches) { + if (stream._ended || stream.destroyed) { + return; + } + stream.push(batch); } - stream.push(batch); + } catch (parseErr) { + stream.destroy(parseErr as Error); + return; } + } + + if (!hasRemainingRows) { if (!stream._ended && !stream.destroyed) { stream.end(); } - } catch (parseErr) { - stream.destroy(parseErr as Error); + return; } + + if (stream._ended || stream.destroyed) { + return; + } + + streamDefaultJobStream( + ctx, + res?.jobReference, + Number(inlineRowsCount), + res?.arrowSchema?.serializedSchema, + ); }); }; diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index f313c08846c..76cb1d0645e 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -387,6 +387,12 @@ export class BigQuery extends Service { * Requires the optional `@google-cloud/bigquery-storage` (>= 6.1.0) package * to be installed. * + * Rows that don't fit in the first `jobs.query` response are read with the + * BigQuery Storage Read API, which must be enabled in the project and is + * billed separately. The stream reads more rows only as the caller consumes + * batches, so iterate over batches to keep memory bounded; collecting every + * batch (for example with `toArray()`) holds the whole result in memory. + * * @param {string|object} [options] A string SQL query or configuration * object. See {@link BigQuery#query} for a complete list of options. * @returns {ArrowRecordBatchStream} A stream of `RecordBatch` objects. The diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts index 0fedc298f1f..2614a19cbb9 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -17,6 +17,7 @@ import {randomUUID} from 'crypto'; import {after, before, describe, it} from 'mocha'; import * as arrow from 'apache-arrow'; import {BigQuery, Job} from '../src'; +import {checkArrowDependencies} from '../src/arrow'; const bigquery = new BigQuery(); const GCLOUD_TESTS_PREFIX = 'nodejs_bq_test'; @@ -179,4 +180,222 @@ describe('Arrow Query Results', () => { }, ); }); + + it('should stream multi-page results across inline jobs.query and streams/_default and support early cancellation', async () => { + const totalRows = 20000; + const sql = `SELECT id, REPEAT('x', 1024) AS payload FROM UNNEST(GENERATE_ARRAY(1, ${totalRows})) AS id ORDER BY id`; + + // First, break out of iteration after the first batch to verify early cancellation keeps the shared client usable. + const earlyStream = bigquery.createQueryArrowStream({ + query: sql, + useQueryCache: false, + }); + let earlyBatches = 0; + for await (const batch of earlyStream) { + assert.ok(batch.numRows > 0); + earlyBatches++; + break; + } + assert.strictEqual(earlyBatches, 1); + + // Then read the full multi-page result on the same BigQuery client. + const stream = bigquery.createQueryArrowStream({ + query: sql, + useQueryCache: false, + }); + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify the result spanned the inline jobs.query batch and streams/_default. + assert.ok( + batches.length >= 2, + `Expected >= 2 batches, got ${batches.length}`, + ); + assert.ok(stream.jobReference?.jobId); + + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, totalRows); + const ids = Array.from(table.getChild('id')!.toArray()); + for (let i = 0; i < totalRows; i++) { + if (ids[i] !== BigInt(i + 1)) { + assert.fail( + `Expected id ${i + 1} at index ${i}, got ${String(ids[i])}`, + ); + } + } + }); + + it('should stream multi-page cached results from streams/_default', async () => { + // A cache hit can return totalRows "0" with a pageToken and no inline + // rows, so the second run reads every row from streams/_default. + const totalRows = 20000; + const sql = `SELECT id, REPEAT('c', 1024) AS payload FROM UNNEST(GENERATE_ARRAY(1, ${totalRows})) AS id ORDER BY id`; + for (let run = 0; run < 2; run++) { + const table = new arrow.Table( + await bigquery + .createQueryArrowStream({query: sql, useQueryCache: true}) + .toArray(), + ); + assert.strictEqual(table.numRows, totalRows); + assert.strictEqual(Number(table.getChild('id')!.get(0)), 1); + assert.strictEqual( + Number(table.getChild('id')!.get(totalRows - 1)), + totalRows, + ); + } + }); + + it('should cancel only the ReadRows stream of a read that ends early while another read on the same client continues', async () => { + // Concurrent reads share one BigQueryReadClient, so ending one read early + // must cancel only its own ReadRows stream. The read that ends early uses + // a 3M-row result so its stream is still open when it's cancelled. + const endedEarlySql = + 'SELECT a * 1000 + b AS id FROM UNNEST(GENERATE_ARRAY(0, 2999)) AS a, UNNEST(GENERATE_ARRAY(0, 999)) AS b'; + const totalRows = 20000; + const readFullySql = `SELECT id, REPEAT('x', 1024) AS payload FROM UNNEST(GENERATE_ARRAY(1, ${totalRows})) AS id ORDER BY id`; + const readClientProto = + checkArrowDependencies().storage.BigQueryReadClient.prototype; + const origReadRows = readClientProto.readRows; + let openedStreams = 0; + let cancelledStreams = 0; + readClientProto.readRows = function ( + this: unknown, + ...args: Parameters + ) { + const grpcStream = origReadRows.apply(this, args); + openedStreams++; + const origCancel = grpcStream.cancel?.bind(grpcStream); + grpcStream.cancel = () => { + cancelledStreams++; + origCancel?.(); + }; + return grpcStream; + }; + + try { + const endedEarly = bigquery + .createQueryArrowStream({query: endedEarlySql, useQueryCache: false}) + [Symbol.asyncIterator](); + const readFully = bigquery + .createQueryArrowStream({query: readFullySql, useQueryCache: false}) + [Symbol.asyncIterator](); + + // Advance each read until it has opened its own ReadRows stream. + while (openedStreams === 0) { + assert.ok(!(await endedEarly.next()).done); + } + const fullBatches: arrow.RecordBatch[] = []; + while (openedStreams === 1) { + const next = await readFully.next(); + assert.ok(!next.done); + fullBatches.push(next.value); + } + + // End the first read early, then finish the second. + await endedEarly.return!(); + let next = await readFully.next(); + while (!next.done) { + fullBatches.push(next.value); + next = await readFully.next(); + } + + // Verify only the ended read was cancelled and the other read every row. + assert.strictEqual(cancelledStreams, 1); + const ids = Array.from(new arrow.Table(fullBatches).getChild('id')!); + assert.strictEqual(ids.length, totalRows); + assert.strictEqual(ids[0], 1n); + assert.strictEqual(ids[totalRows - 1], BigInt(totalRows)); + } finally { + readClientProto.readRows = origReadRows; + } + }); + + it('should stream multi-page results from streams/_default in a non-US location (EU)', async () => { + const totalRows = 15000; + const stream = bigquery.createQueryArrowStream({ + query: `SELECT id, REPEAT('y', 1024) AS payload FROM UNNEST(GENERATE_ARRAY(1, ${totalRows})) AS id ORDER BY id`, + location: 'EU', + useQueryCache: false, + }); + + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify the job ran in EU and streamed all rows across inline + streams/_default. + assert.strictEqual(stream.jobReference?.location, 'EU'); + assert.ok( + batches.length >= 2, + `Expected >= 2 batches, got ${batches.length}`, + ); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, totalRows); + }); + + it('should decode all BigQuery data types with rows split between the inline first page and streams/_default', async () => { + const totalRows = 15000; + const sql = ` + SELECT + id, + IF(MOD(id, 2) = 1, TRUE, NULL) AS bool_col, + IF(MOD(id, 2) = 1, 42, NULL) AS int64_col, + IF(MOD(id, 2) = 1, 3.14, NULL) AS float64_col, + IF(MOD(id, 2) = 1, NUMERIC '123456789.123456789', NULL) AS numeric_col, + IF(MOD(id, 2) = 1, BIGNUMERIC '12345678901234567890.123456789012345678', NULL) AS bignumeric_col, + IF(MOD(id, 2) = 1, REPEAT('hello_', 200), NULL) AS string_col, + IF(MOD(id, 2) = 1, b'\\x01\\x02\\x03', NULL) AS bytes_col, + IF(MOD(id, 2) = 1, DATE '2026-01-15', NULL) AS date_col, + IF(MOD(id, 2) = 1, TIME '12:34:56.789012', NULL) AS time_col, + IF(MOD(id, 2) = 1, DATETIME '2026-01-15 12:34:56.789012', NULL) AS datetime_col, + IF(MOD(id, 2) = 1, TIMESTAMP '2026-01-15 12:34:56.789012Z', NULL) AS timestamp_col, + IF(MOD(id, 2) = 1, ST_GEOGPOINT(-122.4194, 37.7749), NULL) AS geography_col, + IF(MOD(id, 2) = 1, JSON '{"a": 1, "b": [true, "two"]}', NULL) AS json_col, + IF(MOD(id, 2) = 1, RANGE '[2026-01-01, 2026-02-01)', NULL) AS range_col, + [id, id + 1] AS array_col, + STRUCT(id AS x, IF(MOD(id, 2) = 1, 'nested', NULL) AS y) AS struct_col + FROM UNNEST(GENERATE_ARRAY(1, ${totalRows})) AS id + ORDER BY id + `; + + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({query: sql, useQueryCache: false}) + .toArray(); + + // Verify rows were split across the inline first page and streams/_default and combine cleanly into a Table. + assert.ok( + batches.length >= 2, + `Expected >= 2 batches across inline and stream, got ${batches.length}`, + ); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, totalRows); + assert.strictEqual(table.numCols, 17); + + // Check row 0 (id = 1, non-null values from inline page) and row 1 (id = 2, NULL values). + assert.strictEqual(table.getChild('bool_col')!.get(0), true); + assert.strictEqual(table.getChild('bool_col')!.get(1), null); + assert.strictEqual(table.getChild('int64_col')!.get(0), 42n); + assert.strictEqual(table.getChild('int64_col')!.get(1), null); + assert.strictEqual(table.getChild('float64_col')!.get(0), 3.14); + assert.strictEqual( + table.getChild('json_col')!.get(0), + '{"a":1,"b":[true,"two"]}', + ); + assert.deepStrictEqual( + Array.from(table.getChild('array_col')!.get(0).toArray()), + [1n, 2n], + ); + assert.deepStrictEqual(table.getChild('struct_col')!.get(0).toJSON(), { + x: 1n, + y: 'nested', + }); + + // Check the last row (id = 15000, streamed from streams/_default). + const lastIdx = totalRows - 1; + assert.strictEqual(table.getChild('id')!.get(lastIdx), BigInt(totalRows)); + assert.strictEqual(table.getChild('bool_col')!.get(lastIdx), null); + assert.deepStrictEqual( + table.getChild('struct_col')!.get(lastIdx).toJSON(), + { + x: BigInt(totalRows), + y: null, + }, + ); + }); }); diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 088a395b6d4..974956a8917 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -13,11 +13,17 @@ // limitations under the License. import * as assert from 'assert'; +import {EventEmitter} from 'events'; import {describe, it, before, after, beforeEach} from 'mocha'; import * as Module from 'module'; import * as arrow from 'apache-arrow'; import {BigQuery, Job, Query} from '../src'; -import {checkArrowDependencies} from '../src/arrow'; +import { + BigQueryReadClient, + checkArrowDependencies, + getStorageReadClient, + StorageModule, +} from '../src/arrow'; import bigquery from '../src/types'; /** @@ -25,7 +31,9 @@ import bigquery from '../src/types'; * on which copy of the package is installed (for example, an unbuilt in-repo * workspace link). */ -const storageStub = {}; +const storageStub = { + BigQueryReadClient: class {}, +}; /** * Serializes an Arrow Table into separate base64 `serializedSchema` and @@ -82,6 +90,40 @@ function stubJobsQuery( return () => capturedReq; } +/** + * Registers a stub `BigQueryReadClient` on `bq` via `getStorageReadClient`. + */ +type ReadRowsRequest = Parameters[0]; + +/** + * A fake ReadRows stream. google-gax streams always have `cancel`, `pause`, + * and `resume`; `stubStorageReadClient` adds no-op ones when a fake omits them. + */ +type FakeReadRowsStream = EventEmitter & + Partial void>>; + +/** + * Registers a fake `BigQueryReadClient` as the shared client for `bq`. + */ +function stubStorageReadClient( + bq: BigQuery, + client: {readRows: (req: ReadRowsRequest) => FakeReadRowsStream}, +): BigQueryReadClient { + const noop = () => {}; + return getStorageReadClient(bq, { + BigQueryReadClient: class { + readRows(req: ReadRowsRequest) { + const grpcStream = client.readRows(req); + return Object.assign(grpcStream, { + cancel: grpcStream.cancel ?? noop, + pause: grpcStream.pause ?? noop, + resume: grpcStream.resume ?? noop, + }); + } + }, + }); +} + describe('Arrow Support', () => { let originalRequire: typeof Module.prototype.require; @@ -107,9 +149,10 @@ describe('Arrow Support', () => { describe('checkArrowDependencies', () => { it('should succeed when @google-cloud/bigquery-storage >= 6.1.0 is installed', () => { - const arrowMod = checkArrowDependencies(); - // Verify the returned module exposes the Arrow RecordBatchReader API. + const {arrow: arrowMod, storage} = checkArrowDependencies(); + // Verify the returned modules expose RecordBatchReader and BigQueryReadClient. assert.strictEqual(typeof arrowMod.RecordBatchReader.from, 'function'); + assert.strictEqual(typeof storage.BigQueryReadClient, 'function'); }); it('should throw an actionable error when @google-cloud/bigquery-storage is missing', () => { @@ -406,30 +449,422 @@ describe('Arrow Support', () => { /Query did not complete within the jobs\.query timeout\./, ); }); + }); + + describe('BigQuery#createQueryArrowStream (default job stream)', () => { + let bq: BigQuery; + + beforeEach(() => { + bq = new BigQuery({ + projectId: 'test-project', + universeDomain: 'example.universe.goog', + }); + }); + + it('should lazily create and reuse one BigQueryReadClient with credentials, universeDomain, and keepalive', () => { + const capturedOpts: Array> = []; + const fakeStorage: StorageModule = { + BigQueryReadClient: class FakeReadClient { + constructor(opts?: Record) { + capturedOpts.push(opts || {}); + } + readRows(): never { + throw new Error('not called'); + } + }, + }; + + const client1 = getStorageReadClient(bq, fakeStorage); + const client2 = getStorageReadClient(bq, fakeStorage); + + // Verify the same BigQueryReadClient instance is reused across queries on the same BigQuery client. + assert.strictEqual(client1, client2); + assert.strictEqual(capturedOpts.length, 1); + assert.strictEqual(capturedOpts[0].auth, bq.authClient); + assert.strictEqual(capturedOpts[0].projectId, 'test-project'); + assert.strictEqual( + capturedOpts[0].universeDomain, + 'example.universe.goog', + ); + assert.strictEqual(capturedOpts[0]['grpc.keepalive_time_ms'], 60_000); + assert.strictEqual(capturedOpts[0]['grpc.keepalive_timeout_ms'], 60_000); + }); + + it('should emit inline batches first and stream remaining rows from streams/_default at offset = pageRowCount', async () => { + const page2Payload = serializeArrowTable( + makeSampleTable([3, 4], ['carol', 'dave']), + ); + + stubJobsQuery(bq, { + res: { + jobComplete: true, + queryId: 'multi-page-query', + jobReference: { + projectId: 'test-project', + jobId: 'job-multi-1', + location: 'EU', + }, + pageRowCount: '2', + totalRows: '4', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, + }); + + const capturedReadRowsReqs: Array<{ + readStream?: string; + offset?: number; + }> = []; + stubStorageReadClient(bq, { + readRows: (req: {readStream?: string; offset?: number}) => { + capturedReadRowsReqs.push(req); + const grpcStream = new EventEmitter(); + process.nextTick(() => { + grpcStream.emit('data', { + arrowSchema: { + serializedSchema: Buffer.from( + page2Payload.serializedSchema, + 'base64', + ), + }, + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + page2Payload.serializedRecordBatch, + 'base64', + ), + }, + }); + grpcStream.emit('end'); + }); + return grpcStream; + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl') + .toArray(); + + // Verify ReadRows was called on streams/_default using the job's EU location and offset = pageRowCount (2). + assert.deepStrictEqual(capturedReadRowsReqs, [ + { + readStream: + 'projects/test-project/locations/EU/jobs/job-multi-1/streams/_default', + offset: 2, + }, + ]); + + // Verify both the inline batch and streamed batch are emitted in order. + assert.strictEqual(batches.length, 2); + const table = new arrow.Table(batches); + assert.deepStrictEqual( + Array.from(table.getChild('id')!.toArray()), + [1, 2, 3, 4], + ); + assert.deepStrictEqual(table.getChild('name')!.toArray(), [ + 'alice', + 'bob', + 'carol', + 'dave', + ]); + }); + + it('should pause ReadRows while the consumer is behind and resume it when the consumer reads', async () => { + stubJobsQuery(bq, { + res: { + jobComplete: true, + jobReference: { + projectId: 'test-project', + jobId: 'job-backpressure-1', + location: 'US', + }, + pageRowCount: '2', + totalRows: '42', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, + }); + + const streamedBatches = 20; + const page2Payload = serializeArrowTable( + makeSampleTable([3, 4], ['carol', 'dave']), + ); + const flowEvents: string[] = []; + stubStorageReadClient(bq, { + readRows: () => { + let ended = false; + const grpcStream = Object.assign(new EventEmitter(), { + pause: () => flowEvents.push('pause'), + // Like a real paused gRPC stream, only finish once resumed. + resume: () => { + flowEvents.push('resume'); + if (flowEvents.includes('pause') && !ended) { + ended = true; + process.nextTick(() => grpcStream.emit('end')); + } + }, + }); + process.nextTick(() => { + // Emit more batches in one tick than the output stream buffers + // (highWaterMark 16), before the consumer can read any of them. + for (let i = 0; i < streamedBatches; i++) { + grpcStream.emit('data', { + arrowSchema: { + serializedSchema: Buffer.from( + page2Payload.serializedSchema, + 'base64', + ), + }, + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + page2Payload.serializedRecordBatch, + 'base64', + ), + }, + }); + } + }); + return grpcStream; + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl') + .toArray(); + + // Verify ReadRows was paused once the buffer filled, resumed when the + // consumer read, and no batches were lost. + const firstPause = flowEvents.indexOf('pause'); + assert.ok(firstPause >= 0, 'expected ReadRows to be paused'); + assert.ok( + flowEvents.indexOf('resume', firstPause) > firstPause, + 'expected ReadRows to be resumed after pausing', + ); + assert.strictEqual(batches.length, 1 + streamedBatches); + }); + + it('should cancel only the active ReadRows stream when the caller ends iteration early', async () => { + stubJobsQuery(bq, { + res: { + jobComplete: true, + jobReference: { + projectId: 'test-project', + jobId: 'job-cancel-1', + location: 'US', + }, + pageRowCount: '2', + totalRows: '6', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, + }); + + let cancelCount = 0; + const sharedClient = stubStorageReadClient(bq, { + readRows: () => { + const grpcStream = new EventEmitter() as EventEmitter & { + cancel: () => void; + }; + grpcStream.cancel = () => { + cancelCount++; + }; + process.nextTick(() => { + grpcStream.emit('data', { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + SAMPLE_PAYLOAD.serializedRecordBatch, + 'base64', + ), + }, + }); + }); + return grpcStream; + }, + }); + + const stream = bq.createQueryArrowStream('SELECT id, name FROM tbl'); + const received: arrow.RecordBatch[] = []; + for await (const batch of stream) { + received.push(batch); + if (received.length === 2) { + break; + } + } + + // ReadRows is cancelled from the stream's 'close' event, which Node emits + // on the next tick after the loop destroys the stream. + await new Promise(resolve => setImmediate(resolve)); + // Verify breaking out of the stream cancels the active gRPC stream while keeping the shared client. + assert.strictEqual(received.length, 2); + assert.strictEqual(cancelCount, 1); + assert.strictEqual( + getStorageReadClient(bq, { + BigQueryReadClient: class { + readRows(): never { + throw new Error('should not instantiate a new client'); + } + }, + }), + sharedClient, + ); + }); + + it('should raise a ReadRows error after the batches already emitted', async () => { + stubJobsQuery(bq, { + res: { + jobComplete: true, + jobReference: { + projectId: 'test-project', + jobId: 'job-err-1', + location: 'US', + }, + pageRowCount: '2', + totalRows: '4', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, + }); + const readRowsError = Object.assign(new Error('permission denied'), { + code: 7, + }); + stubStorageReadClient(bq, { + readRows: () => { + const grpcStream = new EventEmitter(); + process.nextTick(() => grpcStream.emit('error', readRowsError)); + return grpcStream; + }, + }); + + const stream = bq.createQueryArrowStream('SELECT id, name FROM tbl'); + const received: arrow.RecordBatch[] = []; + const errors: Error[] = []; + stream.on('data', batch => received.push(batch)); + stream.on('error', err => errors.push(err)); + stream.on('end', () => errors.push(new Error('unexpected end'))); + // The error and 'close' are emitted on later ticks. + await new Promise(resolve => stream.once('close', resolve)); + + // Verify the inline batch was delivered, then exactly one terminal event: + // the ReadRows error. + assert.strictEqual(received.length, 1); + assert.deepStrictEqual(errors, [readRowsError]); + }); - it('should fail when jobs.query does not return all rows inline', async () => { + it('should cancel ReadRows and raise the error when a streamed batch fails to decode', async () => { stubJobsQuery(bq, { - res: {jobComplete: true, pageRowCount: '2', totalRows: '5'}, + res: { + jobComplete: true, + jobReference: { + projectId: 'test-project', + jobId: 'job-bad-1', + location: 'US', + }, + totalRows: '2', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + }, + }); + let cancelCount = 0; + stubStorageReadClient(bq, { + readRows: () => { + const grpcStream = Object.assign(new EventEmitter(), { + cancel: () => cancelCount++, + }); + process.nextTick(() => + grpcStream.emit('data', { + arrowRecordBatch: { + serializedRecordBatch: Buffer.from('not an arrow batch'), + }, + }), + ); + return grpcStream; + }, }); - // TODO: Replace with default job stream (pageRowCount < totalRows) tests once added. + // Verify the decode error is raised and the gRPC stream is cancelled. await assert.rejects( - bq.createQueryArrowStream('SELECT 1').toArray(), - /Query results were not returned inline in the jobs\.query response\./, + bq.createQueryArrowStream('SELECT id, name FROM tbl').toArray(), ); + assert.strictEqual(cancelCount, 1); }); - it('should fail when jobs.query returns a pageToken with totalRows = 0', async () => { - // Cached multi-page results can report totalRows "0" and omit pageRowCount. + it('should fail when remaining rows must be streamed but jobReference is missing', async () => { stubJobsQuery(bq, { - res: {jobComplete: true, totalRows: '0', pageToken: 'token-1'}, + res: { + jobComplete: true, + pageRowCount: '2', + totalRows: '5', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, }); - // TODO: Replace with default job stream (pageToken) tests once added. + // Verify an actionable error is raised if BigQuery omits jobReference on a multi-page response. await assert.rejects( bq.createQueryArrowStream('SELECT 1').toArray(), - /Query results were not returned inline in the jobs\.query response\./, + /Cannot stream remaining Arrow rows without a jobReference\./, ); }); + + it('should stream from streams/_default at offset 0 when jobs.query returns a pageToken with totalRows = 0', async () => { + // Cached multi-page ARROW results can report totalRows "0", omit + // pageRowCount, and return no inline rows while setting pageToken. + stubJobsQuery(bq, { + res: { + jobComplete: true, + jobReference: { + projectId: 'test-project', + jobId: 'job-cached-1', + location: 'US', + }, + totalRows: '0', + pageToken: 'token-1', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + }, + }); + + const capturedOffsets: Array = []; + stubStorageReadClient(bq, { + readRows: (req: {offset?: number}) => { + capturedOffsets.push(req.offset); + const grpcStream = new EventEmitter(); + process.nextTick(() => { + grpcStream.emit('data', { + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + SAMPLE_PAYLOAD.serializedRecordBatch, + 'base64', + ), + }, + }); + grpcStream.emit('end'); + }); + return grpcStream; + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl') + .toArray(); + + // Verify the rows are read from the start of the default stream instead + // of returning an empty result. + assert.deepStrictEqual(capturedOffsets, [0]); + assert.strictEqual(new arrow.Table(batches).numRows, 2); + }); }); }); From c4b2070d5712ab7f677a3c6cb3c9a5229913b8ba Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 18:07:32 +0000 Subject: [PATCH 2/4] refactor(bigquery): centralize Arrow stream end and ReadRows cleanup Run the query in one async function whose outcome is the only place that ends or destroys the stream, so helpers no longer check the private ResourceStream _ended flag. Wrap ReadRows in a promise that settles through one settle() on end, error, decode failure, or early destroy. settle() removes every listener, keeps a no-op 'error' listener so the CANCELLED error that follows cancel() can't crash the process, and cancels the call unless it already finished. --- handwritten/bigquery/src/arrow.ts | 348 ++++++++++++++--------------- handwritten/bigquery/test/arrow.ts | 32 ++- 2 files changed, 194 insertions(+), 186 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 7ba0dec4218..9466a6f08cd 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -16,7 +16,7 @@ import {ResourceStream} from '@google-cloud/paginator'; import type {RecordBatch, Schema} from 'apache-arrow'; -import type {BigQuery, Query} from './bigquery'; +import type {BigQuery, JobsQueryResponse, Query} from './bigquery'; import type bigquery from './types'; type ArrowModule = typeof import('apache-arrow'); @@ -262,32 +262,35 @@ export function decodeArrowPayload( * * @private */ -function streamDefaultJobStream( +async function streamDefaultJobStream( ctx: ArrowStreamContext, jobRef: bigquery.IJobReference | undefined, offset: number, initialSchema?: Uint8Array | string, -): void { - const {bq, stream} = ctx; +): Promise { + const {bq} = ctx; const projectId = resolveProjectId(bq, jobRef?.projectId); const jobId = jobRef?.jobId; const location = jobRef?.location || bq.location; if (!projectId || !jobId || !location) { - stream.destroy( - new Error('Cannot stream remaining Arrow rows without a jobReference.'), + throw new Error( + 'Cannot stream remaining Arrow rows without a jobReference.', ); - return; } const readStream = `projects/${projectId}/locations/${location}/jobs/${jobId}/streams/_default`; - streamReadRows(ctx, {readStream, offset, initialSchema}); + await streamReadRows(ctx, {readStream, offset, initialSchema}); } /** * Streams Arrow `RecordBatch`es from a Storage Read API stream starting at * `offset`, pausing `ReadRows` while the consumer is behind. * + * Resolves when `ReadRows` ends or the consumer destroys the stream, and + * rejects on a `ReadRows` or decode error. In every case, the listeners added + * here are removed, and `ReadRows` is cancelled unless it already finished. + * * @private */ function streamReadRows( @@ -297,93 +300,94 @@ function streamReadRows( offset: number; initialSchema?: Uint8Array | string; }, -): void { +): Promise { const {bq, arrow, storage, stream} = ctx; const {readStream, offset, initialSchema} = options; - let readClient: BigQueryReadClient; - try { - readClient = getStorageReadClient(bq, storage); - } catch (err) { - stream.destroy(err as Error); - return; - } - - const grpcStream = readClient.readRows({ - readStream, - offset, - }); - - // Cancel only the active ReadRows gRPC stream if the caller ends iteration - // early, keeping the shared BigQueryReadClient open for subsequent queries. - let finished = false; - const cancelGrpcStream = () => { - if (!finished) { - finished = true; - grpcStream.cancel(); - } - }; - stream.once('close', cancelGrpcStream); + return new Promise((resolve, reject) => { + const grpcStream = getStorageReadClient(bq, storage).readRows({ + readStream, + offset, + }); - // ReadRows is paused when push() reports a full buffer; resume it when the - // consumer asks for more so slow consumers don't buffer the whole result. - ctx.onRead = () => grpcStream.resume(); + let serializedSchema: Uint8Array | undefined = initialSchema + ? toUint8Array(initialSchema) + : undefined; - let serializedSchema: Uint8Array | undefined = initialSchema - ? toUint8Array(initialSchema) - : undefined; + let settled = false; + // Detaches every listener added here and settles the promise. Cancels + // only this ReadRows call, keeping the shared BigQueryReadClient open. + const settle = (opts: {cancel: boolean; err?: Error}) => { + if (settled) { + return; + } + settled = true; + grpcStream.removeListener('data', onData); + grpcStream.removeListener('error', onError); + grpcStream.removeListener('end', onEnd); + stream.removeListener('close', onClose); + ctx.onRead = undefined; + // cancel() makes gRPC emit a CANCELLED 'error' on a later tick, which + // would crash the process if no 'error' listener were attached. + grpcStream.on('error', () => {}); + if (opts.cancel) { + grpcStream.cancel(); + } + if (opts.err) { + reject(opts.err); + } else { + resolve(); + } + }; - grpcStream.on('data', (resp: IReadRowsResponse) => { - if (stream._ended || stream.destroyed) { - cancelGrpcStream(); - return; - } - // The Storage Read API sends arrowSchema on the first ReadRowsResponse and - // may also emit responses without record batch bytes (e.g., heartbeats). - if (resp.arrowSchema?.serializedSchema) { - serializedSchema = toUint8Array(resp.arrowSchema.serializedSchema); - } - const rawBatch = resp.arrowRecordBatch?.serializedRecordBatch; - const hasBatchBytes = - rawBatch && - (typeof rawBatch === 'string' - ? rawBatch.length > 0 - : rawBatch.byteLength > 0); - if (!hasBatchBytes) { - return; - } - try { - const batches = decodeArrowPayload(arrow, serializedSchema, rawBatch); - for (const batch of batches) { - if (stream._ended || stream.destroyed) { - cancelGrpcStream(); - return; - } - if (!stream.push(batch)) { - grpcStream.pause(); + const onData = (resp: IReadRowsResponse) => { + if (stream.destroyed) { + settle({cancel: true}); + return; + } + // The Storage Read API sends arrowSchema on the first ReadRowsResponse + // and may also emit responses without record batch bytes (e.g., + // heartbeats). + if (resp.arrowSchema?.serializedSchema) { + serializedSchema = toUint8Array(resp.arrowSchema.serializedSchema); + } + const rawBatch = resp.arrowRecordBatch?.serializedRecordBatch; + const hasBatchBytes = + rawBatch && + (typeof rawBatch === 'string' + ? rawBatch.length > 0 + : rawBatch.byteLength > 0); + if (!hasBatchBytes) { + return; + } + try { + const batches = decodeArrowPayload(arrow, serializedSchema, rawBatch); + for (const batch of batches) { + if (stream.destroyed) { + settle({cancel: true}); + return; + } + if (!stream.push(batch)) { + grpcStream.pause(); + } } + } catch (parseErr) { + settle({cancel: true, err: parseErr as Error}); } - } catch (parseErr) { - cancelGrpcStream(); - stream.destroy(parseErr as Error); - } - }); - - grpcStream.on('error', (err: Error) => { - stream.removeListener('close', cancelGrpcStream); - if (finished || stream._ended || stream.destroyed) { - return; - } - finished = true; - stream.destroy(err); - }); - - grpcStream.on('end', () => { - stream.removeListener('close', cancelGrpcStream); - finished = true; - if (!stream._ended && !stream.destroyed) { - stream.end(); - } + }; + const onError = (err: Error) => settle({cancel: false, err}); + const onEnd = () => settle({cancel: false}); + // The consumer destroyed the stream (e.g., ended iteration early). + const onClose = () => settle({cancel: true}); + + // ReadRows is paused when push() reports a full buffer; resume it when the + // consumer asks for more so slow consumers don't buffer the whole result. + ctx.onRead = () => grpcStream.resume(); + + grpcStream.on('data', onData); + grpcStream.on('error', onError); + grpcStream.on('end', onEnd); + stream.on('close', onClose); }); } @@ -404,66 +408,40 @@ export function createQueryArrowStream_( let started = false; // Set once the query starts, so later reads can resume the active ReadRows. let readCtx: ArrowStreamContext | undefined; - stream._read = () => { - if (started) { - readCtx?.onRead?.(); - return; - } - started = true; - let arrow: ArrowModule; - let storage: StorageModule; - try { - ({arrow, storage} = checkArrowDependencies()); - } catch (err) { - stream.destroy(err as Error); - return; - } + // Runs the query and pushes its batches. The returned promise settles once + // all batches are pushed, the consumer destroys the stream, or a step fails. + const runQuery = async (): Promise => { + const {arrow, storage} = checkArrowDependencies(); if (!query || (typeof query === 'object' && !query.query)) { - stream.destroy(new Error('A SQL query string is required.')); - return; + throw new Error('A SQL query string is required.'); } const queryObj: Query = typeof query === 'string' ? {query} : {...query}; if (queryObj.writeDisposition === 'WRITE_APPEND') { - stream.destroy( - new Error( - 'WRITE_APPEND is not supported with Arrow query results because reading ' + - 'the destination table via the Storage Read API would include pre-existing rows.', - ), + throw new Error( + 'WRITE_APPEND is not supported with Arrow query results because reading ' + + 'the destination table via the Storage Read API would include pre-existing rows.', ); - return; } if (queryObj.destination) { // TODO: Support explicit destination tables via CreateReadSession. - stream.destroy( - new Error( - 'Explicit destination tables are not yet supported in createQueryArrowStream.', - ), + throw new Error( + 'Explicit destination tables are not yet supported in createQueryArrowStream.', ); - return; } const ctx: ArrowStreamContext = {bq, arrow, storage, stream}; readCtx = ctx; - let queryReq: bigquery.IQueryRequest | undefined; - try { - queryReq = bq.buildQueryRequest_(queryObj, {}); - } catch (err) { - // For example, invalid query parameters. - stream.destroy(err as Error); - return; - } + // Throws for an invalid request, for example invalid query parameters. + const queryReq = bq.buildQueryRequest_(queryObj, {}); if (!queryReq) { // TODO: Route queries that cannot use jobs.query through jobs.insert (createQueryJob). - stream.destroy( - new Error('Query configuration cannot be executed via jobs.query.'), - ); - return; + throw new Error('Query configuration cannot be executed via jobs.query.'); } // Default to JOB_CREATION_OPTIONAL unless overridden on the query or client. @@ -472,75 +450,75 @@ export function createQueryArrowStream_( } queryReq.queryResultsFormat = 'ARROW'; - void bq.runJobsQuery(queryReq, (err, job, res) => { - if (err) { - stream.destroy(err); - return; - } + const [job, res] = (await bq.runJobsQuery(queryReq)) as JobsQueryResponse; - if (res?.queryId) { - stream.queryId = res.queryId; - } - if (res?.jobReference && job) { - stream.jobReference = res.jobReference; - stream.emit('job', job); - } + if (res?.queryId) { + stream.queryId = res.queryId; + } + if (res?.jobReference && job) { + stream.jobReference = res.jobReference; + stream.emit('job', job); + } - if (res?.jobComplete === false) { - // TODO: Wait for unfinished jobs before streaming Arrow results. - stream.destroy( - new Error('Query did not complete within the jobs.query timeout.'), - ); - return; - } + if (res?.jobComplete === false) { + // TODO: Wait for unfinished jobs before streaming Arrow results. + throw new Error('Query did not complete within the jobs.query timeout.'); + } - // pageRowCount and totalRows are returned as uint64 strings in REST JSON. - const inlineRowsCount = BigInt(res?.pageRowCount ?? 0); - const totalRowsCount = BigInt(res?.totalRows ?? 0); - // Also check pageToken: on cached multi-page results, jobs.query can - // return totalRows "0" and omit pageRowCount while setting pageToken. - const hasRemainingRows = - Boolean(res?.pageToken) || inlineRowsCount < totalRowsCount; - - // Emit the inline first page (or 0-row schema-preserving batch when totalRows is 0) - // before resuming remaining rows from streams/_default at offset = pageRowCount. - if (!hasRemainingRows || inlineRowsCount > 0n) { - try { - const batches = decodeArrowPayload( - arrow, - res?.arrowSchema?.serializedSchema, - res?.arrowRecordBatch?.serializedRecordBatch, - ); - for (const batch of batches) { - if (stream._ended || stream.destroyed) { - return; - } - stream.push(batch); - } - } catch (parseErr) { - stream.destroy(parseErr as Error); + // pageRowCount and totalRows are returned as uint64 strings in REST JSON. + const inlineRowsCount = BigInt(res?.pageRowCount ?? 0); + const totalRowsCount = BigInt(res?.totalRows ?? 0); + // Also check pageToken: on cached multi-page results, jobs.query can + // return totalRows "0" and omit pageRowCount while setting pageToken. + const hasRemainingRows = + Boolean(res?.pageToken) || inlineRowsCount < totalRowsCount; + + // Emit the inline first page (or 0-row schema-preserving batch when totalRows is 0) + // before resuming remaining rows from streams/_default at offset = pageRowCount. + if (!hasRemainingRows || inlineRowsCount > 0n) { + const batches = decodeArrowPayload( + arrow, + res?.arrowSchema?.serializedSchema, + res?.arrowRecordBatch?.serializedRecordBatch, + ); + for (const batch of batches) { + if (stream.destroyed) { return; } + stream.push(batch); } + } - if (!hasRemainingRows) { - if (!stream._ended && !stream.destroyed) { - stream.end(); - } - return; - } + if (!hasRemainingRows || stream.destroyed) { + return; + } - if (stream._ended || stream.destroyed) { + await streamDefaultJobStream( + ctx, + res?.jobReference, + Number(inlineRowsCount), + res?.arrowSchema?.serializedSchema, + ); + }; + + stream._read = () => { + if (started) { + readCtx?.onRead?.(); + return; + } + started = true; + // The only place that ends or destroys the stream. + void (async () => { + try { + await runQuery(); + } catch (err) { + stream.destroy(err as Error); return; } - - streamDefaultJobStream( - ctx, - res?.jobReference, - Number(inlineRowsCount), - res?.arrowSchema?.serializedSchema, - ); - }); + if (!stream.destroyed) { + stream.end(); + } + })(); }; return stream; diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 974956a8917..49491697b52 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -124,6 +124,19 @@ function stubStorageReadClient( }); } +/** + * Asserts that only a no-op 'error' listener remains on a finished ReadRows + * stream, so an error gRPC emits afterwards (e.g., CANCELLED after cancel()) + * is swallowed instead of crashing the process. + */ +function assertReadRowsListenersRemoved(grpcStream: EventEmitter): void { + assert.strictEqual(grpcStream.listenerCount('data'), 0); + assert.strictEqual(grpcStream.listenerCount('end'), 0); + assert.strictEqual(grpcStream.listenerCount('error'), 1); + const cancelled = Object.assign(new Error('1 CANCELLED'), {code: 1}); + assert.doesNotThrow(() => grpcStream.emit('error', cancelled)); +} + describe('Arrow Support', () => { let originalRequire: typeof Module.prototype.require; @@ -517,10 +530,12 @@ describe('Arrow Support', () => { readStream?: string; offset?: number; }> = []; + let readRowsStream: EventEmitter | undefined; stubStorageReadClient(bq, { readRows: (req: {readStream?: string; offset?: number}) => { capturedReadRowsReqs.push(req); const grpcStream = new EventEmitter(); + readRowsStream = grpcStream; process.nextTick(() => { grpcStream.emit('data', { arrowSchema: { @@ -568,6 +583,8 @@ describe('Arrow Support', () => { 'carol', 'dave', ]); + // Verify the ReadRows listeners were removed once it ended. + assertReadRowsListenersRemoved(readRowsStream!); }); it('should pause ReadRows while the consumer is behind and resume it when the consumer reads', async () => { @@ -665,11 +682,13 @@ describe('Arrow Support', () => { }); let cancelCount = 0; + let readRowsStream: EventEmitter | undefined; const sharedClient = stubStorageReadClient(bq, { readRows: () => { const grpcStream = new EventEmitter() as EventEmitter & { cancel: () => void; }; + readRowsStream = grpcStream; grpcStream.cancel = () => { cancelCount++; }; @@ -718,6 +737,9 @@ describe('Arrow Support', () => { }), sharedClient, ); + // Verify the ReadRows listeners were removed and the CANCELLED error + // that follows cancel() doesn't crash the process. + assertReadRowsListenersRemoved(readRowsStream!); }); it('should raise a ReadRows error after the batches already emitted', async () => { @@ -740,9 +762,11 @@ describe('Arrow Support', () => { const readRowsError = Object.assign(new Error('permission denied'), { code: 7, }); + let readRowsStream: EventEmitter | undefined; stubStorageReadClient(bq, { readRows: () => { const grpcStream = new EventEmitter(); + readRowsStream = grpcStream; process.nextTick(() => grpcStream.emit('error', readRowsError)); return grpcStream; }, @@ -761,6 +785,8 @@ describe('Arrow Support', () => { // the ReadRows error. assert.strictEqual(received.length, 1); assert.deepStrictEqual(errors, [readRowsError]); + // Verify the ReadRows listeners were removed after the error. + assertReadRowsListenersRemoved(readRowsStream!); }); it('should cancel ReadRows and raise the error when a streamed batch fails to decode', async () => { @@ -777,11 +803,13 @@ describe('Arrow Support', () => { }, }); let cancelCount = 0; + let readRowsStream: EventEmitter | undefined; stubStorageReadClient(bq, { readRows: () => { const grpcStream = Object.assign(new EventEmitter(), { cancel: () => cancelCount++, }); + readRowsStream = grpcStream; process.nextTick(() => grpcStream.emit('data', { arrowRecordBatch: { @@ -793,11 +821,13 @@ describe('Arrow Support', () => { }, }); - // Verify the decode error is raised and the gRPC stream is cancelled. + // Verify the decode error is raised, the gRPC stream is cancelled, and + // its listeners are removed. await assert.rejects( bq.createQueryArrowStream('SELECT id, name FROM tbl').toArray(), ); assert.strictEqual(cancelCount, 1); + assertReadRowsListenersRemoved(readRowsStream!); }); it('should fail when remaining rows must be streamed but jobReference is missing', async () => { From d1472dbdde6f9dc17312743f0f1c528ca430d28c Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 20:42:06 +0000 Subject: [PATCH 3/4] docs(bigquery): mark test-only Arrow types private BigQueryReadClient and StorageModule are exported only so the unit tests can stub them. Document them as @private, like the exported helper functions in arrow.ts, so they stay out of the public API docs. --- handwritten/bigquery/src/arrow.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 9466a6f08cd..dd99d285354 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -36,10 +36,22 @@ interface ReadRowsStream extends NodeJS.EventEmitter { resume(): void; } +/** + * The part of `@google-cloud/bigquery-storage`'s `BigQueryReadClient` that + * {@link BigQuery#createQueryArrowStream} uses. Exported for tests. + * + * @private + */ export interface BigQueryReadClient { readRows(request: {readStream: string; offset?: number}): ReadRowsStream; } +/** + * The part of the `@google-cloud/bigquery-storage` module that + * {@link BigQuery#createQueryArrowStream} uses. Exported for tests. + * + * @private + */ export interface StorageModule { BigQueryReadClient: new ( options?: Record, From 8c53d44578aa6bd308b60bf14b9fce5d427ea1d0 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 21 Oct 2026 00:25:45 +0000 Subject: [PATCH 4/4] fix(bigquery): stop Arrow stream work after destroy and limit buffering Address review feedback on the default job stream: - Return after jobs.query if the consumer destroyed the stream, so no 'job' event or batches are emitted on a closed stream. - Buffer at most one RecordBatch (highWaterMark 1) instead of 16. - Pass read() through streamOptions instead of assigning stream._read, and move the try/catch that ends or destroys the stream into runQuery. - Mark streamReadRows async. - Document calling destroy() when stopping early, and that the no-op ReadRows 'error' listener swallows every later error. --- handwritten/bigquery/src/arrow.ts | 64 +++++++++++++++++----------- handwritten/bigquery/src/bigquery.ts | 2 + handwritten/bigquery/test/arrow.ts | 44 ++++++++++++++++++- 3 files changed, 85 insertions(+), 25 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index f36a9e4a42c..59ccf505fd3 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -302,7 +302,7 @@ async function streamDefaultJobStream( * * @private */ -function streamReadRows( +async function streamReadRows( ctx: ArrowStreamContext, options: { readStream: string; @@ -337,7 +337,10 @@ function streamReadRows( stream.removeListener('close', onClose); ctx.onRead = undefined; // cancel() makes gRPC emit a CANCELLED 'error' on a later tick, which - // would crash the process if no 'error' listener were attached. + // would crash the process if no 'error' listener were attached. This + // listener swallows every later error on this ReadRows call, not just + // CANCELLED: we're done with the call, so there's no value in + // surfacing them. grpcStream.on('error', () => {}); if (opts.cancel) { grpcStream.cancel(); @@ -409,18 +412,34 @@ export function createQueryArrowStream_( bq: BigQuery, query?: Query | string, ): ArrowRecordBatchStream { - const stream = new ResourceStream( - {}, - () => {}, - ) as ArrowRecordBatchStream; - let started = false; // Set once the query starts, so later reads can resume the active ReadRows. let readCtx: ArrowStreamContext | undefined; + const stream = new ResourceStream( + { + streamOptions: { + // Buffer at most one decoded RecordBatch ahead of the consumer, since + // each batch can hold many rows. + highWaterMark: 1, + // Called when the consumer wants more batches: the first call starts + // the query, and later calls resume a paused ReadRows call. + read: () => { + if (started) { + readCtx?.onRead?.(); + return; + } + started = true; + void runQuery(); + }, + }, + }, + () => {}, + ) as ArrowRecordBatchStream; + // Runs the query and pushes its batches. The returned promise settles once // all batches are pushed, the consumer destroys the stream, or a step fails. - const runQuery = async (): Promise => { + const streamQueryResults = async (): Promise => { const {arrow, storage} = checkArrowDependencies(); if (!query || (typeof query === 'object' && !query.query)) { @@ -460,6 +479,10 @@ export function createQueryArrowStream_( queryReq.queryResultsFormat = 'ARROW'; const [job, res] = (await bq.runJobsQuery(queryReq)) as JobsQueryResponse; + // The consumer may have destroyed the stream while jobs.query ran. + if (stream.destroyed) { + return; + } if (res?.queryId) { stream.queryId = res.queryId; @@ -510,24 +533,17 @@ export function createQueryArrowStream_( ); }; - stream._read = () => { - if (started) { - readCtx?.onRead?.(); + // The only place that ends or destroys the stream. + const runQuery = async (): Promise => { + try { + await streamQueryResults(); + } catch (err) { + stream.destroy(err as Error); return; } - started = true; - // The only place that ends or destroys the stream. - void (async () => { - try { - await runQuery(); - } catch (err) { - stream.destroy(err as Error); - return; - } - if (!stream.destroyed) { - stream.end(); - } - })(); + if (!stream.destroyed) { + stream.end(); + } }; return stream; diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index 4c405d8dc6c..5d6e201ec05 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -392,6 +392,8 @@ export class BigQuery extends Service { * billed separately. The stream reads more rows only as the caller consumes * batches, so iterate over batches to keep memory bounded; collecting every * batch (for example with `toArray()`) holds the whole result in memory. + * If you stop reading before the stream ends, call `stream.destroy()` to + * cancel the read; breaking out of `for await` does this for you. * * @param {string|object} [options] A string SQL query or configuration * object. See {@link BigQuery#query} for a complete list of options. diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 047e279de55..ce5c4293969 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -413,6 +413,45 @@ describe('Arrow Support', () => { ); }); + it('should not emit job or push batches when the stream is destroyed during jobs.query', async () => { + const fakeJob = bq.job('job-destroyed-1', {location: 'US'}); + bq.job = () => fakeJob; + const stream = bq.createQueryArrowStream('SELECT 1'); + bq.request = (reqOpts, callback) => { + // The consumer gives up while jobs.query is in flight. + stream.destroy(); + callback(null, { + jobComplete: true, + jobReference: { + projectId: 'test-project', + jobId: 'job-destroyed-1', + location: 'US', + }, + totalRows: '2', + pageRowCount: '2', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }); + }; + let jobEmitted = false; + stream.on('job', () => { + jobEmitted = true; + }); + const batches: arrow.RecordBatch[] = []; + stream.on('data', batch => batches.push(batch)); + stream.resume(); + await new Promise(resolve => stream.once('close', resolve)); + // Let the query finish its work after jobs.query returns. + await new Promise(resolve => setImmediate(resolve)); + + // Verify nothing is emitted on the destroyed stream. + assert.strictEqual(jobEmitted, false); + assert.strictEqual(stream.jobReference, undefined); + assert.strictEqual(batches.length, 0); + }); + it('should forward jobs.query errors to the stream error event', async () => { stubJobsQuery(bq, {err: new Error('syntax error in SQL')}); @@ -626,7 +665,7 @@ describe('Arrow Support', () => { }); process.nextTick(() => { // Emit more batches in one tick than the output stream buffers - // (highWaterMark 16), before the consumer can read any of them. + // (highWaterMark 1), before the consumer can read any of them. for (let i = 0; i < streamedBatches; i++) { grpcStream.emit('data', { arrowSchema: { @@ -787,6 +826,9 @@ describe('Arrow Support', () => { assert.deepStrictEqual(errors, [readRowsError]); // Verify the ReadRows listeners were removed after the error. assertReadRowsListenersRemoved(readRowsStream!); + // Verify the 'close' listener added to the Arrow stream was removed too + // (the test's own once('close') listener has already fired). + assert.strictEqual(stream.listenerCount('close'), 0); }); it('should cancel ReadRows and raise the error when a streamed batch fails to decode', async () => {