Skip to content

Commit

Permalink
feat: Add new fields to copy job statistics (#3205)
Browse files Browse the repository at this point in the history 
* 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.
  • Loading branch information
@PhongChuong
PhongChuong authored Mar 25, 2024
1 parent 7a24d3e commit 64bdda8
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 2 deletions.
54 changes: 52 additions & 2 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/JobStatistics.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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() {
Expand Down
18 changes: 18 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/JobStatisticsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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()));
Expand Down Expand Up @@ -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);
Expand Down
39 changes: 39 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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());
Expand All @@ -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";
Expand Down

0 comments on commit 64bdda8

Please sign in to comment.