diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java index 16ed4b5c0c0d..69f438055da9 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java @@ -953,9 +953,23 @@ public SpannerWriteResult expand(PCollection input) { @Override public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); + populateDisplayDataWithParamaters(builder); + } + + private void populateDisplayDataWithParamaters(DisplayData.Builder builder) { getSpannerConfig().populateDisplayData(builder); builder.add( - DisplayData.item("batchSizeBytes", getBatchSizeBytes()).withLabel("Batch Size in Bytes")); + DisplayData.item("batchSizeBytes", getBatchSizeBytes()) + .withLabel("Max batch size in bytes")); + builder.add( + DisplayData.item("maxNumMutations", getMaxNumMutations()) + .withLabel("Max number of mutated cells in each batch")); + builder.add( + DisplayData.item("maxNumRows", getMaxNumRows()) + .withLabel("Max number of rows in each batch")); + builder.add( + DisplayData.item("groupingFactor", getGroupingFactor()) + .withLabel("Number of batches to sort over")); } } @@ -992,6 +1006,12 @@ public WriteGrouped(Write spec) { this.spec = spec; } + @Override + public void populateDisplayData(DisplayData.Builder builder) { + super.populateDisplayData(builder); + spec.populateDisplayDataWithParamaters(builder); + } + @Override public SpannerWriteResult expand(PCollection input) { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java index 6129882b4cf7..aaca469105c4 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java @@ -589,20 +589,50 @@ public void retryOnAbortedAndDeadlineExceeded() throws InterruptedException { } @Test - public void displayData() throws Exception { + public void displayDataWrite() throws Exception { SpannerIO.Write write = SpannerIO.write() .withProjectId("test-project") .withInstanceId("test-instance") .withDatabaseId("test-database") - .withBatchSizeBytes(123); + .withBatchSizeBytes(123) + .withMaxNumMutations(456) + .withMaxNumRows(789) + .withGroupingFactor(100); DisplayData data = DisplayData.from(write); - assertThat(data.items(), hasSize(4)); + assertThat(data.items(), hasSize(7)); assertThat(data, hasDisplayItem("projectId", "test-project")); assertThat(data, hasDisplayItem("instanceId", "test-instance")); assertThat(data, hasDisplayItem("databaseId", "test-database")); assertThat(data, hasDisplayItem("batchSizeBytes", 123)); + assertThat(data, hasDisplayItem("maxNumMutations", 456)); + assertThat(data, hasDisplayItem("maxNumRows", 789)); + assertThat(data, hasDisplayItem("groupingFactor", 100)); + } + + @Test + public void displayDataWriteGrouped() throws Exception { + SpannerIO.WriteGrouped writeGrouped = + SpannerIO.write() + .withProjectId("test-project") + .withInstanceId("test-instance") + .withDatabaseId("test-database") + .withBatchSizeBytes(123) + .withMaxNumMutations(456) + .withMaxNumRows(789) + .withGroupingFactor(100) + .grouped(); + + DisplayData data = DisplayData.from(writeGrouped); + assertThat(data.items(), hasSize(7)); + assertThat(data, hasDisplayItem("projectId", "test-project")); + assertThat(data, hasDisplayItem("instanceId", "test-instance")); + assertThat(data, hasDisplayItem("databaseId", "test-database")); + assertThat(data, hasDisplayItem("batchSizeBytes", 123)); + assertThat(data, hasDisplayItem("maxNumMutations", 456)); + assertThat(data, hasDisplayItem("maxNumRows", 789)); + assertThat(data, hasDisplayItem("groupingFactor", 100)); } @Test