diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 12c5b938dc3..59ccf505fd3 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -16,11 +16,48 @@ 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'); +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; +} + +/** + * 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, + ) => BigQueryReadClient; +} + /** * [Beta] A stream of Apache Arrow `RecordBatch` objects returned by * {@link BigQuery#createQueryArrowStream}. @@ -37,6 +74,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,23 +152,56 @@ function requireWithMinVersion( } /** - * Loads `apache-arrow` and verifies that the optional - * `@google-cloud/bigquery-storage` (>= 6.3.0) package is installed. + * Loads `apache-arrow` and the optional `@google-cloud/bigquery-storage` + * (>= 6.3.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 // request fields, so check the version here. - requireWithMinVersion(requireFn, '@google-cloud/bigquery-storage', 6, 3); - return arrow; + const storage = requireWithMinVersion( + requireFn, + '@google-cloud/bigquery-storage', + 6, + 3, + ); + 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; } /** @@ -160,6 +264,145 @@ 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 + */ +async function streamDefaultJobStream( + ctx: ArrowStreamContext, + jobRef: bigquery.IJobReference | undefined, + offset: number, + initialSchema?: Uint8Array | string, +): Promise { + const {bq} = ctx; + const projectId = resolveProjectId(bq, jobRef?.projectId); + const jobId = jobRef?.jobId; + const location = jobRef?.location || bq.location; + + if (!projectId || !jobId || !location) { + throw new Error( + 'Cannot stream remaining Arrow rows without a jobReference.', + ); + } + + const readStream = `projects/${projectId}/locations/${location}/jobs/${jobId}/streams/_default`; + 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 + */ +async function streamReadRows( + ctx: ArrowStreamContext, + options: { + readStream: string; + offset: number; + initialSchema?: Uint8Array | string; + }, +): Promise { + const {bq, arrow, storage, stream} = ctx; + const {readStream, offset, initialSchema} = options; + + return new Promise((resolve, reject) => { + const grpcStream = getStorageReadClient(bq, storage).readRows({ + readStream, + offset, + }); + + 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. 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(); + } + if (opts.err) { + reject(opts.err); + } else { + resolve(); + } + }; + + 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}); + } + }; + 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); + }); +} + /** * Executes a query and returns a stream of `apache-arrow` `RecordBatch` objects. * @@ -169,67 +412,64 @@ export function createQueryArrowStream_( bq: BigQuery, query?: Query | string, ): 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; - let started = false; - stream._read = () => { - if (started) { - return; - } - started = true; - - let arrow: ArrowModule; - try { - arrow = 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 streamQueryResults = 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; } - let queryReq: bigquery.IQueryRequest | undefined; - try { - queryReq = bq.buildQueryRequest_(queryObj, {}); - } catch (err) { - // For example, invalid query parameters. - stream.destroy(err as Error); - return; - } + const ctx: ArrowStreamContext = {bq, arrow, storage, stream}; + readCtx = ctx; + + // 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. @@ -238,63 +478,72 @@ 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; + // The consumer may have destroyed the stream while jobs.query ran. + if (stream.destroyed) { + return; + } - 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. - // 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; - } + // 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; - 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); - } - if (!stream._ended && !stream.destroyed) { - stream.end(); + // 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; } - } catch (parseErr) { - stream.destroy(parseErr as Error); + stream.push(batch); } - }); + } + + if (!hasRemainingRows || stream.destroyed) { + return; + } + + await streamDefaultJobStream( + ctx, + res?.jobReference, + Number(inlineRowsCount), + res?.arrowSchema?.serializedSchema, + ); + }; + + // The only place that ends or destroys the stream. + const runQuery = async (): Promise => { + try { + await streamQueryResults(); + } catch (err) { + stream.destroy(err as Error); + return; + } + if (!stream.destroyed) { + stream.end(); + } }; return stream; diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index 4ce370e2720..5d6e201ec05 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -387,6 +387,14 @@ export class BigQuery extends Service { * Requires the optional `@google-cloud/bigquery-storage` (>= 6.3.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. + * 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. * @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 85d4e1609a0..ce5c4293969 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,53 @@ 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, + }); + } + }, + }); +} + +/** + * 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; @@ -107,9 +162,10 @@ describe('Arrow Support', () => { describe('checkArrowDependencies', () => { it('should succeed when @google-cloud/bigquery-storage >= 6.3.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', () => { @@ -357,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')}); @@ -406,30 +501,442 @@ 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']), + ); - it('should fail when jobs.query does not return all rows inline', async () => { stubJobsQuery(bq, { - res: {jobComplete: true, pageRowCount: '2', totalRows: '5'}, + 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, + }, + }, }); - // TODO: Replace with default job stream (pageRowCount < totalRows) tests once added. + const capturedReadRowsReqs: Array<{ + 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: { + 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', + ]); + // 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 () => { + 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 1), 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; + let readRowsStream: EventEmitter | undefined; + const sharedClient = stubStorageReadClient(bq, { + readRows: () => { + const grpcStream = new EventEmitter() as EventEmitter & { + cancel: () => void; + }; + readRowsStream = grpcStream; + 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, + ); + // 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 () => { + 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, + }); + let readRowsStream: EventEmitter | undefined; + stubStorageReadClient(bq, { + readRows: () => { + const grpcStream = new EventEmitter(); + readRowsStream = grpcStream; + 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]); + // 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 () => { + stubJobsQuery(bq, { + res: { + jobComplete: true, + jobReference: { + projectId: 'test-project', + jobId: 'job-bad-1', + location: 'US', + }, + totalRows: '2', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + }, + }); + 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: { + serializedRecordBatch: Buffer.from('not an arrow batch'), + }, + }), + ); + return grpcStream; + }, + }); + + // Verify the decode error is raised, the gRPC stream is cancelled, and + // its listeners are removed. 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); + assertReadRowsListenersRemoved(readRowsStream!); }); - 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); + }); }); });