Repository navigation
feat(bigquery): add createQueryArrowStream for inline jobs.query results - #9596
Conversation
product-auto-label
Bot
added
the
api: bigquery
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Apache Arrow support to the BigQuery client library, allowing users to stream query results as Arrow RecordBatch objects. It adds the necessary dependencies, implements the Arrow streaming pipeline, and exposes the createQueryArrowStream method on the BigQuery class. The feedback suggests addressing potential silent data loss by explicitly handling asynchronous query jobs and paginated results, as well as simplifying the dependency verification logic.
71dff93 to
af4b542
Compare
af4b542 to
36ae146
Compare
36ae146 to
2ec010b
Compare
bed860e to
a756517
Compare
a9ab952 to
9b3d5e9
Compare
| "@google-cloud/paginator": "^7.0.0", | ||
| "@google-cloud/precise-date": "^6.0.0", | ||
| "@google-cloud/promisify": "^6.0.0", | ||
| "apache-arrow": "^21.0.0", |
There was a problem hiding this comment.
All arrow dependencies (apache-arrow and @google-cloud/bigquery-storage) should be under peerDependencies with peerDependenciesMeta marking them "optional": true (while keeping both in devDependencies for CI) so we bundle Arrow support as an opt-in "plugin" and don't install apache-arrow by default for non-Arrow users. (Note: peerDependencies without peerDependenciesMeta are auto-installed by npm v7+.)
https://docs.npmjs.com/cli/v11/configuring-npm/package-json#peerdependenciesmeta
There was a problem hiding this comment.
I considered this approach, but apache-arrow can't be an optional peer as long as our public typings reference it. ArrowRecordBatchStream is ResourceStream<RecordBatch>, and bigquery.d.ts re-exports it. TypeScript can't make that import optional, so any TS user who doesn't install apache-arrow breaks, even if they never call the Arrow API:
build/src/arrow.d.ts(17,42): error TS2307: Cannot find module 'apache-arrow' or its corresponding type declarations.
build/src/arrow.d.ts(20,34): error TS2307: Cannot find module 'apache-arrow' or its corresponding type declarations.
With skipLibCheck it compiles, but RecordBatch becomes any. pack-n-play (system-test/install.ts) caught exactly this, which is why I moved apache-arrow to dependencies.
To make it optional, we'd need to replace RecordBatch in the public typings with our own minimal interface, so the .d.ts files no longer import apache-arrow. That's about 45 lines, and we'd have to keep it in sync with apache-arrow's API. The trade-off is roughly 8.4 MB on install (flatbuffers, json-with-bigint, tslib) against users getting the real RecordBatch types. I went with the dependency for simplicity, but I'm happy to switch if you'd rather keep the install lean.
For @google-cloud/bigquery-storage I agree, and I'll declare it as an optional peer (peerDependenciesMeta.optional) while keeping it in devDependencies for CI.
There was a problem hiding this comment.
Added @google-cloud/bigquery-storage as an optional peer (kept in devDependencies for CI) in c1fc109. apache-arrow stays a regular dependency for the TS2307 reason above.
There was a problem hiding this comment.
Nit (for future consideration): Keeping apache-arrow in dependencies (while lazy-loading it at runtime) and moving @google-cloud/bigquery-storage to optional peerDependencies makes sense for now so pack-n-play doesn't fail with TS2307.
In the future, if we want to make apache-arrow a true optional peerDependency as well (so default installs don't download it at all), we could consider exposing Arrow via a subpath export in package.json:
- Add
"exports"topackage.jsonwith"."(./build/src/index.js) and"./arrow"(./build/src/arrow.js). - Stop re-exporting
./arrowfromsrc/index.tsandsrc/bigquery.tsso the root.d.tsfiles neverimport type {RecordBatch} from 'apache-arrow'. - Optionally use TypeScript module augmentation (
declare module './bigquery' { interface BigQuery { ... } }) insidesrc/arrow.tsif we also wantbq.queryAsArrowStream()typed when a user imports@google-cloud/bigquery/arrow.
That way, TypeScript only resolves apache-arrow types for users who actually import @google-cloud/bigquery/arrow, and pack-n-play stays green without apache-arrow installed. Totally fine to keep the current setup for this PR stack!
| function loadPackageJson( | ||
| requireFn: NodeRequire, | ||
| pkgName: string, | ||
| ): {version?: string} { | ||
| try { | ||
| return requireFn(`${pkgName}/package.json`); | ||
| } catch { | ||
| const entryPath = (requireFn.resolve || require.resolve)(pkgName); | ||
| return requireFn(path.join(path.dirname(entryPath), 'package.json')); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Requires a dependency and verifies its minimum major.minor version. | ||
| * | ||
| * @private | ||
| */ | ||
| function requireWithMinVersion<T = unknown>( | ||
| 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 = loadPackageJson(requireFn, pkgName); | ||
| } 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; | ||
| } | ||
|
|
||
| /** | ||
| * Verifies that `apache-arrow` (>= 21.0.0) and the optional | ||
| * `@google-cloud/bigquery-storage` (>= 6.1.0) package are installed and meet | ||
| * minimum version requirements. | ||
| * | ||
| * @private | ||
| */ | ||
| export function checkArrowDependencies( | ||
| requireFn: NodeRequire = require, | ||
| ): ArrowModule { | ||
| const arrow = requireWithMinVersion<ArrowModule>( | ||
| requireFn, | ||
| 'apache-arrow', | ||
| 21, | ||
| ); | ||
| // 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; | ||
| } |
There was a problem hiding this comment.
Saw your note that loadPackageJson will be removed once bigquery-storage is released.
Do we need requireWithMinVersion at all once 6.3.0 is published? Note that require('apache-arrow/package.json') also throws MODULE_NOT_FOUND because apache-arrow's exports map rewrites ./* to ./*.js. If both packages are declared in peerDependencies, the package manager handles version constraints at install time and we only need a simple lazy require() here.
There was a problem hiding this comment.
Good catch on the exports map. That also means my earlier note was wrong: loadPackageJson wasn't test-only, it ran for every user because of apache-arrow.
Since apache-arrow is a regular ^21 dependency, I dropped its version check, which removes loadPackageJson entirely (2d466a9). It's now a plain lazy require().
I kept the check for bigquery-storage, though. Package managers only warn when an installed version is outside an optional peer range (I checked npm 10 too), and an older version doesn't fail loudly: its protos silently drop newer request fields like arrowSerializationOptions. The check turns that into a clear error.
| function toBuffer(serialized: Uint8Array | string): Buffer { | ||
| if (typeof serialized === 'string') { | ||
| return Buffer.from(serialized, 'base64'); | ||
| } | ||
| return Buffer.from( | ||
| serialized.buffer, | ||
| serialized.byteOffset, | ||
| serialized.byteLength, | ||
| ); | ||
| } | ||
|
|
||
| /** | ||
| * 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 [createEmptyRecordBatch(arrow, new arrow.Schema([]))]; | ||
| } | ||
| // BigQuery returns the Arrow IPC Schema message and RecordBatch message as | ||
| // separate fields; concatenate them into a single IPC stream for the reader. | ||
| const schemaBuf = toBuffer(serializedSchema); | ||
| const ipcBuf = serializedRecordBatch | ||
| ? Buffer.concat([schemaBuf, toBuffer(serializedRecordBatch)]) | ||
| : schemaBuf; | ||
| const reader = arrow.RecordBatchReader.from(ipcBuf); | ||
| reader.open(); | ||
| const schema = reader.schema; | ||
| const batches = [...reader]; | ||
| return batches.length > 0 ? batches : [createEmptyRecordBatch(arrow, schema)]; | ||
| } |
There was a problem hiding this comment.
We can pass [schemaBuf, batchBuf] directly to arrow.RecordBatchReader.from to avoid copying buffers with Buffer.concat. Also, is createEmptyRecordBatch needed? RecordBatchReader.from([schemaBuf]).readAll() already returns a 0-row RecordBatch with the schema when no record batch is present, and new arrow.RecordBatch({}) handles the empty schema case:
| function toBuffer(serialized: Uint8Array | string): Buffer { | |
| if (typeof serialized === 'string') { | |
| return Buffer.from(serialized, 'base64'); | |
| } | |
| return Buffer.from( | |
| serialized.buffer, | |
| serialized.byteOffset, | |
| serialized.byteLength, | |
| ); | |
| } | |
| /** | |
| * 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 [createEmptyRecordBatch(arrow, new arrow.Schema([]))]; | |
| } | |
| // BigQuery returns the Arrow IPC Schema message and RecordBatch message as | |
| // separate fields; concatenate them into a single IPC stream for the reader. | |
| const schemaBuf = toBuffer(serializedSchema); | |
| const ipcBuf = serializedRecordBatch | |
| ? Buffer.concat([schemaBuf, toBuffer(serializedRecordBatch)]) | |
| : schemaBuf; | |
| const reader = arrow.RecordBatchReader.from(ipcBuf); | |
| reader.open(); | |
| const schema = reader.schema; | |
| const batches = [...reader]; | |
| return batches.length > 0 ? batches : [createEmptyRecordBatch(arrow, schema)]; | |
| } | |
| function toUint8Array(serialized: Uint8Array | string): Uint8Array { | |
| return typeof serialized === 'string' | |
| ? Buffer.from(serialized, 'base64') | |
| : serialized; | |
| } | |
| /** | |
| * 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({})]; | |
| } | |
| const chunks = serializedRecordBatch | |
| ? [toUint8Array(serializedSchema), toUint8Array(serializedRecordBatch)] | |
| : [toUint8Array(serializedSchema)]; | |
| return arrow.RecordBatchReader.from(chunks).readAll(); | |
| } |
There was a problem hiding this comment.
Took the chunk array and new arrow.RecordBatch({}) (caec977). No more Buffer.concat copy per batch.
On dropping createEmptyRecordBatch, you're right that the reader already returns a 0-row batch with the schema, but that batch is incomplete for nested columns. For a STRUCT column it has no child fields, so on an empty result:
batch.getChild('address').getChild('city');
// reader's batch: TypeError: Vector constructor expects an Array of Data instances
// createEmptyRecordBatch: empty vectorLIST columns have the same problem (the child data is undefined). Looking into this showed we had the bug too: [...reader] already returned that placeholder, so our fallback never ran. Fixed in 6c36eb1, and the unit test now reads into a nested field so it can't regress quietly.
| bq: BigQuery, | ||
| query?: Query | string, | ||
| ): ArrowRecordBatchStream { | ||
| const bqInternal = bq as unknown as BigQueryInternal; |
There was a problem hiding this comment.
Can we mark buildQueryRequest_ and runJobsQuery as @internal (non-private) on BigQuery instead of casting bq as unknown as BigQueryInternal?
There was a problem hiding this comment.
Done in 69bbe38: dropped private on both and removed the cast. I tagged them @private rather than @internal to match how the rest of BigQuery marks members shared across files (queryAsStream_, mergeSchemaWithRows_). Happy to switch to @internal if you prefer.
| return; | ||
| } | ||
|
|
||
| const queryReq = bqInternal.buildQueryRequest_(queryObj, {}); |
There was a problem hiding this comment.
Wrap buildQueryRequest_ in a try / catch and route any error to stream.destroy(err as Error):
| const queryReq = bqInternal.buildQueryRequest_(queryObj, {}); | |
| let queryReq: bigquery.IQueryRequest | undefined; | |
| try { | |
| queryReq = bqInternal.buildQueryRequest_(queryObj, {}); | |
| } catch (err) { | |
| stream.destroy(err as Error); | |
| return; | |
| } |
There was a problem hiding this comment.
Done in 60f6082, with a test for a null param without types.
| // pageRowCount and totalRows are returned as uint64 strings in REST JSON. | ||
| if (BigInt(res?.pageRowCount ?? 0) < BigInt(res?.totalRows ?? 0)) { |
There was a problem hiding this comment.
Also check Boolean(res?.pageToken) here—on cached multi-page queries, jobs.query can return totalRows: "0" and omit pageRowCount while setting res.pageToken:
| // pageRowCount and totalRows are returned as uint64 strings in REST JSON. | |
| if (BigInt(res?.pageRowCount ?? 0) < BigInt(res?.totalRows ?? 0)) { | |
| // pageRowCount and totalRows are returned as uint64 strings in REST JSON. | |
| if ( | |
| Boolean(res?.pageToken) || | |
| BigInt(res?.pageRowCount ?? 0) < BigInt(res?.totalRows ?? 0) | |
| ) { |
There was a problem hiding this comment.
Good catch, done in 27bed0b with a test for totalRows: '0' plus a pageToken. We've hit this too: on ARROW cache hits, jobs.query reports totalRows: "0" even when rows remain, so without this check the stream silently returned nothing.
|
#9617 is a relevant PR for the system tests. |
quirogas
left a comment
quirogas
left a comment
There was a problem hiding this comment.
LGTM! Thanks for addressing all the review comments.
|
I added one docs commit after your approval: |
ohmayr
added
the
do not merge
|
Adding a Googlers, see: b/573764383 |
Introduce bigquery.createQueryArrowStream(query) to execute SQL queries with queryResultsFormat set to ARROW and stream decoded Apache Arrow RecordBatch objects: - Add apache-arrow (>= 21.0.0) to dependencies so published TypeScript declarations resolve RecordBatch, and verify apache-arrow (>= 21.0.0) and optional @google-cloud/bigquery-storage (>= 6.1.0) lazily at runtime. - Default jobCreationMode to JOB_CREATION_OPTIONAL when not configured on the request or client. - Reject writeDisposition: WRITE_APPEND before executing the query. - Decode inline Arrow IPC schema and record batch payloads, emitting a 0-row RecordBatch preserving the schema (or 0-column schema) when no rows are returned.
On query cache hits whose results don't fit inline, jobs.query with queryResultsFormat ARROW returns totalRows "0" and omits pageRowCount while still setting pageToken. Comparing pageRowCount with totalRows alone saw 0 < 0, so createQueryArrowStream ended with no rows and no error. Also treat a non-empty pageToken as remaining rows, so these responses fail with the "not returned inline" error instead of silently returning an empty result.
List @google-cloud/bigquery-storage under peerDependencies with peerDependenciesMeta optional, so the supported range is published and package managers check it when the package is installed, without installing it for users who don't use Arrow. Keep it in devDependencies because optional peers aren't installed for tests.
createQueryArrowStream reached BigQuery's private buildQueryRequest_ and runJobsQuery through an `as unknown as BigQueryInternal` cast, the only such cast in src/. A cast also hides signature changes from the compiler. Drop the TypeScript private modifier on both methods and tag them @Private in JSDoc, matching how the package already shares internals across files (Table.formatMetadata_, BigQuery.mergeSchemaWithRows_), and call them directly from arrow.ts.
buildQueryRequest_ throws on invalid query parameters, for example a null parameter without a type. The throw happened inside _read and only reached the stream because Node's Readable catches it. Catch it explicitly and destroy the stream with the error, matching the checkArrowDependencies() handling, so the behavior doesn't depend on stream internals.
decodeArrowPayload joined the IPC schema and record batch messages with
Buffer.concat before decoding, copying every batch. Pass them to
RecordBatchReader.from() as consecutive chunks instead, and use
new RecordBatch({}) when the response has no schema.
For a response with a schema but no record batch, RecordBatchReader yields its own 0-row placeholder batch. decodeArrowPayload returned it as is, so the createEmptyRecordBatch fallback never ran. The placeholder has no Struct children and undefined List children, so reading a nested column of a 0-row result threw "Vector constructor expects an Array of Data instances". Build the batch with createEmptyRecordBatch whenever the response has no record batch, and assert in the unit test that nested children are readable.
apache-arrow is a regular ^21.0.0 dependency, so the install already
guarantees its version. Its runtime check also needed loadPackageJson,
because apache-arrow's exports map blocks require('apache-arrow/
package.json'). Load it with a plain lazy require and remove
loadPackageJson.
Keep the @google-cloud/bigquery-storage check. 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 such as
arrowSerializationOptions.
The ARROW query results format of the BigQuery API isn't generally available yet; the API reference still describes it as not yet available. Label createQueryArrowStream and ArrowRecordBatchStream as [Beta] with @beta, as the Java and Python clients label their Arrow query APIs, so callers know the method may change.
8f12a09 to
79608b0
Compare
|
The decision remains the same as what we have here. So no longer blocked. |
ohmayr
removed
the
do not merge
2 participants
Adds
BigQuery#createQueryArrowStream(query)to stream query results asapache-arrowRecordBatchobjects when all rows are returned inline byjobs.query(queryResultsFormat: 'ARROW').apache-arrowand verifies that@google-cloud/bigquery-storage(>= 6.1.0) is installed.jobCreationModetoJOB_CREATION_OPTIONALunless configured on the request or client, and rejectswriteDisposition: 'WRITE_APPEND'before sending the request.arrowSchemaandarrowRecordBatchIPC bytes intoRecordBatchobjects, emitting a 0-rowRecordBatchpreserving the schema (or a 0-column schema for statements without a result schema) when no rows are returned.createQueryArrowStreamandArrowRecordBatchStreamas[Beta](@beta), since theARROWresults format onjobs.queryisn't generally available yet.