diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 12c5b938dc3..81291718e8c 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -142,20 +142,29 @@ export function decodeArrowPayload( serializedSchema?: Uint8Array | string, serializedRecordBatch?: Uint8Array | string, ): RecordBatch[] { - if (!serializedSchema) { + const hasSchema = + serializedSchema !== undefined && serializedSchema.length > 0; + const hasBatch = + serializedRecordBatch !== undefined && serializedRecordBatch.length > 0; + if (!hasSchema) { + // Without a schema the batch can't be decoded; returning an empty batch + // would silently drop its rows. + if (hasBatch) { + throw new Error('Got an Arrow record batch before its schema.'); + } 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 + const chunks = hasBatch ? [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 + return hasBatch ? [...reader] : [createEmptyRecordBatch(arrow, reader.schema)]; } diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 85d4e1609a0..b313e8dbc0b 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -268,6 +268,46 @@ describe('Arrow Support', () => { assert.strictEqual(batches[0].numRows, 2); }); + it('should fail instead of dropping rows when jobs.query returns a record batch without a schema', async () => { + stubJobsQuery(bq, { + res: { + jobComplete: true, + totalRows: '2', + pageRowCount: '2', + arrowRecordBatch: { + serializedRecordBatch: SAMPLE_PAYLOAD.serializedRecordBatch, + }, + }, + }); + + // Verify the batch isn't silently replaced with an empty RecordBatch. + await assert.rejects( + bq.createQueryArrowStream('SELECT 1').toArray(), + /Got an Arrow record batch before its schema\./, + ); + }); + + it('should treat an empty record batch without a schema like a missing one', async () => { + stubJobsQuery(bq, { + res: { + jobComplete: true, + totalRows: '0', + arrowRecordBatch: { + // An empty Uint8Array is truthy but carries no rows. + serializedRecordBatch: new Uint8Array(0) as unknown as string, + }, + }, + }); + + // Verify one 0-row, 0-column RecordBatch is emitted instead of an error. + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT 1') + .toArray(); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.strictEqual(batches[0].numCols, 0); + }); + 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([