chore: Add comments and refactor constants

Author: lqiu96

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageReadClientTest.java

@@ -81,7 +81,6 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import com.google.protobuf.Int64Value;
-import com.google.protobuf.Timestamp;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
@@ -114,7 +113,6 @@
 import org.apache.avro.LogicalTypes;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericData;
-import org.apache.avro.generic.GenericRecordBuilder;
 import org.apache.avro.util.Utf8;
 import org.json.JSONArray;
 import org.json.JSONObject;
@@ -530,6 +528,7 @@ public static void beforeClass()
 
   private static void setupTimestampTable()
       throws DescriptorValidationException, IOException, InterruptedException {
+    // Schema to create a BQ table
     com.google.cloud.bigquery.Schema timestampSchema =
         com.google.cloud.bigquery.Schema.of(
             Field.newBuilder(TIMESTAMP_COLUMN_NAME, StandardSQLTypeName.TIMESTAMP)
@@ -540,6 +539,13 @@ private static void setupTimestampTable()
                 .setMode(Mode.NULLABLE)
                 .build());
 
+    // Create BQ table with timestamps
+    TableId tableId = TableId.of(DATASET, BQSTORAGE_TIMESTAMP_READ_TABLE);
+    bigquery.create(TableInfo.of(tableId, StandardTableDefinition.of(timestampSchema)));
+
+    TableName parentTable = TableName.of(projectName, DATASET, BQSTORAGE_TIMESTAMP_READ_TABLE);
+
+    // Define the BQStorage schema to write to
     TableSchema timestampTableSchema =
         TableSchema.newBuilder()
             .addFields(
@@ -557,23 +563,8 @@ private static void setupTimestampTable()
                     .build())
             .build();
 
-    // Create table with Range fields.
-    TableId tableId = TableId.of(DATASET, BQSTORAGE_TIMESTAMP_READ_TABLE);
-    bigquery.create(TableInfo.of(tableId, StandardTableDefinition.of(timestampSchema)));
-
-    TableName parentTable = TableName.of(projectName, DATASET, BQSTORAGE_TIMESTAMP_READ_TABLE);
-    RetrySettings retrySettings =
-        RetrySettings.newBuilder()
-            .setInitialRetryDelayDuration(Duration.ofMillis(500))
-            .setRetryDelayMultiplier(1.1)
-            .setMaxAttempts(5)
-            .setMaxRetryDelayDuration(Duration.ofSeconds(10))
-            .build();
-
     try (JsonStreamWriter writer =
-        JsonStreamWriter.newBuilder(parentTable.toString(), timestampTableSchema)
-            .setRetrySettings(retrySettings)
-            .build()) {
+        JsonStreamWriter.newBuilder(parentTable.toString(), timestampTableSchema).build()) {
       JSONArray data = new JSONArray();
       for (Object[] timestampData : Helper.INPUT_TIMESTAMPS) {
         JSONObject row = new JSONObject();
@@ -874,6 +865,7 @@ public void testRangeTypeWrite()
     }
   }
 
+  // Tests that inputs for micros and picos can be read properly via Arrow
   @Test
   public void timestamp_readArrow() throws IOException {
     String table =
@@ -925,13 +917,15 @@ public void timestamp_readArrow() throws IOException {
         Preconditions.checkState(response.hasArrowRecordBatch());
         reader.processRows(
             response.getArrowRecordBatch(),
-            new SimpleRowReaderArrow.ArrowTimestampBatchConsumer(Helper.INPUT_TIMESTAMPS));
+            new SimpleRowReaderArrow.ArrowTimestampBatchConsumer(
+                Helper.EXPECTED_TIMESTAMPS_HIGHER_PRECISION_ISO_OUTPUT));
         rowCount += response.getRowCount();
       }
-      assertEquals(Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION.length, rowCount);
+      assertEquals(Helper.EXPECTED_TIMESTAMPS_HIGHER_PRECISION_ISO_OUTPUT.length, rowCount);
     }
   }
 
+  // Tests that inputs for micros and picos can be read properly via Avro
   @Test
   public void timestamp_readAvro() throws IOException {
     String table =
@@ -944,11 +938,12 @@ public void timestamp_readAvro() throws IOException {
             .map(x -> x.get(TIMESTAMP_HIGHER_PRECISION_COLUMN_NAME).toString())
             .collect(Collectors.toList());
     for (int i = 0; i < timestamps.size(); i++) {
-      assertEquals(Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION[i][0], timestamps.get(i));
+      assertEquals(Helper.EXPECTED_TIMESTAMPS_HIGHER_PRECISION_ISO_OUTPUT[i][0], timestamps.get(i));
     }
     for (int i = 0; i < timestampHigherPrecision.size(); i++) {
       assertEquals(
-          Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION[i][1], timestampHigherPrecision.get(i));
+          Helper.EXPECTED_TIMESTAMPS_HIGHER_PRECISION_ISO_OUTPUT[i][1],
+          timestampHigherPrecision.get(i));
     }
   }
 
@@ -1887,90 +1882,6 @@ private long readStreamToOffset(ReadStream readStream, long rowOffset) {
     return rowCount;
   }
 
-  /**
-   * Reads all the rows from the specified table.
-   *
-   * <p>For every row, the consumer is called for processing.
-   *
-   * @param table
-   * @param snapshotInMillis Optional. If specified, all rows up to timestamp will be returned.
-   * @param filter Optional. If specified, it will be used to restrict returned data.
-   * @param consumer that receives all Avro rows.
-   * @throws IOException
-   */
-  private void processRowsAtSnapshot(
-      String table, Long snapshotInMillis, String filter, AvroRowConsumer consumer)
-      throws IOException {
-    Preconditions.checkNotNull(table);
-    Preconditions.checkNotNull(consumer);
-
-    CreateReadSessionRequest.Builder createSessionRequestBuilder =
-        CreateReadSessionRequest.newBuilder()
-            .setParent(parentProjectId)
-            .setMaxStreamCount(1)
-            .setReadSession(
-                ReadSession.newBuilder().setTable(table).setDataFormat(DataFormat.AVRO).build());
-
-    if (snapshotInMillis != null) {
-      Timestamp snapshotTimestamp =
-          Timestamp.newBuilder()
-              .setSeconds(snapshotInMillis / 1_000)
-              .setNanos((int) ((snapshotInMillis % 1000) * 1000000))
-              .build();
-      createSessionRequestBuilder
-          .getReadSessionBuilder()
-          .setTableModifiers(
-              ReadSession.TableModifiers.newBuilder().setSnapshotTime(snapshotTimestamp).build());
-    }
-
-    if (filter != null && !filter.isEmpty()) {
-      createSessionRequestBuilder
-          .getReadSessionBuilder()
-          .setReadOptions(TableReadOptions.newBuilder().setRowRestriction(filter).build());
-    }
-
-    ReadSession session = readClient.createReadSession(createSessionRequestBuilder.build());
-    assertEquals(
-        String.format(
-            "Did not receive expected number of streams for table '%s' CreateReadSession"
-                + " response:%n%s",
-            table, session.toString()),
-        1,
-        session.getStreamsCount());
-
-    ReadRowsRequest readRowsRequest =
-        ReadRowsRequest.newBuilder().setReadStream(session.getStreams(0).getName()).build();
-
-    SimpleRowReaderAvro reader =
-        new SimpleRowReaderAvro(new Schema.Parser().parse(session.getAvroSchema().getSchema()));
-
-    ServerStream<ReadRowsResponse> stream = readClient.readRowsCallable().call(readRowsRequest);
-    for (ReadRowsResponse response : stream) {
-      reader.processRows(response.getAvroRows(), consumer);
-    }
-  }
-
-  /**
-   * Reads all the rows from the specified table and returns a list as generic Avro records.
-   *
-   * @param table
-   * @param filter Optional. If specified, it will be used to restrict returned data.
-   * @return
-   */
-  List<GenericData.Record> readAllRows(String table, String filter) throws IOException {
-    final List<GenericData.Record> rows = new ArrayList<>();
-    processRowsAtSnapshot(
-        /* table= */ table,
-        /* snapshotInMillis= */ null,
-        /* filter= */ filter,
-        (AvroRowConsumer)
-            record -> {
-              // clone the record since that reference will be reused by the reader.
-              rows.add(new GenericRecordBuilder(record).build());
-            });
-    return rows;
-  }
-
   /**
    * Runs a query job with WRITE_APPEND disposition to the destination table and returns the
    * successfully completed job.

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryStorageWriteClientTest.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.bigquery.storage.v1.it;
 
+import static com.google.cloud.bigquery.storage.v1.it.util.Helper.EXPECTED_TIMESTAMPS_HIGHER_PRECISION_ISO_OUTPUT;
 import static com.google.cloud.bigquery.storage.v1.it.util.Helper.TIMESTAMP_COLUMN_NAME;
 import static com.google.cloud.bigquery.storage.v1.it.util.Helper.TIMESTAMP_HIGHER_PRECISION_COLUMN_NAME;
 import static com.google.common.truth.Truth.assertThat;
@@ -131,8 +132,9 @@ public class ITBigQueryStorageWriteClientTest {
         {165846896123456L /* 1975-04-04T12:34:56.123456Z */, 165846896123456789L}
       };
 
-  // Arrow's higher precision column is padded with extra 0's.
-  public static final Object[][] EXPECTED_ARROW_WRITE_TIMESTAMPS =
+  // Arrow's higher precision column is padded with extra 0's if configured to return
+  // ISO as output for any picosecond enabled column.
+  public static final Object[][] EXPECTED_ARROW_WRITE_TIMESTAMPS_ISO_OUTPUT =
       new Object[][] {
         {1735734896123456L /* 2025-01-01T12:34:56.123456Z */, "2025-01-01T12:34:56.123456789000Z"},
         {1580646896123456L /* 2020-02-02T12:34:56.123456Z */, "2020-02-02T12:34:56.123456789000Z"},
@@ -2296,25 +2298,18 @@ public void testLargeRequest() throws IOException, InterruptedException, Executi
     }
   }
 
+  // Tests that inputs for micro and picos are able to use Arrow to write
+  // to BQ
   @Test
   public void timestamp_arrowWrite() throws IOException {
-    com.google.cloud.bigquery.Schema bigqueryTableSchema =
-        com.google.cloud.bigquery.Schema.of(
-            Field.newBuilder(TIMESTAMP_COLUMN_NAME, StandardSQLTypeName.TIMESTAMP)
-                .setMode(Mode.NULLABLE)
-                .build(),
-            Field.newBuilder(TIMESTAMP_HIGHER_PRECISION_COLUMN_NAME, StandardSQLTypeName.TIMESTAMP)
-                .setMode(Mode.NULLABLE)
-                .setTimestampPrecision(12L)
-                .build());
-
     String tableName = "bqstorage_timestamp_write_arrow";
-    TableId testTableId = TableId.of(DATASET, tableName);
-    bigquery.create(
-        TableInfo.of(
-            testTableId,
-            StandardTableDefinition.newBuilder().setSchema(bigqueryTableSchema).build()));
+    // Opt to create a new table to write to instead of re-using table to prevent
+    // the test from failing due to any issues with deleting data after test.
+    // Increases the test time duration, but would be more resilient to transient
+    // failures
+    createTimestampTable(tableName);
 
+    // Define the fields as Arrow types that are compatible with BQ Schema types
     List<org.apache.arrow.vector.types.pojo.Field> fields =
         ImmutableList.of(
             new org.apache.arrow.vector.types.pojo.Field(
@@ -2332,7 +2327,7 @@ public void timestamp_arrowWrite() throws IOException {
     org.apache.arrow.vector.types.pojo.Schema arrowSchema =
         new org.apache.arrow.vector.types.pojo.Schema(fields, null);
 
-    int numRows = Helper.INPUT_TIMESTAMPS.length;
+    int numRows = INPUT_ARROW_WRITE_TIMESTAMPS.length;
     TableName parent = TableName.of(ServiceOptions.getDefaultProjectId(), DATASET, tableName);
     try (StreamWriter streamWriter =
         StreamWriter.newBuilder(parent.toString() + "/_default")
@@ -2365,44 +2360,23 @@ public void timestamp_arrowWrite() throws IOException {
             future, new Helper.AppendCompleteCallback(), MoreExecutors.directExecutor());
       }
     }
-    String table =
-        BigQueryResource.formatTableResource(
-            ServiceOptions.getDefaultProjectId(), DATASET, tableName);
-    List<GenericData.Record> rows = Helper.readAllRows(readClient, parentProjectId, table, null);
-    List<Long> timestamps =
-        rows.stream().map(x -> (Long) x.get(TIMESTAMP_COLUMN_NAME)).collect(Collectors.toList());
-    List<String> timestampHigherPrecision =
-        rows.stream()
-            .map(x -> x.get(TIMESTAMP_HIGHER_PRECISION_COLUMN_NAME).toString())
-            .collect(Collectors.toList());
-    assertEquals(timestamps.size(), Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION.length);
-    assertEquals(
-        timestampHigherPrecision.size(), Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION.length);
-    for (int i = 0; i < timestampHigherPrecision.size(); i++) {
-      assertEquals(timestamps.get(i), Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION[i][0]);
-      assertEquals(timestampHigherPrecision.get(i), EXPECTED_ARROW_WRITE_TIMESTAMPS[i][1]);
-    }
+    assertTimestamps(tableName, EXPECTED_ARROW_WRITE_TIMESTAMPS_ISO_OUTPUT);
   }
 
+  // Tests that inputs for micro and picos are able to converted to protobuf
+  // and written to BQ
   @Test
   public void timestamp_protobufWrite()
       throws IOException, DescriptorValidationException, InterruptedException {
-    com.google.cloud.bigquery.Schema bqTableSchema =
-        com.google.cloud.bigquery.Schema.of(
-            Field.newBuilder(TIMESTAMP_COLUMN_NAME, StandardSQLTypeName.TIMESTAMP)
-                .setMode(Mode.NULLABLE)
-                .build(),
-            Field.newBuilder(TIMESTAMP_HIGHER_PRECISION_COLUMN_NAME, StandardSQLTypeName.TIMESTAMP)
-                .setMode(Mode.NULLABLE)
-                .setTimestampPrecision(12L)
-                .build());
-
     String tableName = "bqstorage_timestamp_write_protobuf";
-    TableId testTableId = TableId.of(DATASET, tableName);
-    bigquery.create(
-        TableInfo.of(
-            testTableId, StandardTableDefinition.newBuilder().setSchema(bqTableSchema).build()));
-
+    // Opt to create a new table to write to instead of re-using table to prevent
+    // the test from failing due to any issues with deleting data after test.
+    // Increases the test time duration, but would be more resilient to transient
+    // failures
+    createTimestampTable(tableName);
+
+    // Define the table schema so that the automatic converter is able to
+    // determine how to convert from Json -> Protobuf
     TableFieldSchema testTimestamp =
         TableFieldSchema.newBuilder()
             .setName(TIMESTAMP_COLUMN_NAME)
@@ -2436,24 +2410,48 @@ public void timestamp_protobufWrite()
       ApiFutures.addCallback(
           future, new Helper.AppendCompleteCallback(), MoreExecutors.directExecutor());
     }
+    assertTimestamps(tableName, EXPECTED_TIMESTAMPS_HIGHER_PRECISION_ISO_OUTPUT);
+  }
 
+  private void createTimestampTable(String tableName) {
+    Schema bqTableSchema =
+        Schema.of(
+            Field.newBuilder(TIMESTAMP_COLUMN_NAME, StandardSQLTypeName.TIMESTAMP)
+                .setMode(Mode.NULLABLE)
+                .build(),
+            Field.newBuilder(TIMESTAMP_HIGHER_PRECISION_COLUMN_NAME, StandardSQLTypeName.TIMESTAMP)
+                .setMode(Mode.NULLABLE)
+                .setTimestampPrecision(12L)
+                .build());
+
+    TableId testTableId = TableId.of(DATASET, tableName);
+    bigquery.create(
+        TableInfo.of(
+            testTableId, StandardTableDefinition.newBuilder().setSchema(bqTableSchema).build()));
+  }
+
+  private void assertTimestamps(String tableName, Object[][] expected) throws IOException {
     String table =
         BigQueryResource.formatTableResource(
             ServiceOptions.getDefaultProjectId(), DATASET, tableName);
+
+    // Read all the data as Avro GenericRecords
     List<GenericData.Record> rows = Helper.readAllRows(readClient, parentProjectId, table, null);
+
+    // Each timestamp response is expected to contain two fields:
+    // 1. Micros from timestamp as a Long and 2. ISO8601 instant with picos precision
     List<Long> timestamps =
         rows.stream().map(x -> (Long) x.get(TIMESTAMP_COLUMN_NAME)).collect(Collectors.toList());
     List<String> timestampHigherPrecision =
         rows.stream()
             .map(x -> x.get(TIMESTAMP_HIGHER_PRECISION_COLUMN_NAME).toString())
             .collect(Collectors.toList());
-    assertEquals(timestamps.size(), Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION.length);
-    assertEquals(
-        timestampHigherPrecision.size(), Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION.length);
-    for (int i = 0; i < timestamps.size(); i++) {
-      assertEquals(timestamps.get(i), Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION[i][0]);
-      assertEquals(
-          timestampHigherPrecision.get(i), Helper.EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION[i][1]);
+
+    assertEquals(timestamps.size(), expected.length);
+    assertEquals(timestampHigherPrecision.size(), expected.length);
+    for (int i = 0; i < timestampHigherPrecision.size(); i++) {
+      assertEquals(timestamps.get(i), expected[i][0]);
+      assertEquals(timestampHigherPrecision.get(i), expected[i][1]);
     }
   }
 }

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/util/Helper.java

@@ -56,9 +56,9 @@ public class Helper {
         {165846896123456L /* 1975-04-04T12:34:56.123456Z */, "1975-04-04T12:34:56.123456789123Z"}
       };
 
-  // Expected response for timestamps from the input. The output is configured to return ISO8601
-  // format for any picosecond enabled column.
-  public static final Object[][] EXPECTED_TIMESTAMPS_ISO_HIGHER_PRECISION =
+  // Expected response for timestamps from the input. If enabled with ISO as output, it will
+  // ISO8601 format for any picosecond enabled column.
+  public static final Object[][] EXPECTED_TIMESTAMPS_HIGHER_PRECISION_ISO_OUTPUT =
       new Object[][] {
         {1735734896123456L /* 2025-01-01T12:34:56.123456Z */, "2025-01-01T12:34:56.123456789123Z"},
         {1580646896123456L /* 2020-02-02T12:34:56.123456Z */, "2020-02-02T12:34:56.123456789123Z"},