How to use

DisplayDataEvaluator

in

org.apache.beam.sdk.transforms.display

 @Test
 public void testSourceTransform() {
  PTransform<? super PBegin, ? extends POutput> myTransform = TextIO.read().from("foo.*");

  DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
  Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(myTransform);

  assertThat(displayData, hasItem(hasDisplayItem("filePattern", "foo.*")));
 }
}

/**
 * Traverse the specified {@link PTransform}, collecting {@link DisplayData} registered on the
 * inner primitive {@link PTransform PTransforms}.
 *
 * @return the set of {@link DisplayData} for primitive {@link PTransform PTransforms}.
 */
public <InputT> Set<DisplayData> displayDataForPrimitiveTransforms(
  final PTransform<? super PCollection<InputT>, ? extends POutput> root) {
 return displayDataForPrimitiveTransforms(root, null);
}

/** Create a new {@link DisplayDataEvaluator} using the specified {@link PipelineOptions}. */
public static DisplayDataEvaluator create(PipelineOptions pipelineOptions) {
 return new DisplayDataEvaluator(pipelineOptions);
}

/**
 * Create a new {@link DisplayDataEvaluator} using options returned from {@link
 * #getDefaultOptions()}.
 */
public static DisplayDataEvaluator create() {
 return create(getDefaultOptions());
}

@Test
public void testPrimitiveTransform() {
 PTransform<? super PCollection<Integer>, ? super PCollection<Integer>> myTransform =
   ParDo.of(
     new DoFn<Integer, Integer>() {
      @ProcessElement
      public void processElement(ProcessContext c) throws Exception {}
      @Override
      public void populateDisplayData(DisplayData.Builder builder) {
       builder.add(DisplayData.item("foo", "bar"));
      }
     });
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(myTransform);
 assertThat(displayData, hasItem(hasDisplayItem("foo")));
}

/**
 * Traverse the specified source {@link PTransform}, collecting {@link DisplayData} registered on
 * the inner primitive {@link PTransform PTransforms}.
 *
 * @param root The source root {@link PTransform} to traverse
 * @return the set of {@link DisplayData} for primitive source {@link PTransform PTransforms}.
 */
public Set<DisplayData> displayDataForPrimitiveSourceTransforms(
  final PTransform<? super PBegin, ? extends POutput> root) {
 Pipeline pipeline = Pipeline.create(options);
 pipeline.apply("SourceTransform", root);
 return displayDataForPipeline(pipeline, root);
}

@Test
@Category(ValidatesRunner.class)
public void testPrimitiveDisplayData() {
 SimpleFunction<Integer, ?> mapFn =
   new SimpleFunction<Integer, Integer>() {
    @Override
    public Integer apply(Integer input) {
     return input;
    }
   };
 MapElements<Integer, ?> map = MapElements.via(mapFn);
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(map);
 assertThat(
   "MapElements should include the mapFn in its primitive display data",
   displayData,
   hasItem(hasDisplayItem("class", mapFn.getClass())));
}

/**
 * Traverse the specified {@link PTransform}, collecting {@link DisplayData} registered on the
 * inner primitive {@link PTransform PTransforms}.
 *
 * @param root The root {@link PTransform} to traverse
 * @param inputCoder The coder to set for the {@link PTransform} input, or null to infer the
 *     default coder.
 * @return the set of {@link DisplayData} for primitive {@link PTransform PTransforms}.
 */
public <InputT> Set<DisplayData> displayDataForPrimitiveTransforms(
  final PTransform<? super PCollection<InputT>, ? extends POutput> root,
  Coder<InputT> inputCoder) {
 Create.Values<InputT> input;
 if (inputCoder != null) {
  input = Create.empty(inputCoder);
 } else {
  // These types don't actually work, but the pipeline will never be run
  input = (Create.Values<InputT>) Create.empty(VoidCoder.of());
 }
 Pipeline pipeline = Pipeline.create(options);
 pipeline.apply("Input", input).apply("Transform", root);
 return displayDataForPipeline(pipeline, root);
}

 @Test
 @Category(ValidatesRunner.class)
 public void testPrimitiveWriteDisplayData() {
  DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
  PubsubIO.Write<?> write = PubsubIO.writeStrings().to("projects/project/topics/topic");

  Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write);
  assertThat(
    "PubsubIO.Write should include the topic in its primitive display data",
    displayData,
    hasItem(hasDisplayItem("topic")));
 }
}

@Test
@Category(ValidatesRunner.class)
public void testPrimitiveReadDisplayData() {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 AvroIO.Read<GenericRecord> read =
   AvroIO.readGenericRecords(Schema.create(Schema.Type.STRING)).from("/foo.*");
 Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
 assertThat(
   "AvroIO.Read should include the file pattern in its primitive transform",
   displayData,
   hasItem(hasDisplayItem("filePattern")));
}

@Test
public void testCompositeTransform() {
 PTransform<? super PCollection<String>, ? super POutput> myTransform =
   new PTransform<PCollection<String>, POutput>() {
    @Override
    public PCollection<String> expand(PCollection<String> input) {
     return input.apply(
       ParDo.of(
         new DoFn<String, String>() {
          @ProcessElement
          public void processElement(ProcessContext c) throws Exception {
           c.output(c.element());
          }
          @Override
          public void populateDisplayData(DisplayData.Builder builder) {
           builder.add(DisplayData.item("primitiveKey", "primitiveValue"));
          }
         }));
    }
    @Override
    public void populateDisplayData(DisplayData.Builder builder) {
     builder.add(DisplayData.item("compositeKey", "compositeValue"));
    }
   };
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(myTransform);
 assertThat(displayData, not(hasItem(hasDisplayItem("compositeKey", "compositeValue"))));
 assertThat(displayData, hasItem(hasDisplayItem("primitiveKey", "primitiveValue")));
}

@Test
@Category(ValidatesRunner.class)
public void testPrimitiveReadDisplayData() {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 TextIO.Read read = TextIO.read().from("foobar");
 Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
 assertThat(
   "TextIO.Read should include the file prefix in its primitive display data",
   displayData,
   hasItem(hasDisplayItem(hasValue(startsWith("foobar")))));
}

@Test
public void testDeleteKeyPrimitiveDisplayData() {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 PTransform<PCollection<Key>, ?> write = DatastoreIO.v1().deleteKey().withProjectId("myProject");
 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")));
}

@Test
@Category({ValidatesRunner.class, UsesUnboundedPCollections.class})
public void testPrimitiveReadDisplayData() {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 Set<DisplayData> displayData;
 PubsubIO.Read<String> baseRead = PubsubIO.readStrings();
 // Reading from a subscription.
 PubsubIO.Read<String> read =
   baseRead.fromSubscription("projects/project/subscriptions/subscription");
 displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
 assertThat(
   "PubsubIO.Read should include the subscription in its primitive display data",
   displayData,
   hasItem(hasDisplayItem("subscription")));
 // Reading from a topic.
 read = baseRead.fromTopic("projects/project/topics/topic");
 displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
 assertThat(
   "PubsubIO.Read should include the topic in its primitive display data",
   displayData,
   hasItem(hasDisplayItem("topic")));
}

@Test
public void testWritePrimitiveDisplayData() {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 PTransform<PCollection<Entity>, ?> write = DatastoreIO.v1().write().withProjectId("myProject");
 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 upsertFn in its primitive display data",
   displayData,
   hasItem(hasDisplayItem("upsertFn")));
}

@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)));
}

@Test
public void testDeleteEntityPrimitiveDisplayData() {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 PTransform<PCollection<Entity>, ?> write =
   DatastoreIO.v1().deleteEntity().withProjectId("myProject");
 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 deleteEntityFn in its primitive display data",
   displayData,
   hasItem(hasDisplayItem("deleteEntityFn")));
}

@Test
public void testTableSourcePrimitiveDisplayData() throws IOException, InterruptedException {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 BigQueryIO.Read read =
   BigQueryIO.read()
     .from("project:dataset.tableId")
     .withTestServices(
       new FakeBigQueryServices()
         .withDatasetService(new FakeDatasetService())
         .withJobService(new FakeJobService()))
     .withoutValidation();
 Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
 assertThat(
   "BigQueryIO.Read should include the table spec in its primitive display data",
   displayData,
   hasItem(hasDisplayItem("table")));
}

@Test
@Category(ValidatesRunner.class)
public void testCombinePerKeyPrimitiveDisplayData() {
 DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
 UniqueInts combineFn = new UniqueInts();
 PTransform<PCollection<KV<Integer, Integer>>, ? extends POutput> combine =
   Combine.perKey(combineFn);
 Set<DisplayData> displayData =
   evaluator.displayDataForPrimitiveTransforms(
     combine, KvCoder.of(VarIntCoder.of(), VarIntCoder.of()));
 assertThat(
   "Combine.perKey should include the combineFn in its primitive transform",
   displayData,
   hasItem(hasDisplayItem("combineFn", combineFn.getClass())));
}

@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")));
}