Repository: beam Updated Branches: refs/heads/master 16486e2ad -> 0b8932fd3
Adds requested number of splits to display data in DatastoreIO Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/3df8dccd Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/3df8dccd Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/3df8dccd Branch: refs/heads/master Commit: 3df8dccdbe8bd65c45dc0f73f8987946e03c5b93 Parents: 16486e2 Author: Vikas Kedigehalli <[email protected]> Authored: Thu Mar 30 17:31:18 2017 -0700 Committer: Eugene Kirpichov <[email protected]> Committed: Fri Mar 31 16:41:22 2017 -0700 ---------------------------------------------------------------------- .../beam/sdk/io/gcp/datastore/DatastoreV1.java | 4 +++ .../sdk/io/gcp/datastore/DatastoreV1Test.java | 28 +++++++++----------- 2 files changed, 16 insertions(+), 16 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/3df8dccd/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java index 73ac8df..98787fa 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java @@ -793,6 +793,10 @@ public class DatastoreV1 { public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); builder.include("options", options); + if (numSplits > 0) { + builder.add(DisplayData.item("numQuerySplits", numSplits) + .withLabel("Requested number of Query splits")); + } } } http://git-wip-us.apache.org/repos/asf/beam/blob/3df8dccd/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java index af27926..5d3b278 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java @@ -272,23 +272,19 @@ public class DatastoreV1Test { @Test public void testSourcePrimitiveDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); - PTransform<PBegin, PCollection<Entity>> read = DatastoreIO.v1().read().withProjectId( - "myProject").withQuery(Query.newBuilder().build()); - + int numSplits = 98; + PTransform<PBegin, PCollection<Entity>> read = + DatastoreIO.v1().read() + .withProjectId(PROJECT_ID) + .withQuery(Query.newBuilder().build()) + .withNumQuerySplits(numSplits); + + String assertMessage = "DatastoreIO read should include the '%s' in its primitive display data"; Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); - assertThat("DatastoreIO read should include the project in its primitive display data", - displayData, hasItem(hasDisplayItem("projectId"))); - } - - @Test - public void testSourcePrimitiveDisplayDataWithGqlQuery() { - DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); - PTransform<PBegin, PCollection<Entity>> read = DatastoreIO.v1().read().withProjectId( - "myProject").withLiteralGqlQuery(GQL_QUERY); - - Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); - assertThat("DatastoreIO read should include the project in its primitive display data", - displayData, hasItem(hasDisplayItem("projectId"))); + assertThat(String.format(assertMessage, "project id"), + displayData, hasItem(hasDisplayItem("projectId", PROJECT_ID))); + assertThat(String.format(assertMessage, "number of query splits"), + displayData, hasItem(hasDisplayItem("numQuerySplits", numSplits))); } @Test
