feat: Row readers using storage API acceleration can leverage full precision avro values (#7525)

## Description

Copied from https://togithub.com/googleapis/nodejs-bigquery-storage/pull/658

This PR allows users to fetch high precision timestamps for read calls with avro readers. Previous to this PR there was no way users could fetch such timestamps.

## Impact

With a custom Avro reader, now users can fetch and consume data from the server that preserves high precision timestamps on the backend while providing information about the avro reader in the request.

## Testing

A system test is added that proves the timestamp fetched from the server is high precision.

## Additional Information

The source code changes are pretty much exactly in the same place as those from https://togithub.com/googleapis/nodejs-bigquery-storage/pull/656. The main difference is data coming back from the server for arrow readers is very different from data coming back from avro readers. That means instead of using arrow reader transforms from the source code we need to leverage the avsc library to parse server results in the same way that the quickstart guide demonstrates. This allows the test to produce a high precision timestamp.

Author: danieljbruce

handwritten/bigquery-storage/package.json

@@ -50,6 +50,7 @@
     "@types/node": "^22.13.14",
     "@types/sinon": "^17.0.4",
     "@types/uuid": "^10.0.0",
+    "avsc": "^5.7.9",
     "c8": "^10.1.3",
     "gapic-tools": "^1.0.1",
     "gts": "^6.0.2",

handwritten/bigquery-storage/src/reader/read_client.ts

@@ -25,7 +25,7 @@ import {DataFormat} from './data_format';
 type CreateReadSessionRequest =
   protos.google.cloud.bigquery.storage.v1.ICreateReadSessionRequest;
 type ReadSession = protos.google.cloud.bigquery.storage.v1.IReadSession;
-type ArrowSerializationOptions = {
+type AvroArrowSerializationOptions = {
   picosTimestampPrecision: protos.google.cloud.bigquery.storage.v1.ArrowSerializationOptions.PicosTimestampPrecision;
 };
 
@@ -131,7 +131,8 @@ export class ReadClient {
     table: string;
     dataFormat: DataFormat;
     selectedFields?: string[];
-    arrowSerializationOptions?: ArrowSerializationOptions;
+    arrowSerializationOptions?: AvroArrowSerializationOptions;
+    avroSerializationOptions?: AvroArrowSerializationOptions;
   }): Promise<ReadSession> {
     await this.initialize();
     const {table, parent, dataFormat, selectedFields} = request;
@@ -154,6 +155,11 @@ export class ReadClient {
         arrowSerializationOptions: request.arrowSerializationOptions,
       });
     }
+    if (request.avroSerializationOptions) {
+      Object.assign(createReq.readSession.readOptions, {
+        avroSerializationOptions: request.avroSerializationOptions,
+      });
+    }
     const [response] = await this._client.createReadSession(createReq);
     if (typeof [response] === undefined) {
       throw new gax.GoogleError(`${response}`);

handwritten/bigquery-storage/src/reader/read_session.ts

@@ -25,7 +25,7 @@ type ReadRowsResponse =
   protos.google.cloud.bigquery.storage.v1.IReadRowsResponse;
 type ReadSessionInfo = protos.google.cloud.bigquery.storage.v1.IReadSession;
 const ReadSessionInfo = protos.google.cloud.bigquery.storage.v1.ReadSession;
-type ArrowSerializationOptions = {
+type AvroArrowSerializationOptions = {
   picosTimestampPrecision: protos.google.cloud.bigquery.storage.v1.ArrowSerializationOptions.PicosTimestampPrecision;
 };
 
@@ -39,9 +39,13 @@ export type GetStreamOptions = {
    */
   selectedFields?: string;
   /**
-   * Option to opt into higher precision timestamps.
+   * Option to opt into higher precision timestamps for arrow readers.
    */
-  arrowSerializationOptions?: ArrowSerializationOptions;
+  arrowSerializationOptions?: AvroArrowSerializationOptions;
+  /**
+   * Option to opt into higher precision timestamps for avro readers.
+   */
+  avroSerializationOptions?: AvroArrowSerializationOptions;
 };
 
 /**
@@ -93,6 +97,7 @@ export class ReadSession {
       dataFormat: this._format,
       selectedFields: options?.selectedFields?.split(','),
       arrowSerializationOptions: options?.arrowSerializationOptions,
+      avroSerializationOptions: options?.avroSerializationOptions,
     });
     this.trace(
       'session created',

handwritten/bigquery-storage/system-test/reader_client_test.ts

@@ -29,6 +29,8 @@ import {RecordBatch, Table, tableFromIPC} from 'apache-arrow';
 import {ArrowRecordBatchTableRowTransform} from '../src/reader/arrow_transform';
 import {ResourceStream} from '@google-cloud/paginator';
 import {ArrowTableReader} from '../src/reader';
+import {Transform, TransformCallback} from 'stream';
+type ReadSession = protos.google.cloud.bigquery.storage.v1.IReadSession;
 
 type ReadRowsResponse =
   protos.google.cloud.bigquery.storage.v1.IReadRowsResponse;
@@ -230,6 +232,132 @@ describe('reader.ReaderClient', () => {
     });
   });
 
+  describe('AvroReader', () => {
+    it('should read high precision timestamps from an avro stream', async () => {
+      const avro = require('avsc');
+      class AvroRawTransform extends Transform {
+        private session: ReadSession;
+
+        constructor(session: ReadSession) {
+          super({
+            objectMode: true,
+          });
+          this.session = session;
+        }
+
+        _transform(
+          serializedRecordBatch: any,
+          _: BufferEncoding,
+          callback: TransformCallback,
+        ): void {
+          const session = this.session;
+          const schema = JSON.parse(session?.avroSchema?.schema as string);
+          const avroType = avro.Type.forSchema(schema);
+          if (
+            !(
+              serializedRecordBatch.avroRows &&
+              serializedRecordBatch.avroRows.serializedBinaryRows
+            )
+          ) {
+            callback(null);
+            return;
+          }
+          const decodedData = avroType.decode(
+            serializedRecordBatch.avroRows.serializedBinaryRows,
+            0,
+          );
+          callback(null, decodedData.value);
+        }
+      }
+
+      const picosTableId = generateUuid();
+      const picosSchema: any = {
+        fields: [
+          {
+            name: 'customer_name',
+            type: 'STRING',
+            mode: 'REQUIRED',
+          },
+          {
+            name: 'row_num',
+            type: 'INTEGER',
+            mode: 'REQUIRED',
+          },
+          {
+            name: 'created_at',
+            type: 'TIMESTAMP',
+            mode: 'NULLABLE',
+            timestampPrecision: 12,
+          },
+        ],
+      };
+      const expectedTsValue = '2024-04-05T15:45:58.981123456789Z';
+      await bigquery
+        .dataset(datasetId)
+        .createTable(picosTableId, {schema: picosSchema});
+      await bigquery
+        .dataset(datasetId)
+        .table(picosTableId)
+        .insert([
+          {
+            customer_name: 'my-name',
+            row_num: 1,
+            created_at: expectedTsValue,
+          },
+        ]);
+
+      bqReadClient.initialize().catch(err => {
+        throw err;
+      });
+      const client = new ReadClient();
+      client.setClient(bqReadClient);
+
+      try {
+        const session = await client.createReadSession({
+          parent: `projects/${projectId}`,
+          table: `projects/${projectId}/datasets/${datasetId}/tables/${picosTableId}`,
+          dataFormat: AvroFormat,
+          avroSerializationOptions: {
+            picosTimestampPrecision:
+              protos.google.cloud.bigquery.storage.v1.ArrowSerializationOptions
+                .PicosTimestampPrecision.TIMESTAMP_PRECISION_PICOS,
+          },
+        });
+
+        assert.equal(session.dataFormat, AvroFormat);
+        assert.notEqual(session.streams, null);
+        assert.notEqual(session.streams?.length, 0);
+
+        const readStream = session.streams![0];
+        const connection = await client.createReadStream({
+          session,
+          streamName: readStream.name!,
+        });
+
+        const myStream = connection
+          .getRowsStream()
+          .pipe(new AvroRawTransform(session!));
+        const responses: ReadRowsResponse[] = [];
+        await new Promise((resolve, reject) => {
+          myStream.on('data', (data: ReadRowsResponse) => {
+            responses.push(data);
+          });
+          myStream.on('error', reject);
+          myStream.on('end', () => {
+            resolve(null);
+          });
+        });
+
+        assert.equal(responses.length, 1);
+        assert.equal((responses[0] as any)['created_at'], expectedTsValue);
+
+        connection.close();
+        client.close();
+      } finally {
+        client.close();
+      }
+    });
+  });
   describe('ArrowTableReader', () => {
     it('should allow to read a table as an Arrow byte stream', async () => {
       bqReadClient.initialize().catch(err => {