Search in sources :

Example 1 with DisplayDataEvaluator

use of org.apache.beam.sdk.transforms.display.DisplayDataEvaluator in project beam by apache.

the class BigtableIOTest method testReadingPrimitiveDisplayData.

@Test
public void testReadingPrimitiveDisplayData() throws IOException, InterruptedException {
    final String table = "fooTable";
    service.createTable(table);
    RowFilter rowFilter = RowFilter.newBuilder().setRowKeyRegexFilter(ByteString.copyFromUtf8("foo.*")).build();
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
    BigtableIO.Read read = BigtableIO.read().withBigtableOptions(BIGTABLE_OPTIONS).withTableId(table).withRowFilter(rowFilter).withBigtableService(service);
    Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
    assertThat("BigtableIO.Read should include the table id in its primitive display data", displayData, Matchers.hasItem(hasDisplayItem("tableId")));
    assertThat("BigtableIO.Read should include the row filter, if it exists, in its primitive " + "display data", displayData, Matchers.hasItem(hasDisplayItem("rowFilter")));
}
Also used : RowFilter(com.google.bigtable.v2.RowFilter) ByteString(com.google.protobuf.ByteString) DisplayDataEvaluator(org.apache.beam.sdk.transforms.display.DisplayDataEvaluator) DisplayData(org.apache.beam.sdk.transforms.display.DisplayData) Test(org.junit.Test)

Example 2 with DisplayDataEvaluator

use of org.apache.beam.sdk.transforms.display.DisplayDataEvaluator in project beam by apache.

the class BigQueryIOTest method testQuerySourcePrimitiveDisplayData.

@Test
public void testQuerySourcePrimitiveDisplayData() throws IOException, InterruptedException {
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
    BigQueryIO.Read read = BigQueryIO.read().fromQuery("foobar").withTestServices(new FakeBigQueryServices().withDatasetService(new FakeDatasetService()).withJobService(new FakeJobService())).withoutValidation();
    Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
    assertThat("BigQueryIO.Read should include the query in its primitive display data", displayData, hasItem(hasDisplayItem("query")));
}
Also used : DisplayDataEvaluator(org.apache.beam.sdk.transforms.display.DisplayDataEvaluator) DisplayData(org.apache.beam.sdk.transforms.display.DisplayData) Test(org.junit.Test)

Example 3 with DisplayDataEvaluator

use of org.apache.beam.sdk.transforms.display.DisplayDataEvaluator in project beam by apache.

the class ReadTest method testPrimitiveDisplayData.

private void testPrimitiveDisplayData(boolean isStreaming) {
    PipelineOptions options = DisplayDataEvaluator.getDefaultOptions();
    options.as(StreamingOptions.class).setStreaming(isStreaming);
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(options);
    SerializableBoundedSource boundedSource = new SerializableBoundedSource() {

        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
            builder.add(DisplayData.item("foo", "bar"));
        }
    };
    SerializableUnboundedSource unboundedSource = new SerializableUnboundedSource() {

        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
            builder.add(DisplayData.item("foo", "bar"));
        }
    };
    Read.Bounded<String> bounded = Read.from(boundedSource);
    BoundedReadFromUnboundedSource<String> unbounded = Read.from(unboundedSource).withMaxNumRecords(1234);
    Set<DisplayData> boundedDisplayData = evaluator.displayDataForPrimitiveSourceTransforms(bounded);
    assertThat(boundedDisplayData, hasItem(hasDisplayItem("source", boundedSource.getClass())));
    assertThat(boundedDisplayData, hasItem(includesDisplayDataFor("source", boundedSource)));
    Set<DisplayData> unboundedDisplayData = evaluator.displayDataForPrimitiveSourceTransforms(unbounded);
    assertThat(unboundedDisplayData, hasItem(hasDisplayItem("source")));
    assertThat(unboundedDisplayData, hasItem(includesDisplayDataFor("source", unboundedSource)));
}
Also used : PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) StreamingOptions(org.apache.beam.sdk.options.StreamingOptions) DisplayDataEvaluator(org.apache.beam.sdk.transforms.display.DisplayDataEvaluator) DisplayData(org.apache.beam.sdk.transforms.display.DisplayData)

Example 4 with DisplayDataEvaluator

use of org.apache.beam.sdk.transforms.display.DisplayDataEvaluator in project beam by apache.

the class DatastoreV1Test method testSourcePrimitiveDisplayData.

@Test
public void testSourcePrimitiveDisplayData() {
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
    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(String.format(assertMessage, "project id"), displayData, hasItem(hasDisplayItem("projectId", PROJECT_ID)));
    assertThat(String.format(assertMessage, "number of query splits"), displayData, hasItem(hasDisplayItem("numQuerySplits", numSplits)));
}
Also used : PCollection(org.apache.beam.sdk.values.PCollection) DisplayDataEvaluator(org.apache.beam.sdk.transforms.display.DisplayDataEvaluator) PBegin(org.apache.beam.sdk.values.PBegin) DisplayData(org.apache.beam.sdk.transforms.display.DisplayData) Test(org.junit.Test)

Example 5 with DisplayDataEvaluator

use of org.apache.beam.sdk.transforms.display.DisplayDataEvaluator in project beam by apache.

the class DatastoreV1Test method testDeleteKeyPrimitiveDisplayData.

@Test
public void testDeleteKeyPrimitiveDisplayData() {
    int hintNumWorkers = 10;
    DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
    PTransform<PCollection<Key>, ?> write = DatastoreIO.v1().deleteKey().withProjectId("myProject").withHintNumWorkers(hintNumWorkers);
    Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write);
    assertThat("DatastoreIO write should include the project in its primitive display data", displayData, hasItem(hasDisplayItem("projectId")));
    assertThat("DatastoreIO write should include the deleteKeyFn in its primitive display data", displayData, hasItem(hasDisplayItem("deleteKeyFn")));
    assertThat("DatastoreIO write should include ramp-up throttling worker count hint if enabled", displayData, hasItem(hasDisplayItem("hintNumWorkers", hintNumWorkers)));
}
Also used : PCollection(org.apache.beam.sdk.values.PCollection) DisplayDataEvaluator(org.apache.beam.sdk.transforms.display.DisplayDataEvaluator) DisplayData(org.apache.beam.sdk.transforms.display.DisplayData) Test(org.junit.Test)

Aggregations

DisplayData (org.apache.beam.sdk.transforms.display.DisplayData)16 DisplayDataEvaluator (org.apache.beam.sdk.transforms.display.DisplayDataEvaluator)16 Test (org.junit.Test)15 PCollection (org.apache.beam.sdk.values.PCollection)7 Category (org.junit.experimental.categories.Category)5 ByteString (com.google.protobuf.ByteString)2 RowFilter (com.google.bigtable.v2.RowFilter)1 ImmutableSet (com.google.common.collect.ImmutableSet)1 HashSet (java.util.HashSet)1 Set (java.util.Set)1 GenericRecord (org.apache.avro.generic.GenericRecord)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 StreamingOptions (org.apache.beam.sdk.options.StreamingOptions)1 PBegin (org.apache.beam.sdk.values.PBegin)1