feat: Add new fields to copy job statistics (#3205)

* feat: Add new fields to copy job statistics

Specifically, [copiedRows, copiedLogicalBytes] are added to CopyStatistics.

* Fix issue where IT verified copied rows twice instead of copied logical bytes.

* Add additional test to verified non-zero copy statistic case.

Author: PhongChuong

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java

@@ -21,6 +21,7 @@
 import com.google.api.services.bigquery.model.JobStatistics2;
 import com.google.api.services.bigquery.model.JobStatistics3;
 import com.google.api.services.bigquery.model.JobStatistics4;
+import com.google.api.services.bigquery.model.JobStatistics5;
 import com.google.api.services.bigquery.model.QueryParameter;
 import com.google.cloud.StringEnumType;
 import com.google.cloud.StringEnumValue;
@@ -51,14 +52,36 @@ public abstract class JobStatistics implements Serializable {
   /** A Google BigQuery Copy Job statistics. */
   public static class CopyStatistics extends JobStatistics {
 
-    private static final long serialVersionUID = 8218325588441660938L;
+    private static final long serialVersionUID = 8218325588441660939L;
+
+    private final Long copiedLogicalBytes;
+
+    private final Long copiedRows;
 
     static final class Builder extends JobStatistics.Builder<CopyStatistics, Builder> {
 
+      private Long copiedLogicalBytes;
+
+      private Long copiedRows;
+
       private Builder() {}
 
       private Builder(com.google.api.services.bigquery.model.JobStatistics statisticsPb) {
         super(statisticsPb);
+        if (statisticsPb.getCopy() != null) {
+          this.copiedLogicalBytes = statisticsPb.getCopy().getCopiedLogicalBytes();
+          this.copiedRows = statisticsPb.getCopy().getCopiedRows();
+        }
+      }
+
+      Builder setCopiedLogicalBytes(long copiedLogicalBytes) {
+        this.copiedLogicalBytes = copiedLogicalBytes;
+        return self();
+      }
+
+      Builder setCopiedRows(long copiedRows) {
+        this.copiedRows = copiedRows;
+        return self();
       }
 
       @Override
@@ -69,6 +92,25 @@ CopyStatistics build() {
 
     private CopyStatistics(Builder builder) {
       super(builder);
+      this.copiedLogicalBytes = builder.copiedLogicalBytes;
+      this.copiedRows = builder.copiedRows;
+    }
+
+    /** Returns number of logical bytes copied to the destination table. */
+    public Long getCopiedLogicalBytes() {
+      return copiedLogicalBytes;
+    }
+
+    /** Returns number of rows copied to the destination table. */
+    public Long getCopiedRows() {
+      return copiedRows;
+    }
+
+    @Override
+    ToStringHelper toStringHelper() {
+      return super.toStringHelper()
+          .add("copiedLogicalBytes", copiedLogicalBytes)
+          .add("copiedRows", copiedRows);
     }
 
     @Override
@@ -81,7 +123,15 @@ public final boolean equals(Object obj) {
 
     @Override
     public final int hashCode() {
-      return baseHashCode();
+      return Objects.hash(baseHashCode(), copiedLogicalBytes, copiedRows);
+    }
+
+    @Override
+    com.google.api.services.bigquery.model.JobStatistics toPb() {
+      JobStatistics5 copyStatisticsPb = new JobStatistics5();
+      copyStatisticsPb.setCopiedLogicalBytes(copiedLogicalBytes);
+      copyStatisticsPb.setCopiedRows(copiedRows);
+      return super.toPb().setCopy(copyStatisticsPb);
     }
 
     static Builder newBuilder() {

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java

@@ -85,11 +85,15 @@ public class JobStatisticsTest {
   private static final Long SLOTMS = 12545L;
   private static final String TRANSACTION_ID = UUID.randomUUID().toString().substring(0, 8);
   private static final String SESSION_ID = UUID.randomUUID().toString().substring(0, 8);
+  private static final Long COPIED_ROW = 1L;
+  private static final Long COPIED_LOGICAL_BYTES = 2L;
   private static final CopyStatistics COPY_STATISTICS =
       CopyStatistics.newBuilder()
           .setCreationTimestamp(CREATION_TIME)
           .setEndTime(END_TIME)
           .setStartTime(START_TIME)
+          .setCopiedRows(COPIED_ROW)
+          .setCopiedLogicalBytes(COPIED_LOGICAL_BYTES)
           .build();
   private static final ExtractStatistics EXTRACT_STATISTICS =
       ExtractStatistics.newBuilder()
@@ -262,6 +266,12 @@ public void testBuilder() {
     assertEquals(FILE_COUNT, EXTRACT_STATISTICS.getDestinationUriFileCounts());
     assertEquals(INPUT_BYTES, EXTRACT_STATISTICS.getInputBytes());
 
+    assertEquals(CREATION_TIME, COPY_STATISTICS.getCreationTime());
+    assertEquals(START_TIME, COPY_STATISTICS.getStartTime());
+    assertEquals(END_TIME, COPY_STATISTICS.getEndTime());
+    assertEquals(COPIED_LOGICAL_BYTES, COPY_STATISTICS.getCopiedLogicalBytes());
+    assertEquals(COPIED_ROW, COPY_STATISTICS.getCopiedRows());
+
     assertEquals(CREATION_TIME, LOAD_STATISTICS.getCreationTime());
     assertEquals(START_TIME, LOAD_STATISTICS.getStartTime());
     assertEquals(END_TIME, LOAD_STATISTICS.getEndTime());
@@ -334,6 +344,7 @@ public void testBuilder() {
   public void testToPbAndFromPb() {
     compareExtractStatistics(
         EXTRACT_STATISTICS, ExtractStatistics.fromPb(EXTRACT_STATISTICS.toPb()));
+    compareCopyStatistics(COPY_STATISTICS, CopyStatistics.fromPb(COPY_STATISTICS.toPb()));
     compareLoadStatistics(LOAD_STATISTICS, LoadStatistics.fromPb(LOAD_STATISTICS.toPb()));
     compareQueryStatistics(QUERY_STATISTICS, QueryStatistics.fromPb(QUERY_STATISTICS.toPb()));
     compareStatistics(COPY_STATISTICS, CopyStatistics.fromPb(COPY_STATISTICS.toPb()));
@@ -400,6 +411,13 @@ private void compareExtractStatistics(ExtractStatistics expected, ExtractStatist
     assertEquals(expected.getInputBytes(), value.getInputBytes());
   }
 
+  private void compareCopyStatistics(CopyStatistics expected, CopyStatistics value) {
+    assertEquals(expected, value);
+    compareStatistics(expected, value);
+    assertEquals(expected.getCopiedLogicalBytes(), value.getCopiedLogicalBytes());
+    assertEquals(expected.getCopiedRows(), value.getCopiedRows());
+  }
+
   private void compareLoadStatistics(LoadStatistics expected, LoadStatistics value) {
     assertEquals(expected, value);
     compareStatistics(expected, value);

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -90,6 +90,7 @@
 import com.google.cloud.bigquery.JobId;
 import com.google.cloud.bigquery.JobInfo;
 import com.google.cloud.bigquery.JobStatistics;
+import com.google.cloud.bigquery.JobStatistics.CopyStatistics;
 import com.google.cloud.bigquery.JobStatistics.ExtractStatistics;
 import com.google.cloud.bigquery.JobStatistics.LoadStatistics;
 import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
@@ -5042,11 +5043,18 @@ public void testCopyJob() throws InterruptedException, TimeoutException {
     assertNotNull(createdTable);
     assertEquals(DATASET, createdTable.getTableId().getDataset());
     assertEquals(sourceTableName, createdTable.getTableId().getTable());
+
     TableId destinationTable = TableId.of(DATASET, destinationTableName);
     CopyJobConfiguration configuration = CopyJobConfiguration.of(destinationTable, sourceTable);
     Job remoteJob = bigquery.create(JobInfo.of(configuration));
     remoteJob = remoteJob.waitFor();
     assertNull(remoteJob.getStatus().getError());
+
+    CopyStatistics copyStatistics = remoteJob.getStatistics();
+    assertNotNull(copyStatistics);
+    assertEquals(0, copyStatistics.getCopiedRows().longValue());
+    assertEquals(0, copyStatistics.getCopiedLogicalBytes().longValue());
+
     Table remoteTable = bigquery.getTable(DATASET, destinationTableName);
     assertNotNull(remoteTable);
     assertEquals(destinationTable.getDataset(), remoteTable.getTableId().getDataset());
@@ -5056,6 +5064,37 @@ public void testCopyJob() throws InterruptedException, TimeoutException {
     assertTrue(remoteTable.delete());
   }
 
+  @Test
+  public void testCopyJobStatistics() throws InterruptedException, TimeoutException {
+    String sourceTableName = "test_copy_job_statistics_source_table";
+    String destinationTableName = "test_copy_job_statistics_destination_table";
+
+    QueryJobConfiguration createTable =
+        QueryJobConfiguration.newBuilder(
+                String.format(
+                    "CREATE TABLE %s AS SELECT num FROM UNNEST(GENERATE_ARRAY(0,5)) as num",
+                    sourceTableName))
+            .setDefaultDataset(DatasetId.of(DATASET))
+            .setUseLegacySql(false)
+            .build();
+    bigquery.query(createTable);
+
+    // Copy the created table.
+    TableId sourceTable = TableId.of(DATASET, sourceTableName);
+    TableId destinationTable = TableId.of(DATASET, destinationTableName);
+    CopyJobConfiguration configuration = CopyJobConfiguration.of(destinationTable, sourceTable);
+    Job remoteJob = bigquery.create(JobInfo.of(configuration));
+    remoteJob = remoteJob.waitFor();
+    assertNull(remoteJob.getStatus().getError());
+
+    CopyStatistics copyStatistics = remoteJob.getStatistics();
+    assertNotNull(copyStatistics);
+    assertEquals(6, copyStatistics.getCopiedRows().longValue());
+    // Assert != 0 since copied logical bytes is may return non-deterministic value due to how the
+    // data is represented.
+    assertNotEquals(0, copyStatistics.getCopiedLogicalBytes().longValue());
+  }
+
   @Test
   public void testSnapshotTableCopyJob() throws InterruptedException {
     String sourceTableName = "test_copy_job_base_table";