diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 8f9ea5dba71..8701835099e 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -484,7 +484,8 @@ async function createArrowReadSession( dataFormat: 'ARROW', readOptions: {arrowSerializationOptions: ctx.arrowSerializationOptions}, }, - // One stream keeps rows in the query's order (e.g., ORDER BY). + // One stream keeps rows in the query's order (e.g., ORDER BY) for a table + // the query just wrote. maxStreamCount: 1, }); return session ?? undefined; @@ -945,13 +946,6 @@ export function createQueryArrowStream_( throw new Error('maxResults must be a non-negative integer.'); } - if (queryObj.writeDisposition === 'WRITE_APPEND') { - 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.', - ); - } - const ctx: ArrowStreamContext = { bq, arrow, diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index 3d0e5e86e76..ced33b8cc1a 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -403,7 +403,9 @@ export class BigQuery extends Service { * `for await` does this for you. * * With an explicit `destination` table, the stream reads that whole table - * through a Storage Read API session once the query finishes. + * through a Storage Read API session once the query finishes. With + * `writeDisposition: 'WRITE_APPEND'`, that includes rows the table already + * had, as with {@link BigQuery#query}. * * @param {string|ArrowQuery} [options] A string SQL query or configuration * object. See {@link BigQuery#query} for a complete list of options. diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts index 3ea868aa7e3..2781172b106 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -138,28 +138,6 @@ describe('Arrow Query Results', () => { 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; @@ -961,4 +939,41 @@ describe('Arrow Query Results', () => { .toArray(); assert.strictEqual(new arrow.Table(allBatches).numRows, totalRows); }); + + it('should return existing and new rows for WRITE_APPEND and fail WRITE_EMPTY on a non-empty destination table', async () => { + const tableRef = `\`${dataset.id}.append_target\``; + const destTable = dataset.table('append_target'); + await bigquery.query(`CREATE TABLE ${tableRef} AS SELECT 1 AS x`); + + // 1. WRITE_APPEND adds the new row and, like query(), reads back the whole + // destination table, including the row that was already there. + const appendBatches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({ + query: 'SELECT 2 AS x', + destination: destTable, + writeDisposition: 'WRITE_APPEND', + }) + .toArray(); + const appendTable = new arrow.Table(appendBatches); + assert.deepStrictEqual( + Array.from(appendTable.getChild('x')!.toArray()).sort(), + [1n, 2n], + ); + + // 2. WRITE_EMPTY fails because the destination table already has rows. + await assert.rejects( + bigquery + .createQueryArrowStream({ + query: 'SELECT 3 AS x', + destination: destTable, + writeDisposition: 'WRITE_EMPTY', + }) + .toArray(), + /Already Exists/, + ); + + // Verify the failed WRITE_EMPTY job left the table unchanged. + const [rows] = await bigquery.query(`SELECT x FROM ${tableRef} ORDER BY x`); + assert.deepStrictEqual(rows, [{x: 1}, {x: 2}]); + }); }); diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 0bb1399212b..649132a2289 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -515,22 +515,6 @@ describe('Arrow Support', () => { 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( @@ -1954,6 +1938,54 @@ describe('Arrow Support', () => { assert.strictEqual(batches[0].numRows, 2); }); + it('should send WRITE_APPEND through createQueryJob and read the whole destination table through a read session', async () => { + const {destination, getCapturedJobOpts} = stubDestinationJob( + 'append_tbl', + 'dest-append-1', + ); + const sessionReqs: CreateReadSessionRequest[] = []; + stubStorageReadClient(bq, { + createReadSession: async req => { + sessionReqs.push(req); + return [ + { + arrowSchema: SAMPLE_READ_ROWS_RESPONSE.arrowSchema, + streams: [ + { + name: 'projects/test-project/locations/US/sessions/sess-2/streams/stream-0', + }, + ], + }, + ]; + }, + readRows: () => + fakeReadRowsStream([ + {arrowRecordBatch: SAMPLE_READ_ROWS_RESPONSE.arrowRecordBatch}, + ]), + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream({ + query: 'SELECT id, name FROM tbl', + destination, + writeDisposition: 'WRITE_APPEND', + }) + .toArray(); + + // Verify WRITE_APPEND is sent to jobs.insert instead of being rejected. + const queryConfig = getCapturedJobOpts()?.configuration?.query; + assert.strictEqual(queryConfig?.writeDisposition, 'WRITE_APPEND'); + + // Verify the stream reads the destination table, which holds the + // existing rows plus the appended ones. + assert.deepStrictEqual( + sessionReqs.map(req => req.readSession.table), + ['projects/test-project/datasets/dest_ds/tables/append_tbl'], + ); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + it('should emit a 0-row RecordBatch preserving session.arrowSchema when createReadSession returns no streams', async () => { const {destination} = stubDestinationJob('empty_tbl', 'dest-empty-1');