feat(bigquery): support explicit destination tables in createQueryArrowStream

handwritten/bigquery/src/arrow.ts

@@ -367,6 +367,37 @@ async function streamDefaultJobStream(
   pushBatches(stream, decodeArrowPayload(arrow, schema));
 }
 
+/**
+ * Creates an Arrow read session with a single stream on `table`, billed to
+ * `parentProjectId`.
+ *
+ * @private
+ */
+async function createArrowReadSession(
+  ctx: ArrowStreamContext,
+  parentProjectId: string | undefined,
+  table: bigquery.ITableReference,
+): Promise<IReadSession | undefined> {
+  if (!parentProjectId) {
+    throw new Error(
+      'Cannot create an Arrow read session without a project ID.',
+    );
+  }
+  const [session] = await getStorageReadClient(
+    ctx.bq,
+    ctx.storage,
+  ).createReadSession({
+    parent: `projects/${parentProjectId}`,
+    readSession: {
+      table: `projects/${table.projectId}/datasets/${table.datasetId}/tables/${table.tableId}`,
+      dataFormat: 'ARROW',
+    },
+    // One stream keeps rows in the query's order (e.g., ORDER BY).
+    maxStreamCount: 1,
+  });
+  return session ?? undefined;
+}
+
 /**
  * Reads the Arrow schema of a finished query job's results from its
  * destination table. The default job stream sends no `ReadRowsResponse`, and
@@ -380,7 +411,6 @@ async function readDestinationSchema(
   ctx: ArrowStreamContext,
   job: Job,
 ): Promise<Uint8Array | string | undefined> {
-  const {bq, storage} = ctx;
   // A script's destination table is only set once the job finishes, so read
   // the job's current metadata instead of the response that created it.
   const [metadata] = await job.getMetadata();
@@ -389,18 +419,15 @@ async function readDestinationSchema(
   if (!destRef || ctx.stream.destroyed) {
     return undefined;
   }
-  const [session] = await getStorageReadClient(bq, storage).createReadSession({
-    parent: `projects/${resolveProjectId(
-      bq,
+  const session = await createArrowReadSession(
+    ctx,
+    resolveProjectId(
+      ctx.bq,
       metadata?.jobReference?.projectId,
       destRef.projectId,
-    )}`,
-    readSession: {
-      table: `projects/${destRef.projectId}/datasets/${destRef.datasetId}/tables/${destRef.tableId}`,
-      dataFormat: 'ARROW',
-    },
-    maxStreamCount: 1,
-  });
+    ),
+    destRef,
+  );
   return session?.arrowSchema?.serializedSchema || undefined;
 }
 
@@ -518,6 +545,45 @@ async function streamReadRows(
   });
 }
 
+/**
+ * Creates a Storage Read API session for an explicit destination table and
+ * streams Arrow `RecordBatch`es from its single stream.
+ *
+ * @private
+ */
+async function streamDestinationTable(
+  ctx: ArrowStreamContext,
+  jobRef: bigquery.IJobReference | undefined,
+  table: bigquery.ITableReference,
+): Promise<void> {
+  const {arrow, stream} = ctx;
+  // The session is billed to the job's project; the table can be in another.
+  const session = await createArrowReadSession(
+    ctx,
+    resolveProjectId(ctx.bq, jobRef?.projectId),
+    table,
+  );
+  const sessionSchema = session?.arrowSchema?.serializedSchema || undefined;
+  const streams = session?.streams ?? [];
+  const readStream = streams[0]?.name;
+  if (streams.length === 0 || !readStream) {
+    // An empty destination table returns 0 streams; emit a single 0-row
+    // RecordBatch carrying the session's Arrow schema.
+    pushBatches(stream, decodeArrowPayload(arrow, sessionSchema));
+    return;
+  }
+
+  const {emittedBatch, serializedSchema} = await streamReadRows(ctx, {
+    readStream,
+    offset: 0,
+    initialSchema: sessionSchema,
+  });
+  // A stream that ends without rows still yields a 0-row RecordBatch.
+  if (!emittedBatch) {
+    pushBatches(stream, decodeArrowPayload(arrow, serializedSchema));
+  }
+}
+
 /**
  * Waits for an unfinished BigQuery job to complete using `job.getQueryResults`
  * without fetching any `STRUCT_ENCODING` rows (`maxResults: 0`). Returns early
@@ -550,10 +616,11 @@ async function waitForJobCompletion(
 }
 
 /**
- * Waits for `job` to finish, then streams all of its rows from the default job
- * stream. An unfinished jobs.query and jobs.insert return no rows or schema,
- * so this reads every row from offset 0 and passes the job to look up the
- * schema of an empty result.
+ * Waits for `job` to finish, then streams all of its rows: from
+ * `destinationTable` through a read session when the caller set one, or from
+ * the default job stream otherwise. An unfinished jobs.query and jobs.insert
+ * return no rows or schema, so the default job stream is read from offset 0,
+ * passing the job to look up the schema of an empty result.
  *
  * @private
  */
@@ -562,11 +629,16 @@ async function streamJobWhenComplete(
   job: Job,
   jobRef: bigquery.IJobReference | undefined,
   timeoutMs?: number,
+  destinationTable?: bigquery.ITableReference,
 ): Promise<void> {
   await waitForJobCompletion(job, ctx.stream, timeoutMs);
   if (ctx.stream.destroyed) {
     return;
   }
+  if (destinationTable) {
+    await streamDestinationTable(ctx, jobRef, destinationTable);
+    return;
+  }
   await streamDefaultJobStream(ctx, jobRef, 0, undefined, job);
 }
 
@@ -606,7 +678,18 @@ async function streamQueryJob(
     return;
   }
 
-  await streamJobWhenComplete(ctx, job, resp?.jobReference, timeoutMs);
+  // BigQuery echoes the caller's destination table, with its resolved
+  // project, in the job's configuration.
+  const destinationTable = queryObj.destination
+    ? (resp?.configuration?.query?.destinationTable ?? undefined)
+    : undefined;
+  await streamJobWhenComplete(
+    ctx,
+    job,
+    resp?.jobReference,
+    timeoutMs,
+    destinationTable,
+  );
 }
 
 /**
@@ -662,13 +745,6 @@ export function createQueryArrowStream_(
       );
     }
 
-    if (queryObj.destination) {
-      // TODO: Support explicit destination tables via CreateReadSession.
-      throw new Error(
-        'Explicit destination tables are not yet supported in createQueryArrowStream.',
-      );
-    }
-
     const ctx: ArrowStreamContext = {bq, arrow, storage, stream};
     readCtx = ctx;
 
@@ -678,8 +754,8 @@ export function createQueryArrowStream_(
     });
     if (!queryReq) {
       // buildQueryRequest_ returns undefined when the query sets an option
-      // that only jobs.insert accepts (e.g., jobId, dryRun, or jobTimeoutMs),
-      // or when FAST_QUERY_PATH=DISABLED.
+      // that only jobs.insert accepts (e.g., destination, jobId, dryRun, or
+      // jobTimeoutMs), or when FAST_QUERY_PATH=DISABLED.
       await streamQueryJob(ctx, queryObj);
       return;
     }

handwritten/bigquery/src/bigquery.ts

@@ -375,7 +375,7 @@
   private _universeDomain: string;
   private _defaultJobCreationMode: JobCreationMode;
 
   createQueryStream(options?: Query | string): ResourceStream<RowMetadata> {
     // placeholder body, overwritten in constructor
     return new ResourceStream<RowMetadata>({}, () => {});
   }
@@ -401,6 +401,9 @@
    * ends, call `stream.destroy()` to cancel the read; breaking out of
    * `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.
+   *
    * @param {string|ArrowQuery} [options] A string SQL query or configuration
    *     object. See {@link BigQuery#query} for a complete list of options.
    * @param {number} [options.timeoutMs] How long to wait for the query to
@@ -441,12 +444,12 @@
     return createQueryArrowStream_(this, options);
   }
 
   getDatasetsStream(options?: GetDatasetsOptions): ResourceStream<Dataset> {
     // placeholder body, overwritten in constructor
     return new ResourceStream<Dataset>({}, () => {});
   }
 
   getJobsStream(options?: GetJobsOptions): ResourceStream<Job> {
     // placeholder body, overwritten in constructor
     return new ResourceStream<Job>({}, () => {});
   }
@@ -1665,7 +1668,7 @@
     const parameterMode = isArray(params) ? 'positional' : 'named';
     const queryParameters: bigquery.IQueryParameter[] = [];
     if (parameterMode === 'named') {
       const namedParams = params as {[param: string]: any};
       for (const namedParameter of Object.getOwnPropertyNames(namedParams)) {
         const value = namedParams[namedParameter];
         let queryParameter;
@@ -2312,7 +2315,7 @@
 
       options = extend({job}, queryOpts, options);
       if (res && res.jobComplete) {
         let rows: any = [];
         if (res.schema && res.rows) {
           try {
             /*

handwritten/bigquery/system-test/arrow.ts

@@ -669,4 +669,53 @@ describe('Arrow Query Results', () => {
     assert.ok(arrow.DataType.isInt(batches[0].schema.fields[0].type));
     assert.ok(arrow.DataType.isUtf8(batches[0].schema.fields[1].type));
   });
+
+  it('should stream Arrow results from an explicit destination table via CreateReadSession with default writeDisposition and WRITE_TRUNCATE', async () => {
+    const destTable = dataset.table('dest_results');
+
+    // 1. Default writeDisposition (writes to a new destination table).
+    const initialBatches: arrow.RecordBatch[] = await bigquery
+      .createQueryArrowStream({
+        query:
+          'SELECT x AS id, CONCAT("row_", CAST(x AS STRING)) AS label FROM UNNEST(GENERATE_ARRAY(1, 100)) AS x ORDER BY x',
+        destination: destTable,
+      })
+      .toArray();
+
+    const initialTable = new arrow.Table(initialBatches);
+    assert.strictEqual(initialTable.numRows, 100);
+    assert.strictEqual(initialTable.getChild('id')!.get(0), 1n);
+    assert.strictEqual(initialTable.getChild('id')!.get(99), 100n);
+
+    // 2. WRITE_TRUNCATE overwrites the destination table and reads only the new rows.
+    const truncatedBatches: arrow.RecordBatch[] = await bigquery
+      .createQueryArrowStream({
+        query:
+          'SELECT x AS id, CONCAT("new_", CAST(x AS STRING)) AS label FROM UNNEST(GENERATE_ARRAY(1, 10)) AS x ORDER BY x',
+        destination: destTable,
+        writeDisposition: 'WRITE_TRUNCATE',
+      })
+      .toArray();
+
+    const truncatedTable = new arrow.Table(truncatedBatches);
+    assert.strictEqual(truncatedTable.numRows, 10);
+    assert.strictEqual(truncatedTable.getChild('id')!.get(0), 1n);
+    assert.strictEqual(truncatedTable.getChild('label')!.get(0), 'new_1');
+
+    // 3. 0-row query into an explicit destination table preserves the schema.
+    const emptyBatches: arrow.RecordBatch[] = await bigquery
+      .createQueryArrowStream({
+        query: 'SELECT 1 AS id, "empty" AS label FROM UNNEST([1]) WHERE FALSE',
+        destination: destTable,
+        writeDisposition: 'WRITE_TRUNCATE',
+      })
+      .toArray();
+
+    assert.strictEqual(emptyBatches.length, 1);
+    assert.strictEqual(emptyBatches[0].numRows, 0);
+    assert.deepStrictEqual(
+      emptyBatches[0].schema.fields.map(f => f.name),
+      ['id', 'label'],
+    );
+  });
 });