From 4a82b3cd1aa821d72efc7046329db075c10ae344 Mon Sep 17 00:00:00 2001 From: Praful Makani Date: Tue, 18 Feb 2020 04:49:49 +0530 Subject: [PATCH] feat: expose slotms field in querystage (#172) --- .../com/google/cloud/bigquery/QueryStage.java | 24 ++++++++++++++++--- .../google/cloud/bigquery/QueryStageTest.java | 7 ++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryStage.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryStage.java index 825d653dc..5fc37b98e 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryStage.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryStage.java @@ -163,6 +163,7 @@ static QueryStep fromPb(com.google.api.services.bigquery.model.ExplainQueryStep private final long writeMsMax; private final double writeRatioAvg; private final double writeRatioMax; + private final long slotMs; static final class Builder { @@ -195,6 +196,7 @@ static final class Builder { private long writeMsMax; private double writeRatioAvg; private double writeRatioMax; + private long slotMs; private Builder() {} @@ -343,6 +345,11 @@ Builder setWriteRatioMax(double writeRatioMax) { return this; } + Builder setSlotMs(long slotMs) { + this.slotMs = slotMs; + return this; + } + QueryStage build() { return new QueryStage(this); } @@ -378,6 +385,7 @@ QueryStage build() { writeMsMax = builder.writeMsMax; writeRatioAvg = builder.writeRatioAvg; writeRatioMax = builder.writeRatioMax; + slotMs = builder.slotMs; } /** Returns the number of parallel input segments completed. */ @@ -551,6 +559,11 @@ public double getWriteRatioMax() { return writeRatioMax; } + /** Returns the slot-milliseconds used by the stage. */ + public long getSlotMs() { + return slotMs; + } + @Override public String toString() { return MoreObjects.toStringHelper(this) @@ -583,6 +596,7 @@ public String toString() { .add("writeMsMax", writeMsMax) .add("writeRatioAvg", writeRatioAvg) .add("writeRatioMax", writeRatioMax) + .add("slotMs", slotMs) .toString(); } @@ -617,7 +631,8 @@ public final int hashCode() { writeMsAvg, writeMsMax, writeRatioAvg, - writeRatioMax); + writeRatioMax, + slotMs); } @Override @@ -657,7 +672,8 @@ public final boolean equals(Object obj) { && Objects.equals(steps, other.steps) && Objects.equals(name, other.name) && Objects.equals(status, other.status) - && Objects.equals(inputStages, other.inputStages); + && Objects.equals(inputStages, other.inputStages) + && Objects.equals(slotMs, other.slotMs); } static Builder newBuilder() { @@ -694,7 +710,8 @@ ExplainQueryStage toPb() { .setWriteMsAvg(writeMsAvg) .setWriteMsMax(writeMsMax) .setWriteRatioAvg(writeRatioAvg) - .setWriteRatioMax(writeRatioMax); + .setWriteRatioMax(writeRatioMax) + .setSlotMs(slotMs); if (steps != null) { stagePb.setSteps(Lists.transform(steps, QueryStep.TO_PB_FUNCTION)); } @@ -734,6 +751,7 @@ static QueryStage fromPb(com.google.api.services.bigquery.model.ExplainQueryStag builder.setWriteMsMax(stagePb.getWriteMsMax()); builder.setWriteRatioAvg(stagePb.getWriteRatioAvg()); builder.setWriteRatioMax(stagePb.getWriteRatioMax()); + builder.setSlotMs(stagePb.getSlotMs()); return builder.build(); } } diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryStageTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryStageTest.java index 0c448516d..69edb3988 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryStageTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryStageTest.java @@ -59,6 +59,7 @@ public class QueryStageTest { private static final long WRITE_MS_MAX = 50; private static final double WRITE_RATIO_AVG = 9.9; private static final double WRITE_RATIO_MAX = 10.10; + private static final long SLOTMS = 1522540800000L; private static final QueryStage QUERY_STAGE = QueryStage.newBuilder() .setCompletedParallelInputs(COMPLETED_PARALLEL_INPUTS) @@ -90,6 +91,7 @@ public class QueryStageTest { .setWriteMsMax(WRITE_MS_MAX) .setWriteRatioAvg(WRITE_RATIO_AVG) .setWriteRatioMax(WRITE_RATIO_MAX) + .setSlotMs(SLOTMS) .build(); @Test @@ -112,6 +114,8 @@ public void testBuilder() { assertEquals(INPUT_STAGES, QUERY_STAGE.getInputStages()); assertEquals(PARALLEL_INPUTS, QUERY_STAGE.getParallelInputs()); assertEquals(NAME, QUERY_STAGE.getName()); + assertEquals(READ_MS_AVG, QUERY_STAGE.getReadMsAvg()); + assertEquals(READ_MS_MAX, QUERY_STAGE.getReadMsMax()); assertEquals(READ_RATIO_AVG, QUERY_STAGE.getReadRatioAvg(), 0); assertEquals(READ_RATIO_MAX, QUERY_STAGE.getReadRatioMax(), 0); assertEquals(RECORDS_READ, QUERY_STAGE.getRecordsRead()); @@ -129,6 +133,7 @@ public void testBuilder() { assertEquals(WRITE_MS_MAX, QUERY_STAGE.getWriteMsMax()); assertEquals(WRITE_RATIO_AVG, QUERY_STAGE.getWriteRatioAvg(), 0); assertEquals(WRITE_RATIO_MAX, QUERY_STAGE.getWriteRatioMax(), 0); + assertEquals(SLOTMS, QUERY_STAGE.getSlotMs()); } @Test @@ -178,7 +183,9 @@ private void compareQueryStage(QueryStage expected, QueryStage value) { assertEquals(expected.getWriteMsMax(), expected.getWriteMsMax()); assertEquals(expected.getWriteRatioAvg(), value.getWriteRatioAvg(), 0); assertEquals(expected.getWriteRatioMax(), value.getWriteRatioMax(), 0); + assertEquals(expected.getSlotMs(), value.getSlotMs()); assertEquals(expected.hashCode(), value.hashCode()); + assertEquals(expected.toString(), value.toString()); } private void compareQueryStep(QueryStep expected, QueryStep value) {