Java Code Examples for org.apache.beam.sdk.transforms.display.DisplayData#from()

@Test
public void testDisplayData() {
  DisplayData displayData =
      DisplayData.from(
          XmlIO.<Integer>read()
              .from("foo.xml")
              .withRootElement("bird")
              .withRecordElement("cat")
              .withMinBundleSize(1234)
              .withRecordClass(Integer.class));

  Assert.assertThat(displayData, hasDisplayItem("filePattern", "foo.xml"));
  Assert.assertThat(displayData, hasDisplayItem("rootElement", "bird"));
  Assert.assertThat(displayData, hasDisplayItem("recordElement", "cat"));
  Assert.assertThat(displayData, hasDisplayItem("recordClass", Integer.class));
  Assert.assertThat(displayData, hasDisplayItem("minBundleSize", 1234));
}
 

@Test
public void testWriteWithPubsubGrpcClientFactory() {
  String topic = "projects/project/topics/topic";
  PubsubIO.Write<?> write =
      PubsubIO.writeStrings()
          .to(topic)
          .withClientFactory(PubsubGrpcClient.FACTORY)
          .withTimestampAttribute("myTimestamp")
          .withIdAttribute("myId");

  DisplayData displayData = DisplayData.from(write);

  assertThat(displayData, hasDisplayItem("topic", topic));
  assertThat(displayData, hasDisplayItem("timestampAttribute", "myTimestamp"));
  assertThat(displayData, hasDisplayItem("idAttribute", "myId"));
}
 

@Test
public void testSinkDisplayData() {
  try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
    KafkaIO.Write<Integer, Long> write =
        KafkaIO.<Integer, Long>write()
            .withBootstrapServers("myServerA:9092,myServerB:9092")
            .withTopic("myTopic")
            .withValueSerializer(LongSerializer.class)
            .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))
            .withProducerConfigUpdates(ImmutableMap.of("retry.backoff.ms", 100));

    DisplayData displayData = DisplayData.from(write);

    assertThat(displayData, hasDisplayItem("topic", "myTopic"));
    assertThat(displayData, hasDisplayItem("bootstrap.servers", "myServerA:9092,myServerB:9092"));
    assertThat(displayData, hasDisplayItem("retries", 3));
    assertThat(displayData, hasDisplayItem("retry.backoff.ms", 100));
  }
}
 

@Test
public void testWithOutputTagsDisplayData() {
  DoFn<String, String> fn =
      new DoFn<String, String>() {
        @ProcessElement
        public void proccessElement(ProcessContext c) {}

        @Override
        public void populateDisplayData(Builder builder) {
          builder.add(DisplayData.item("fnMetadata", "foobar"));
        }
      };

  ParDo.MultiOutput<String, String> parDo =
      ParDo.of(fn).withOutputTags(new TupleTag<>(), TupleTagList.empty());

  DisplayData displayData = DisplayData.from(parDo);
  assertThat(displayData, includesDisplayDataFor("fn", fn));
  assertThat(displayData, hasDisplayItem("fn", fn.getClass()));
}
 

@Test
public void testPopulateDisplayData() {
  DisplayData displayData =
      DisplayData.from(
          config.withProjectId(PROJECT_ID).withInstanceId(INSTANCE_ID).withTableId(TABLE_ID)
              ::populateDisplayData);

  assertThat(
      displayData,
      hasDisplayItem(
          allOf(
              hasKey("projectId"), hasLabel("Bigtable Project Id"), hasValue(PROJECT_ID.get()))));

  assertThat(
      displayData,
      hasDisplayItem(
          allOf(
              hasKey("instanceId"),
              hasLabel("Bigtable Instance Id"),
              hasValue(INSTANCE_ID.get()))));

  assertThat(
      displayData,
      hasDisplayItem(
          allOf(hasKey("tableId"), hasLabel("Bigtable Table Id"), hasValue(TABLE_ID.get()))));
}
 

@Test
public void testDoFnWithContextDisplayData() {
  DoFn<String, String> fn =
      new DoFn<String, String>() {
        @ProcessElement
        public void proccessElement(ProcessContext c) {}

        @Override
        public void populateDisplayData(Builder builder) {
          builder.add(DisplayData.item("fnMetadata", "foobar"));
        }
      };

  SingleOutput<String, String> parDo = ParDo.of(fn);

  DisplayData displayData = DisplayData.from(parDo);
  assertThat(displayData, includesDisplayDataFor("fn", fn));
  assertThat(displayData, hasDisplayItem("fn", fn.getClass()));
}
 

@Test
public void testDisplayDataExcludesValuesAccessedButNeverSet() {
  HasDefaults options = PipelineOptionsFactory.as(HasDefaults.class);
  assertEquals("bar", options.getFoo());

  DisplayData data = DisplayData.from(options);
  assertThat(data, not(hasDisplayItem("foo")));
}
 

@Test
public void testDisplayData() {
  Top.Natural<Integer> comparer = new Top.Natural<>();
  PTransform<?, ?> approxQuanitiles = ApproximateQuantiles.globally(20, comparer);
  DisplayData displayData = DisplayData.from(approxQuanitiles);

  assertThat(displayData, hasDisplayItem("numQuantiles", 20));
  assertThat(displayData, hasDisplayItem("comparer", comparer.getClass()));
}
 

@Test
public void testWriteDisplayData() {
  Write write = DatastoreIO.v1().write().withProjectId(PROJECT_ID);

  DisplayData displayData = DisplayData.from(write);

  assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
}
 

@Test
public void testBoundedDisplayDataSubrange() {
  PTransform<?, ?> input = GenerateSequence.from(12).to(1234);
  DisplayData displayData = DisplayData.from(input);
  assertThat(displayData, hasDisplayItem("from", 12));
  assertThat(displayData, hasDisplayItem("to", 1234));
}
 

/**
 * @param ctx       provides translation context
 * @param beamNode  the beam node to be translated
 * @param transform transform which can be obtained from {@code beamNode}
 */
@PrimitiveTransformTranslator(Read.Unbounded.class)
private static void unboundedReadTranslator(final PipelineTranslationContext ctx,
                                            final TransformHierarchy.Node beamNode,
                                            final Read.Unbounded<?> transform) {
  final IRVertex vertex = new BeamUnboundedSourceVertex<>(transform.getSource(), DisplayData.from(transform));
  ctx.addVertex(vertex);
  beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input));
  beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output));
}
 

@Test
public void testDisplayDataForWrappedFn() {
  UniqueInts combineFn =
      new UniqueInts() {
        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
          builder.add(DisplayData.item("foo", "bar"));
        }
      };
  Combine.PerKey<?, ?, ?> combine = Combine.perKey(combineFn);
  DisplayData displayData = DisplayData.from(combine);

  assertThat(displayData, hasDisplayItem("combineFn", combineFn.getClass()));
  assertThat(displayData, hasDisplayItem(hasNamespace(combineFn.getClass())));
}
 

@Test
public void testDisplayDataNullValuesConvertedToEmptyString() throws Exception {
  FooOptions options = PipelineOptionsFactory.as(FooOptions.class);
  options.setFoo(null);

  DisplayData data = DisplayData.from(options);
  assertThat(data, hasDisplayItem("foo", ""));

  FooOptions deserializedOptions = serializeDeserialize(FooOptions.class, options);
  DisplayData deserializedData = DisplayData.from(deserializedOptions);
  assertThat(deserializedData, hasDisplayItem("foo", ""));
}
 

@Test
public void testDisplayDataJsonValueSetAfterDeserialization() throws Exception {
  FooOptions options = PipelineOptionsFactory.as(FooOptions.class);
  options.setFoo("bar");
  DisplayData data = DisplayData.from(options);
  assertThat(data, hasDisplayItem("foo", "bar"));

  FooOptions deserializedOptions = serializeDeserialize(FooOptions.class, options);
  deserializedOptions.setFoo("baz");
  DisplayData dataAfterDeserialization = DisplayData.from(deserializedOptions);
  assertThat(dataAfterDeserialization, hasDisplayItem("foo", "baz"));
}
 

@Test
public void testReadingDisplayData() {
  HBaseIO.Read read = HBaseIO.read().withConfiguration(conf).withTableId("fooTable");
  DisplayData displayData = DisplayData.from(read);
  assertThat(displayData, hasDisplayItem("tableId", "fooTable"));
  assertThat(displayData, hasDisplayItem("configuration"));
}
 

@Test
public void testDisplayDataExcludedFromOverriddenBaseClass() {
  ExtendsBaseOptions options = PipelineOptionsFactory.as(ExtendsBaseOptions.class);
  options.setFoo("bar");

  DisplayData displayData = DisplayData.from(options);
  assertThat(displayData, not(hasDisplayItem(hasNamespace(BaseOptions.class))));
}
 

@Test
public void testDisplayDataArrayValue() throws Exception {
  ArrayOptions options = PipelineOptionsFactory.as(ArrayOptions.class);
  options.setDeepArray(new String[][] {new String[] {"a", "b"}, new String[] {"c"}});
  options.setDeepPrimitiveArray(new int[][] {new int[] {1, 2}, new int[] {3}});

  DisplayData data = DisplayData.from(options);
  assertThat(data, hasDisplayItem("deepArray", "[[a, b], [c]]"));
  assertThat(data, hasDisplayItem("deepPrimitiveArray", "[[1, 2], [3]]"));

  ArrayOptions deserializedOptions = serializeDeserialize(ArrayOptions.class, options);
  DisplayData deserializedData = DisplayData.from(deserializedOptions);
  assertThat(deserializedData, hasDisplayItem("deepPrimitiveArray", "[[1, 2], [3]]"));
}
 

@Test
public void testDisplayData() {
  AvroSource<Bird> source =
      AvroSource.from("foobar.txt").withSchema(Bird.class).withMinBundleSize(1234);

  DisplayData displayData = DisplayData.from(source);
  assertThat(displayData, hasDisplayItem("filePattern", "foobar.txt"));
  assertThat(displayData, hasDisplayItem("minBundleSize", 1234));
}
 

@Test
public void testDeleteKeyDisplayData() {
  DeleteKey deleteKey = DatastoreIO.v1().deleteKey().withProjectId(PROJECT_ID);

  DisplayData displayData = DisplayData.from(deleteKey);

  assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
}
 

@Test
@Category(NeedsRunner.class)
public void testWriteWithWritableByteChannelFactory() throws Exception {
  Coder<String> coder = StringUtf8Coder.of();
  String outputName = "file.txt";
  ResourceId baseDir =
      FileSystems.matchNewResource(
          Files.createTempDirectory(tempFolder.getRoot().toPath(), "testwrite").toString(), true);

  PCollection<String> input = p.apply(Create.of(Arrays.asList(LINES2_ARRAY)).withCoder(coder));

  final WritableByteChannelFactory writableByteChannelFactory =
      new DrunkWritableByteChannelFactory();
  TextIO.Write write =
      TextIO.write()
          .to(
              baseDir
                  .resolve(outputName, ResolveOptions.StandardResolveOptions.RESOLVE_FILE)
                  .toString())
          .withoutSharding()
          .withWritableByteChannelFactory(writableByteChannelFactory);
  DisplayData displayData = DisplayData.from(write);
  assertThat(displayData, hasDisplayItem("writableByteChannelFactory", "DRUNK"));

  input.apply(write);

  p.run();

  final List<String> drunkElems = new ArrayList<>(LINES2_ARRAY.length * 2 + 2);
  for (String elem : LINES2_ARRAY) {
    drunkElems.add(elem);
    drunkElems.add(elem);
  }
  assertOutputFiles(
      drunkElems.toArray(new String[0]),
      null,
      null,
      1,
      baseDir.resolve(
          outputName + writableByteChannelFactory.getSuggestedFilenameSuffix(),
          ResolveOptions.StandardResolveOptions.RESOLVE_FILE),
      write.inner.getShardTemplate());
}