diff --git a/handwritten/bigquery/package.json b/handwritten/bigquery/package.json index a810796404c9..717d3fe6b032 100644 --- a/handwritten/bigquery/package.json +++ b/handwritten/bigquery/package.json @@ -54,6 +54,7 @@ "@google-cloud/paginator": "^7.0.0", "@google-cloud/precise-date": "^6.0.0", "@google-cloud/promisify": "^6.0.0", + "apache-arrow": "^21.0.0", "teeny-request": "^11.0.0", "arrify": "^3.0.0", "big.js": "^7.0.0", @@ -61,6 +62,14 @@ "extend": "^3.0.2", "stream-events": "^1.0.5" }, + "peerDependencies": { + "@google-cloud/bigquery-storage": "^6.1.0" + }, + "peerDependenciesMeta": { + "@google-cloud/bigquery-storage": { + "optional": true + } + }, "overrides": { "@sinonjs/fake-timers": "15.1.1", "@google-cloud/common": { @@ -68,6 +77,7 @@ } }, "devDependencies": { + "@google-cloud/bigquery-storage": "^6.1.0", "@google-cloud/storage": "^7.16.0", "@types/big.js": "^6.2.2", "@types/duplexify": "^3.6.4", diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts new file mode 100644 index 000000000000..223c6af54a9c --- /dev/null +++ b/handwritten/bigquery/src/arrow.ts @@ -0,0 +1,304 @@ +/*! + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import {ResourceStream} from '@google-cloud/paginator'; +import type {RecordBatch, Schema} from 'apache-arrow'; +import type {BigQuery, Query} from './bigquery'; +import type bigquery from './types'; + +type ArrowModule = typeof import('apache-arrow'); + +/** + * [Beta] A stream of Apache Arrow `RecordBatch` objects returned by + * {@link BigQuery#createQueryArrowStream}. + * + * `queryId` and `jobReference` are set once BigQuery responds to the query, + * before the first batch is emitted. + * + * @beta + */ +export type ArrowRecordBatchStream = ResourceStream & { + /** The query ID, when the query runs through `jobs.query`. */ + queryId?: string; + /** The job reference, when BigQuery creates a job for the query. */ + jobReference?: bigquery.IJobReference; +}; + +/** + * Requires a dependency and verifies its minimum major.minor version. + * + * @private + */ +function requireWithMinVersion( + requireFn: NodeRequire, + pkgName: string, + minMajor: number, + minMinor = 0, +): T { + const minVersion = `${minMajor}.${minMinor}.0`; + let mod: T; + let pkg: {version?: string}; + try { + mod = requireFn(pkgName); + pkg = requireFn(`${pkgName}/package.json`); + } catch (err) { + throw new Error( + `Arrow support requires the "${pkgName}" package (>= ${minVersion}), ` + + `but it could not be loaded: ${(err as Error).message}. ` + + `Please install it using: npm install ${pkgName}`, + {cause: err}, + ); + } + + const [majorStr, minorStr] = (pkg.version || '0.0.0').split('.'); + const major = Number.parseInt(majorStr, 10); + const minor = Number.parseInt(minorStr, 10); + const versionOk = + !Number.isNaN(major) && + (major > minMajor || + (major === minMajor && !Number.isNaN(minor) && minor >= minMinor)); + + if (!versionOk) { + throw new Error( + `Arrow support requires "${pkgName}" >= ${minVersion} (found ${pkg.version || 'unknown'}). ` + + `Please upgrade using: npm install ${pkgName}@latest`, + ); + } + return mod; +} + +/** + * Loads `apache-arrow` and verifies that the optional + * `@google-cloud/bigquery-storage` (>= 6.1.0) package is installed. + * + * `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 { + 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. + // 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; +} + +/** + * Converts a serialized Arrow IPC payload (base64 string or Uint8Array) to a + * Uint8Array. + * + * @private + */ +function toUint8Array(serialized: Uint8Array | string): Uint8Array { + return typeof serialized === 'string' + ? Buffer.from(serialized, 'base64') + : serialized; +} + +/** + * Creates a 0-row `RecordBatch` that preserves the provided Arrow `Schema`. + * + * @private + */ +export function createEmptyRecordBatch( + arrow: ArrowModule, + schema: Schema, +): RecordBatch { + // Use makeBuilder so apache-arrow allocates valid 0-length typed array + // buffers for every column (including nested Struct and List fields). + const structData = arrow + .makeBuilder({type: new arrow.Struct(schema.fields)}) + .finish() + .flush(); + return new arrow.RecordBatch(schema, structData); +} + +/** + * Decodes serialized Arrow IPC schema and record batch bytes into `RecordBatch` + * objects, emitting a single 0-row `RecordBatch` (preserving `schema` when + * present) if no record batches are present. + * + * @private + */ +export function decodeArrowPayload( + arrow: ArrowModule, + serializedSchema?: Uint8Array | string, + serializedRecordBatch?: Uint8Array | string, +): RecordBatch[] { + if (!serializedSchema) { + return [new arrow.RecordBatch({})]; + } + // BigQuery returns the Arrow IPC Schema message and RecordBatch message as + // separate fields; the reader accepts them as consecutive chunks of one IPC + // stream, which avoids copying them into a single buffer. + const chunks = serializedRecordBatch + ? [toUint8Array(serializedSchema), toUint8Array(serializedRecordBatch)] + : [toUint8Array(serializedSchema)]; + const reader = arrow.RecordBatchReader.from(chunks); + reader.open(); + // With no record batch, the reader yields its own 0-row placeholder batch, + // which leaves nested Struct children empty and List children undefined. + return serializedRecordBatch + ? [...reader] + : [createEmptyRecordBatch(arrow, reader.schema)]; +} + +/** + * Executes a query and returns a stream of `apache-arrow` `RecordBatch` objects. + * + * @private + */ +export function createQueryArrowStream_( + bq: BigQuery, + query?: Query | string, +): ArrowRecordBatchStream { + const stream = new ResourceStream( + {}, + () => {}, + ) 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; + } + + if (!query || (typeof query === 'object' && !query.query)) { + stream.destroy(new Error('A SQL query string is required.')); + return; + } + + 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.', + ), + ); + return; + } + + if (queryObj.destination) { + // TODO: Support explicit destination tables via CreateReadSession. + stream.destroy( + 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; + } + 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; + } + + // Default to JOB_CREATION_OPTIONAL unless overridden on the query or client. + if (!queryReq.jobCreationMode) { + queryReq.jobCreationMode = 'JOB_CREATION_OPTIONAL'; + } + queryReq.queryResultsFormat = 'ARROW'; + + void bq.runJobsQuery(queryReq, (err, job, res) => { + if (err) { + stream.destroy(err); + return; + } + + 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; + } + + // 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; + } + + 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(); + } + } catch (parseErr) { + stream.destroy(parseErr as Error); + } + }); + }; + + return stream; +} diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index d397e71991da..f313c08846ce 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -56,9 +56,10 @@ import { } from '@google-cloud/common/build/src/util'; import bigquery from './types'; import {logger, setLogFunction} from './logger'; +import {ArrowRecordBatchStream, createQueryArrowStream_} from './arrow'; // Third-Party Re-exports -export {common}; +export {common, ArrowRecordBatchStream}; export interface RequestCallback { (err: Error | null, response?: T | null): void; @@ -375,6 +376,52 @@ export class BigQuery extends Service { return new ResourceStream({}, () => {}); } + /** + * [Beta] Run a query scoped to your project and stream results as Apache + * Arrow `RecordBatch` objects. + * + * This method uses the `ARROW` query results format of the BigQuery API, + * which isn't generally available yet, so it may change in ways that aren't + * backward compatible. + * + * Requires the optional `@google-cloud/bigquery-storage` (>= 6.1.0) package + * to be installed. + * + * @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 + * stream emits a `job` event with the {@link Job} when BigQuery creates + * a job, and sets `stream.queryId` and `stream.jobReference` (when + * available) before the first batch. + * + * @example + * ``` + * const {BigQuery} = require('@google-cloud/bigquery'); + * const bigquery = new BigQuery(); + * + * const query = 'SELECT url FROM `publicdata.samples.github_nested` LIMIT 100'; + * + * bigquery.createQueryArrowStream(query) + * .on('error', console.error) + * .on('data', function(batch) { + * // batch is an apache-arrow RecordBatch from your query. + * }) + * .on('end', function() { + * // All batches retrieved. + * }); + * + * //- + * // You can also consume the stream using `for await...of` or `toArray()`. + * //- + * const batches = await bigquery.createQueryArrowStream(query).toArray(); + * ``` + * + * @beta + */ + createQueryArrowStream(options?: Query | string): ArrowRecordBatchStream { + return createQueryArrowStream_(this, options); + } + getDatasetsStream(options?: GetDatasetsOptions): ResourceStream { // placeholder body, overwritten in constructor return new ResourceStream({}, () => {}); @@ -2310,8 +2357,10 @@ export class BigQuery extends Service { * @param query string | Query * @param options QueryOptions * @returns bigquery.IQueryRequest | undefined + * + * @private */ - private buildQueryRequest_( + buildQueryRequest_( query: string | Query, options: QueryOptions, ): bigquery.IQueryRequest | undefined { @@ -2413,7 +2462,12 @@ export class BigQuery extends Service { return req; } - private runJobsQuery( + /** + * Runs a `jobs.query` request. + * + * @private + */ + runJobsQuery( req: bigquery.IQueryRequest, callback?: JobsQueryCallback, ): void | Promise { diff --git a/handwritten/bigquery/src/index.ts b/handwritten/bigquery/src/index.ts index ee56a62a6b3f..314ec9f5a97a 100644 --- a/handwritten/bigquery/src/index.ts +++ b/handwritten/bigquery/src/index.ts @@ -15,6 +15,7 @@ */ export { + ArrowRecordBatchStream, BigQuery, BigQueryDate, BigQueryDateOptions, diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts new file mode 100644 index 000000000000..0fedc298f1f7 --- /dev/null +++ b/handwritten/bigquery/system-test/arrow.ts @@ -0,0 +1,182 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import * as assert from 'assert'; +import {randomUUID} from 'crypto'; +import {after, before, describe, it} from 'mocha'; +import * as arrow from 'apache-arrow'; +import {BigQuery, Job} from '../src'; + +const bigquery = new BigQuery(); +const GCLOUD_TESTS_PREFIX = 'nodejs_bq_test'; + +describe('Arrow Query Results', () => { + // One dataset for the tests that need tables. It's labeled like the main + // system test dataset so that suite's stale-dataset cleanup also removes it + // if a run is interrupted before after() runs. + const dataset = bigquery.dataset( + `${GCLOUD_TESTS_PREFIX}_arrow_${randomUUID()}`.replace(/-/g, '_'), + ); + + before(async () => { + await dataset.create({labels: {[GCLOUD_TESTS_PREFIX]: ''}}); + }); + + after(async () => { + await dataset.delete({force: true}); + }); + + it('should stream inline Arrow RecordBatches via jobs.query with JOB_CREATION_OPTIONAL and support cached results', async () => { + const sql = + 'SELECT 1 AS id, "alpha" AS label UNION ALL SELECT 2 AS id, "beta" AS label ORDER BY id'; + const [structRows] = await bigquery.query(sql); + + for (const useQueryCache of [false, true]) { + const stream = bigquery.createQueryArrowStream({ + query: sql, + useQueryCache, + }); + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify queryId is populated and decoded Arrow columns match STRUCT_ENCODING rows. + assert.ok(stream.queryId); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, structRows.length); + assert.deepStrictEqual( + Array.from(table.getChild('id')!.toArray()), + structRows.map((r: {id: number}) => BigInt(r.id)), + ); + assert.deepStrictEqual( + table.getChild('label')!.toArray(), + structRows.map((r: {label: string}) => r.label), + ); + } + }); + + it('should populate jobReference and emit job event when jobCreationMode is JOB_CREATION_REQUIRED', async () => { + const stream = bigquery.createQueryArrowStream({ + query: 'SELECT 42 AS answer', + jobCreationMode: 'JOB_CREATION_REQUIRED', + }); + + let emittedJob: Job | undefined; + stream.on('job', (job: Job) => { + emittedJob = job; + }); + + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify jobReference and the 'job' event were populated alongside the inline batch. + assert.ok(stream.jobReference?.jobId); + assert.ok(emittedJob); + assert.strictEqual(emittedJob.id, stream.jobReference?.jobId); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, 1); + assert.deepStrictEqual(Array.from(table.getChild('answer')!.toArray()), [ + 42n, + ]); + }); + + it('should emit a 0-row RecordBatch preserving the schema for an empty result set', async () => { + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream( + 'SELECT * FROM UNNEST(ARRAY>[])', + ) + .toArray(); + + // Verify a single 0-row RecordBatch is emitted with the query's column names intact. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.deepStrictEqual( + batches[0].schema.fields.map(f => f.name), + ['id', 'name'], + ); + assert.ok(arrow.DataType.isInt(batches[0].schema.fields[0].type)); + assert.ok(arrow.DataType.isUtf8(batches[0].schema.fields[1].type)); + }); + + it('should emit a 0-row, 0-column RecordBatch for a statement without a result schema', async () => { + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream('DECLARE x INT64 DEFAULT 1;') + .toArray(); + + // Verify a single 0-row, 0-column RecordBatch is emitted when no schema is returned. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.strictEqual(batches[0].numCols, 0); + }); + + it('should run DDL and DML statements and emit a single 0-row RecordBatch for each', async () => { + // DDL and DML return no result rows, but the stream must still run the + // statement and end with one 0-row RecordBatch. + const tableRef = `\`${dataset.id}.statements\``; + for (const statement of [ + `CREATE TABLE ${tableRef} AS SELECT 1 AS x`, + `UPDATE ${tableRef} SET x = 2 WHERE TRUE`, + ]) { + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream(statement) + .toArray(); + assert.strictEqual(batches.length, 1, statement); + assert.strictEqual(batches[0].numRows, 0, statement); + } + + // Verify both statements ran: UPDATE changed the row CREATE TABLE wrote. + const [rows] = await bigquery.query(`SELECT x FROM ${tableRef}`); + assert.deepStrictEqual(rows, [{x: 2}]); + }); + + it('should reject WRITE_APPEND before running the query and leave the destination table unchanged', async () => { + // Reading a WRITE_APPEND destination would also return the rows already + // in the table, so the stream rejects it before sending the query. + const tableRef = `\`${dataset.id}.append_target\``; + await bigquery.query(`CREATE TABLE ${tableRef} AS SELECT 1 AS x`); + + await assert.rejects( + bigquery + .createQueryArrowStream({ + query: 'SELECT 2 AS x', + destination: dataset.table('append_target'), + writeDisposition: 'WRITE_APPEND', + }) + .toArray(), + /WRITE_APPEND is not supported with Arrow query results/, + ); + + // Verify the query never ran: the table still holds only its original row. + const [rows] = await bigquery.query(`SELECT x FROM ${tableRef}`); + assert.deepStrictEqual(rows, [{x: 1}]); + }); + + it('should raise the same error as query() when a query fails during execution', async () => { + const sql = 'SELECT 1 / 0'; + let structErr: (Error & {code?: number}) | undefined; + try { + await bigquery.query(sql); + } catch (err) { + structErr = err as Error & {code?: number}; + } + assert.ok(structErr); + + // Verify createQueryArrowStream raises an error with the same status code and message. + await assert.rejects( + bigquery.createQueryArrowStream(sql).toArray(), + (arrowErr: Error & {code?: number}) => { + assert.strictEqual(arrowErr.code, structErr!.code); + assert.strictEqual(arrowErr.message, structErr!.message); + return true; + }, + ); + }); +}); diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts new file mode 100644 index 000000000000..088a395b6d48 --- /dev/null +++ b/handwritten/bigquery/test/arrow.ts @@ -0,0 +1,435 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import * as assert from 'assert'; +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 bigquery from '../src/types'; + +/** + * Stand-in for `@google-cloud/bigquery-storage`, so the unit tests don't depend + * on which copy of the package is installed (for example, an unbuilt in-repo + * workspace link). + */ +const storageStub = {}; + +/** + * Serializes an Arrow Table into separate base64 `serializedSchema` and + * `serializedRecordBatch` strings matching BigQuery's REST `jobs.query` format. + */ +function serializeArrowTable(table: arrow.Table): { + serializedSchema: string; + serializedRecordBatch: string; +} { + const emptyTable = new arrow.Table(table.schema); + const schemaWithEos = Buffer.from(arrow.tableToIPC(emptyTable, 'stream')); + // Strip the 8-byte Arrow IPC end-of-stream marker so record batches can follow. + const schemaBuf = schemaWithEos.subarray(0, schemaWithEos.length - 8); + + const fullStream = Buffer.from(arrow.tableToIPC(table, 'stream')); + const batchBuf = fullStream.subarray(schemaBuf.length, fullStream.length - 8); + return { + serializedSchema: schemaBuf.toString('base64'), + serializedRecordBatch: batchBuf.toString('base64'), + }; +} + +function makeSampleTable(ids: number[], names: string[]): arrow.Table { + return new arrow.Table({ + id: arrow.makeVector(new Int32Array(ids)), + name: arrow.vectorFromArray(names, new arrow.Utf8()), + }); +} + +// Shared fixture used across inline jobs.query unit tests. +const SAMPLE_TABLE = makeSampleTable([1, 2], ['alice', 'bob']); +const SAMPLE_PAYLOAD = serializeArrowTable(SAMPLE_TABLE); + +/** + * Stubs `bq.request` for `/queries` (`jobs.query`) to return the given response + * or error and returns a getter for the captured `IQueryRequest`. + */ +function stubJobsQuery( + bq: BigQuery, + outcome: { + err?: Error; + job?: Job | null; + res?: bigquery.IQueryResponse; + }, +): () => bigquery.IQueryRequest | undefined { + let capturedReq: bigquery.IQueryRequest | undefined; + if (outcome.job) { + bq.job = () => outcome.job!; + } + bq.request = (reqOpts, callback) => { + capturedReq = reqOpts.json; + callback(outcome.err || null, outcome.res || {}); + }; + return () => capturedReq; +} + +describe('Arrow Support', () => { + let originalRequire: typeof Module.prototype.require; + + before(() => { + originalRequire = Module.prototype.require; + Module.prototype.require = Object.assign(function ( + this: NodeModule, + id: string, + ) { + if (id === '@google-cloud/bigquery-storage') { + return storageStub; + } + if (id === '@google-cloud/bigquery-storage/package.json') { + return {version: '6.1.0'}; + } + return originalRequire.call(this, id); + }, originalRequire); + }); + + after(() => { + Module.prototype.require = originalRequire; + }); + + 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. + assert.strictEqual(typeof arrowMod.RecordBatchReader.from, 'function'); + }); + + it('should throw an actionable error when @google-cloud/bigquery-storage is missing', () => { + const pkg = '@google-cloud/bigquery-storage'; + const fakeRequire = ((id: string) => { + if (id.startsWith(pkg)) { + throw new Error(`Cannot find module '${id}'`); + } + return require(id); + }) as NodeRequire; + + // Verify the error tells the user how to install the missing package. + assert.throws( + () => checkArrowDependencies(fakeRequire), + /npm install @google-cloud\/bigquery-storage/, + ); + }); + + it('should throw an actionable error when @google-cloud/bigquery-storage is below the minimum version', () => { + const fakeRequire = ((id: string) => { + if (id === '@google-cloud/bigquery-storage/package.json') { + return {version: '6.0.0'}; + } + return require(id); + }) as NodeRequire; + + // Verify outdated package versions are rejected with upgrade instructions. + assert.throws( + () => checkArrowDependencies(fakeRequire), + /requires "@google-cloud\/bigquery-storage" >= 6\.1\.0 \(found 6\.0\.0\)/, + ); + }); + + it('should include the underlying error when an installed package fails to load', () => { + const loadError = new Error( + "Cannot find module '/repo/bigquery-storage/build/src/index.js'", + ); + const fakeRequire = ((id: string) => { + if (id === '@google-cloud/bigquery-storage') { + throw loadError; + } + return require(id); + }) as NodeRequire; + + // Verify the load failure is reported instead of being masked. + assert.throws( + () => checkArrowDependencies(fakeRequire), + (err: Error) => + err.cause === loadError && err.message.includes(loadError.message), + ); + }); + }); + + describe('BigQuery#createQueryArrowStream (inline jobs.query)', () => { + let bq: BigQuery; + + beforeEach(() => { + bq = new BigQuery({projectId: 'test-project'}); + }); + + it('should default jobCreationMode to JOB_CREATION_OPTIONAL and set queryResultsFormat to ARROW', async () => { + const getCapturedReq = stubJobsQuery(bq, { + res: { + jobComplete: true, + queryId: 'stateless-query-123', + totalRows: '2', + pageRowCount: '2', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, + }); + + const stream = bq.createQueryArrowStream( + 'SELECT 1 AS id, "alice" AS name', + ); + let jobEmitted = false; + stream.on('job', () => { + jobEmitted = true; + }); + + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify the request sent to jobs.query used JOB_CREATION_OPTIONAL and ARROW format. + assert.strictEqual( + getCapturedReq()?.jobCreationMode, + 'JOB_CREATION_OPTIONAL', + ); + assert.strictEqual(getCapturedReq()?.queryResultsFormat, 'ARROW'); + + // Verify stateless query metadata is populated without emitting a job event. + assert.strictEqual(stream.queryId, 'stateless-query-123'); + assert.strictEqual(stream.jobReference, undefined); + assert.strictEqual(jobEmitted, false); + + // Verify the decoded Arrow RecordBatch contents match the serialized table. + assert.strictEqual(batches.length, 1); + const resultTable = new arrow.Table(batches); + assert.deepStrictEqual( + Array.from(resultTable.getChild('id')!.toArray()), + [1, 2], + ); + assert.deepStrictEqual(resultTable.getChild('name')?.toArray(), [ + 'alice', + 'bob', + ]); + }); + + it('should preserve client defaultJobCreationMode and emit job when jobReference is returned', async () => { + bq = new BigQuery({ + projectId: 'test-project', + defaultJobCreationMode: 'JOB_CREATION_REQUIRED', + }); + + const expectedJobRef: bigquery.IJobReference = { + projectId: 'test-project', + jobId: 'job-abc', + location: 'US', + }; + const fakeJob = bq.job('job-abc', {location: 'US'}); + const getCapturedReq = stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: true, + queryId: 'query-abc', + jobReference: expectedJobRef, + totalRows: '2', + pageRowCount: '2', + arrowSchema: {serializedSchema: SAMPLE_PAYLOAD.serializedSchema}, + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, + }); + + const stream = bq.createQueryArrowStream( + 'SELECT 1 AS id, "alice" AS name', + ); + let emittedJob: Job | undefined; + stream.on('job', (j: Job) => { + emittedJob = j; + }); + + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify the client's defaultJobCreationMode was not overwritten. + assert.strictEqual( + getCapturedReq()?.jobCreationMode, + 'JOB_CREATION_REQUIRED', + ); + // Verify stream.jobReference and the 'job' event were populated. + assert.deepStrictEqual(stream.jobReference, expectedJobRef); + assert.strictEqual(emittedJob, fakeJob); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + + it('should emit a 0-row RecordBatch preserving the schema when arrowSchema is present without arrowRecordBatch', async () => { + // Include nested Struct and List fields to verify 0-row buffer initialization. + const schema = new arrow.Schema([ + new arrow.Field('id', new arrow.Int64()), + new arrow.Field( + 'nested', + new arrow.Struct([new arrow.Field('sub', new arrow.Int32())]), + ), + new arrow.Field( + 'tags', + new arrow.List(new arrow.Field('item', new arrow.Utf8())), + ), + ]); + const {serializedSchema} = serializeArrowTable(new arrow.Table(schema)); + + stubJobsQuery(bq, { + res: { + jobComplete: true, + queryId: 'empty-query', + totalRows: '0', + arrowSchema: {serializedSchema}, + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, nested, tags FROM tbl WHERE FALSE') + .toArray(); + + // Verify a single 0-row RecordBatch is emitted with the full schema intact + // and valid empty column vectors. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.deepStrictEqual( + batches[0].schema.fields.map(f => f.name), + ['id', 'nested', 'tags'], + ); + assert.strictEqual(batches[0].getChild('id')?.toArray().length, 0); + assert.strictEqual(batches[0].getChild('nested')?.toArray().length, 0); + // Verify nested Struct children are readable; apache-arrow's own 0-row + // placeholder batch has none, so this throws if it is returned. + assert.strictEqual( + batches[0].getChild('nested')?.getChild('sub')?.toArray().length, + 0, + ); + assert.strictEqual(batches[0].getChild('tags')?.toArray().length, 0); + }); + + it('should emit a 0-row, 0-column RecordBatch when both arrowSchema and arrowRecordBatch are absent', async () => { + stubJobsQuery(bq, { + res: { + jobComplete: true, + queryId: 'ddl-query', + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('DECLARE x INT64 DEFAULT 1;') + .toArray(); + + // Verify DDL/script queries emit a single 0-row, 0-column RecordBatch. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.strictEqual(batches[0].numCols, 0); + }); + + it('should fail before sending the query when writeDisposition is WRITE_APPEND', async () => { + const getCapturedReq = stubJobsQuery(bq, {}); + + // Verify WRITE_APPEND is rejected before executing the query on BigQuery. + await assert.rejects( + bq + .createQueryArrowStream({ + query: 'SELECT 1', + writeDisposition: 'WRITE_APPEND', + }) + .toArray(), + /WRITE_APPEND is not supported with Arrow query results/, + ); + assert.strictEqual(getCapturedReq(), undefined); + }); + + it('should fail when a SQL query string is not provided', async () => { + // Verify missing query string emits an error on the stream. + await assert.rejects( + bq.createQueryArrowStream({} as Query).toArray(), + /A SQL query string is required\./, + ); + }); + + it('should forward jobs.query errors to the stream error event', async () => { + stubJobsQuery(bq, {err: new Error('syntax error in SQL')}); + + // Verify API errors from jobs.query reject stream consumption. + await assert.rejects( + bq.createQueryArrowStream('SELECT INVALID').toArray(), + /syntax error in SQL/, + ); + }); + + it('should fail when an explicit destination table is specified', async () => { + const destination = bq.dataset('ds').table('tbl'); + + // TODO: Replace with explicit destination table (CreateReadSession) tests once added. + await assert.rejects( + bq.createQueryArrowStream({query: 'SELECT 1', destination}).toArray(), + /Explicit destination tables are not yet supported in createQueryArrowStream\./, + ); + }); + + it('should fail when the query configuration cannot be executed via jobs.query', async () => { + // TODO: Replace with jobs.insert (createQueryJob) fallback tests once added. + await assert.rejects( + bq + .createQueryArrowStream({query: 'SELECT 1', jobId: 'custom-job-1'}) + .toArray(), + /Query configuration cannot be executed via jobs\.query\./, + ); + }); + + it('should forward errors thrown while building the jobs.query request', async () => { + // A null parameter without an explicit type makes buildQueryRequest_ throw. + await assert.rejects( + bq + .createQueryArrowStream({query: 'SELECT ?', params: [null]}) + .toArray(), + /Parameter types must be provided for null values/, + ); + }); + + it('should fail when jobs.query returns jobComplete = false', async () => { + stubJobsQuery(bq, {res: {jobComplete: false}}); + + // TODO: Replace with job wait (jobComplete === false) tests once unfinished job polling is added. + await assert.rejects( + bq.createQueryArrowStream('SELECT 1').toArray(), + /Query did not complete within the jobs\.query timeout\./, + ); + }); + + it('should fail when jobs.query does not return all rows inline', async () => { + stubJobsQuery(bq, { + res: {jobComplete: true, pageRowCount: '2', totalRows: '5'}, + }); + + // TODO: Replace with default job stream (pageRowCount < totalRows) tests once added. + await assert.rejects( + bq.createQueryArrowStream('SELECT 1').toArray(), + /Query results were not returned inline in the jobs\.query response\./, + ); + }); + + it('should fail when jobs.query returns a pageToken with totalRows = 0', async () => { + // Cached multi-page results can report totalRows "0" and omit pageRowCount. + stubJobsQuery(bq, { + res: {jobComplete: true, totalRows: '0', pageToken: 'token-1'}, + }); + + // TODO: Replace with default job stream (pageToken) tests once added. + await assert.rejects( + bq.createQueryArrowStream('SELECT 1').toArray(), + /Query results were not returned inline in the jobs\.query response\./, + ); + }); + }); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 08a3c497a09c..62f973509c54 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1598,6 +1598,9 @@ importers: '@google-cloud/promisify': specifier: ^6.0.0 version: 6.0.1 + apache-arrow: + specifier: ^21.0.0 + version: 21.2.0 arrify: specifier: ^3.0.0 version: 3.0.0 @@ -1617,6 +1620,9 @@ importers: specifier: ^11.0.0 version: 11.0.1 devDependencies: + '@google-cloud/bigquery-storage': + specifier: ^6.1.0 + version: 6.2.0(protobufjs@8.8.0) '@google-cloud/storage': specifier: ^7.16.0 version: 7.21.0(encoding@0.1.13) @@ -16408,10 +16414,20 @@ packages: '@foliojs-fork/restructure@2.0.2': resolution: {integrity: sha512-59SgoZ3EXbkfSX7b63tsou/SDGzwUEK6MuB5sKqgVK1/XE0fxmpsOb9DQI8LXW3KfGnAjImCGhhEb7uPPAUVNA==} + '@google-cloud/bigquery-storage-api@0.3.0': + resolution: {integrity: sha512-PPBOLw5IKE2wcdz/U4DV4KjeGHmt2Najwq73wrx9Jg0nHJor70jgsrrgdRNJ/4UvZ33ALhH55w0wmr2qk7PFrg==} + engines: {node: '>=22'} + '@google-cloud/bigquery-storage-api@0.5.0': resolution: {integrity: sha512-fN5uZpF13J61ukset7xhaHUweADtqTir5vI9KBq0CnRqirkEnKIvZxn34hoeO2NeP+rk8Ro2SeMVySJ6i16LuA==} engines: {node: '>=22'} + '@google-cloud/bigquery-storage@6.2.0': + resolution: {integrity: sha512-ME41UGUxJJWywspFjcSGkq3ZE8izgBkmjLIyzZ1xqG0wucm9FP6XN+CHzT0vs3GEARguJEt8tuMfvJ7BJH328w==} + engines: {node: '>=22'} + peerDependencies: + protobufjs: ~7.5.5 + '@google-cloud/bigquery@9.0.2': resolution: {integrity: sha512-KLRSKdU+1VsXMDz9tmqtFASQXuBazLFN3VN8fTkuW/jDWWQK7HuNjT2RIKev+5xlBGfBsqQZs4Z5jqT0seV0aw==} engines: {node: '>=22'} @@ -26013,12 +26029,33 @@ snapshots: '@foliojs-fork/restructure@2.0.2': {} + '@google-cloud/bigquery-storage-api@0.3.0': + dependencies: + google-gax: 6.5.0 + transitivePeerDependencies: + - supports-color + '@google-cloud/bigquery-storage-api@0.5.0': dependencies: google-gax: 6.1.0 transitivePeerDependencies: - supports-color + '@google-cloud/bigquery-storage@6.2.0(protobufjs@8.8.0)': + dependencies: + '@google-cloud/bigquery-storage-api': 0.3.0 + '@google-cloud/paginator': 7.0.1 + '@google-cloud/precise-date': 6.0.0 + apache-arrow: 21.2.0 + core-js: 3.49.0 + extend: 3.0.2 + google-auth-library: 11.0.2 + google-gax: 6.5.0 + long: 5.3.2 + protobufjs: 8.8.0 + transitivePeerDependencies: + - supports-color + '@google-cloud/bigquery@9.0.2': dependencies: '@google-cloud/common': 8.0.1