Support picosecond tiemstamps when writing GenericRecord and Beam Rows. (#37294)

Co-authored-by: Claude <cvandermerwe@google.com>

Author: claudevdm

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProto.java

@@ -348,6 +348,9 @@ private static TableFieldSchema fieldDescriptorFromAvroField(org.apache.avro.Sch
             fieldDescriptorFromAvroField(
                 new Schema.Field(field.name(), elementType, field.doc(), field.defaultVal()));
         builder = builder.setType(elementFieldSchema.getType());
+        if (elementFieldSchema.hasTimestampPrecision()) {
+          builder.setTimestampPrecision(elementFieldSchema.getTimestampPrecision());
+        }
         builder.addAllFields(elementFieldSchema.getFieldsList());
         builder = builder.setMode(TableFieldSchema.Mode.REPEATED);
         break;

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProto.java

@@ -237,6 +237,9 @@ private static TableFieldSchema fieldDescriptorFromBeamField(Field field) {
         TableFieldSchema elementFieldSchema =
             fieldDescriptorFromBeamField(Field.of(field.getName(), elementType));
         builder = builder.setType(elementFieldSchema.getType());
+        if (elementFieldSchema.hasTimestampPrecision()) {
+          builder = builder.setTimestampPrecision(elementFieldSchema.getTimestampPrecision());
+        }
         builder.addAllFields(elementFieldSchema.getFieldsList());
         builder = builder.setMode(TableFieldSchema.Mode.REPEATED);
         break;

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProtoTest.java

@@ -348,7 +348,23 @@ private static Schema createTimestampNanosSchema() {
         .endRecord();
   }
 
+  private static Schema createRepeatedTimestampNanosSchema() {
+    Schema longSchema = Schema.create(Schema.Type.LONG);
+    longSchema.addProp("logicalType", "timestamp-nanos");
+
+    Schema arraySchema = Schema.createArray(longSchema);
+
+    return SchemaBuilder.record("RepeatedTimestampNanosRecord")
+        .fields()
+        .name("timestampNanosArray")
+        .type(arraySchema)
+        .noDefault()
+        .endRecord();
+  }
+
   private static final Schema TIMESTAMP_NANOS_SCHEMA = createTimestampNanosSchema();
+  private static final Schema REPEATED_TIMESTAMP_NANOS_SCHEMA =
+      createRepeatedTimestampNanosSchema();
 
   private static GenericRecord baseRecord;
   private static GenericRecord rawLogicalTypesRecord;
@@ -885,4 +901,22 @@ public void testProtoTableSchemaFromAvroSchemaTimestampNanos() {
     assertTrue(field.hasTimestampPrecision());
     assertEquals(12L, field.getTimestampPrecision().getValue());
   }
+
+  @Test
+  public void testProtoTableSchemaFromAvroSchemaRepeatedTimestampNanos() {
+    com.google.cloud.bigquery.storage.v1.TableSchema protoSchema =
+        AvroGenericRecordToStorageApiProto.protoTableSchemaFromAvroSchema(
+            REPEATED_TIMESTAMP_NANOS_SCHEMA);
+
+    assertEquals(1, protoSchema.getFieldsCount());
+    com.google.cloud.bigquery.storage.v1.TableFieldSchema field = protoSchema.getFields(0);
+
+    assertEquals("timestampnanosarray", field.getName());
+    assertEquals(
+        com.google.cloud.bigquery.storage.v1.TableFieldSchema.Type.TIMESTAMP, field.getType());
+    assertEquals(
+        com.google.cloud.bigquery.storage.v1.TableFieldSchema.Mode.REPEATED, field.getMode());
+
+    assertEquals(12L, field.getTimestampPrecision().getValue());
+  }
 }

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BeamRowToStorageApiProtoTest.java

@@ -69,6 +69,10 @@ public class BeamRowToStorageApiProtoTest {
       Schema.builder()
           .addField("timestampNanos", FieldType.logicalType(Timestamp.NANOS).withNullable(true))
           .build();
+  private static final Schema TIMESTAMP_NANOS_ARRAY_SCHEMA =
+      Schema.builder()
+          .addField("timestampNanosArray", FieldType.array(FieldType.logicalType(Timestamp.NANOS)))
+          .build();
   private static final EnumerationType TEST_ENUM =
       EnumerationType.create("ONE", "TWO", "RED", "BLUE");
   private static final Schema BASE_SCHEMA =
@@ -614,6 +618,19 @@ public void testTimestampNanosSchema() {
     assertEquals(12L, field.getTimestampPrecision().getValue());
   }
 
+  @Test
+  public void testTimestampNanosArraySchema() {
+    com.google.cloud.bigquery.storage.v1.TableSchema protoSchema =
+        BeamRowToStorageApiProto.protoTableSchemaFromBeamSchema(TIMESTAMP_NANOS_ARRAY_SCHEMA);
+
+    assertEquals(1, protoSchema.getFieldsCount());
+    TableFieldSchema field = protoSchema.getFields(0);
+    assertEquals(TableFieldSchema.Type.TIMESTAMP, field.getType());
+    assertEquals(
+        com.google.cloud.bigquery.storage.v1.TableFieldSchema.Mode.REPEATED, field.getMode());
+    assertEquals(12L, field.getTimestampPrecision().getValue());
+  }
+
   @Test
   public void testTimestampNanosDescriptor() throws Exception {
     DescriptorProto descriptor =