org.apache.beam.sdk.transforms.display.DisplayData Java Examples

@Test
public void testDisplayData() {
  final ApproximateUnique.Globally<Integer> specifiedSampleSize =
      ApproximateUnique.globally(1234);
  final ApproximateUnique.PerKey<String, Integer> specifiedMaxError =
      ApproximateUnique.perKey(0.1234);

  assertThat(DisplayData.from(specifiedSampleSize), hasDisplayItem("sampleSize", 1234));

  final DisplayData maxErrorDisplayData = DisplayData.from(specifiedMaxError);
  assertThat(maxErrorDisplayData, hasDisplayItem("maximumEstimationError", 0.1234));
  assertThat(
      "calculated sampleSize should be included",
      maxErrorDisplayData,
      hasDisplayItem("sampleSize"));
}
 

@Test
public void testReadWithPubsubGrpcClientFactory() {
  String topic = "projects/project/topics/topic";
  PubsubIO.Read<String> read =
      PubsubIO.readStrings()
          .fromTopic(StaticValueProvider.of(topic))
          .withClientFactory(PubsubGrpcClient.FACTORY)
          .withTimestampAttribute("myTimestamp")
          .withIdAttribute("myId");

  DisplayData displayData = DisplayData.from(read);

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

@Test
public void testAssignDisplayDataUnchanged() {
  FixedWindows windowFn = FixedWindows.of(Duration.standardHours(5));

  Window<Object> original = Window.into(windowFn);
  WindowingStrategy<?, ?> updated = WindowingStrategy.globalDefault().withWindowFn(windowFn);

  DisplayData displayData = DisplayData.from(new Window.Assign<>(original, updated));

  assertThat(displayData, hasDisplayItem("windowFn", windowFn.getClass()));
  assertThat(displayData, includesDisplayDataFor("windowFn", windowFn));

  assertThat(displayData, not(hasDisplayItem("trigger")));
  assertThat(displayData, not(hasDisplayItem("accumulationMode")));
  assertThat(displayData, not(hasDisplayItem("allowedLateness")));
  assertThat(displayData, not(hasDisplayItem("closingBehavior")));
  assertThat(displayData, not(hasDisplayItem("timestampCombiner")));
}
 

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

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  builder.addIfNotNull(DisplayData.item("topic", getTopic()).withLabel("Topic"));
  Set<String> ignoredProducerPropertiesKeys = IGNORED_PRODUCER_PROPERTIES.keySet();
  for (Map.Entry<String, Object> conf : getProducerConfig().entrySet()) {
    String key = conf.getKey();
    if (!ignoredProducerPropertiesKeys.contains(key)) {
      Object value =
          DisplayData.inferType(conf.getValue()) != null
              ? conf.getValue()
              : String.valueOf(conf.getValue());
      builder.add(DisplayData.item(key, ValueProvider.StaticValueProvider.of(value)));
    }
  }
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  if (getQuery() != null) {
    builder.add(DisplayData.item("query", getQuery()));
  }
  if (getTable() != null) {
    builder.add(DisplayData.item("table", getTable()));
  }
  builder.add(DisplayData.item("storageIntegrationName", getStagingBucketName()));
  builder.add(DisplayData.item("stagingBucketName", getStagingBucketName()));
  builder.add(DisplayData.item("csvMapper", getCsvMapper().getClass().getName()));
  builder.add(DisplayData.item("coder", getCoder().getClass().getName()));
  if (getDataSourceProviderFn() instanceof HasDisplayData) {
    ((HasDisplayData) getDataSourceProviderFn()).populateDisplayData(builder);
  }
}
 

/**
 * Create a group by key transform.
 * It returns GroupByKeyAndWindowDoFnTransform if window function is not default.
 *
 * @param ctx      translation context
 * @param beamNode the beam node to be translated
 * @return group by key transform
 */
private static Transform createGBKTransform(
  final PipelineTranslationContext ctx,
  final TransformHierarchy.Node beamNode) {
  final AppliedPTransform pTransform = beamNode.toAppliedPTransform(ctx.getPipeline());
  final PCollection<?> mainInput = (PCollection<?>)
    Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(pTransform));
  final TupleTag mainOutputTag = new TupleTag<>();

  if (isGlobalWindow(beamNode, ctx.getPipeline())) {
    return new GroupByKeyTransform();
  } else {
    return new GroupByKeyAndWindowDoFnTransform(
      getOutputCoders(pTransform),
      mainOutputTag,
      mainInput.getWindowingStrategy(),
      ctx.getPipelineOptions(),
      SystemReduceFn.buffering(mainInput.getCoder()),
      DisplayData.from(beamNode.getTransform()));
  }
}
 

/**
 * @param ctx       provides translation context
 * @param beamNode  the beam node to be translated
 * @param transform transform which can be obtained from {@code beamNode}
 */
@PrimitiveTransformTranslator({Window.class, Window.Assign.class})
private static void windowTranslator(final PipelineTranslationContext ctx,
                                     final TransformHierarchy.Node beamNode,
                                     final PTransform<?, ?> transform) {
  final WindowFn windowFn;
  if (transform instanceof Window) {
    windowFn = ((Window) transform).getWindowFn();
  } else if (transform instanceof Window.Assign) {
    windowFn = ((Window.Assign) transform).getWindowFn();
  } else {
    throw new UnsupportedOperationException(String.format("%s is not supported", transform));
  }
  final IRVertex vertex = new OperatorVertex(
    new WindowFnTransform(windowFn, DisplayData.from(beamNode.getTransform())));
  ctx.addVertex(vertex);
  beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input));
  beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output));
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  String displayBaseFilename =
      params.baseFilename.isAccessible()
          ? params.baseFilename.get().toString()
          : ("(" + params.baseFilename + ")");
  builder.add(
      DisplayData.item(
              "filenamePattern",
              String.format("%s%s%s", displayBaseFilename, params.shardTemplate, params.suffix))
          .withLabel("Filename pattern"));
  builder.add(
      DisplayData.item("filePrefix", params.baseFilename).withLabel("Output File Prefix"));
  builder.add(
      DisplayData.item("shardNameTemplate", params.shardTemplate)
          .withLabel("Output Shard Name Template"));
  builder.add(DisplayData.item("fileSuffix", params.suffix).withLabel("Output file Suffix"));
}
 

@Test
public void testDisplayData() {
  DynamicDestinations<String, Void, String> dynamicDestinations =
      DynamicFileDestinations.constant(
          DefaultFilenamePolicy.fromParams(
              new Params()
                  .withBaseFilename(
                      getBaseOutputDirectory()
                          .resolve("file", StandardResolveOptions.RESOLVE_FILE))
                  .withShardTemplate("-SS-of-NN")));
  SimpleSink<Void> sink =
      new SimpleSink<Void>(
          getBaseOutputDirectory(), dynamicDestinations, Compression.UNCOMPRESSED) {
        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
          builder.add(DisplayData.item("foo", "bar"));
        }
      };
  WriteFiles<String, ?, String> write = WriteFiles.to(sink);

  DisplayData displayData = DisplayData.from(write);

  assertThat(displayData, hasDisplayItem("sink", sink.getClass()));
  assertThat(displayData, includesDisplayDataFor("sink", sink));
}
 

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

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);

  builder
      .add(DisplayData.item("numMonths", number).withLabel("Window Months"))
      .addIfNotDefault(
          DisplayData.item("startDate", new DateTime(startDate, timeZone).toInstant())
              .withLabel("Window Start Date"),
          new DateTime(DEFAULT_START_DATE, DateTimeZone.UTC).toInstant());
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);

  Combine.populateDisplayData(builder, fn, fnDisplayData);
  if (hotKeyFanout instanceof HasDisplayData) {
    builder.include("hotKeyFanout", (HasDisplayData) hotKeyFanout);
  }
  builder.add(
      DisplayData.item("fanoutFn", hotKeyFanout.getClass()).withLabel("Fanout Function"));
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  builder
      .addIfNotNull(
          DisplayData.item("table", BigQueryHelpers.toTableSpec(readSession.getTableReference()))
              .withLabel("Table"))
      .add(DisplayData.item("readSession", readSession.getName()).withLabel("Read session"))
      .add(DisplayData.item("stream", stream.getName()).withLabel("Stream"));
}
 

/**
 * Translates an {@link AppliedPTransform} by:
 *
 * <ul>
 *   <li>adding an input to the PTransform for each {@link AppliedPTransform#getInputs()}.
 *   <li>adding an output to the PTransform for each {@link AppliedPTransform#getOutputs()}.
 *   <li>adding a PCollection for each {@link AppliedPTransform#getOutputs()}.
 *   <li>adding a reference to each subtransform.
 *   <li>set the unique name.
 *   <li>set the display data.
 * </ul>
 */
static RunnerApi.PTransform.Builder translateAppliedPTransform(
    AppliedPTransform<?, ?, ?> appliedPTransform,
    List<AppliedPTransform<?, ?, ?>> subtransforms,
    SdkComponents components)
    throws IOException {
  RunnerApi.PTransform.Builder transformBuilder = RunnerApi.PTransform.newBuilder();
  for (Map.Entry<TupleTag<?>, PValue> taggedInput : appliedPTransform.getInputs().entrySet()) {
    checkArgument(
        taggedInput.getValue() instanceof PCollection,
        "Unexpected input type %s",
        taggedInput.getValue().getClass());
    transformBuilder.putInputs(
        toProto(taggedInput.getKey()),
        components.registerPCollection((PCollection<?>) taggedInput.getValue()));
  }
  for (Map.Entry<TupleTag<?>, PValue> taggedOutput : appliedPTransform.getOutputs().entrySet()) {
    // TODO: Remove gating
    if (taggedOutput.getValue() instanceof PCollection) {
      checkArgument(
          taggedOutput.getValue() instanceof PCollection,
          "Unexpected output type %s",
          taggedOutput.getValue().getClass());
      transformBuilder.putOutputs(
          toProto(taggedOutput.getKey()),
          components.registerPCollection((PCollection<?>) taggedOutput.getValue()));
    }
  }
  for (AppliedPTransform<?, ?, ?> subtransform : subtransforms) {
    transformBuilder.addSubtransforms(components.getExistingPTransformId(subtransform));
  }

  transformBuilder.setUniqueName(appliedPTransform.getFullName());
  transformBuilder.addAllDisplayData(
      DisplayDataTranslation.toProto(DisplayData.from(appliedPTransform.getTransform())));
  return transformBuilder;
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  if (fewKeys) {
    builder.add(DisplayData.item("fewKeys", true).withLabel("Has Few Keys"));
  }
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  builder.add(DisplayData.item("class", originalFnForDisplayData.getClass()));
  if (originalFnForDisplayData instanceof HasDisplayData) {
    builder.include("fn", (HasDisplayData) originalFnForDisplayData);
  }
  builder.add(DisplayData.item("exceptionHandler.class", exceptionHandler.getClass()));
  if (exceptionHandler instanceof HasDisplayData) {
    builder.include("exceptionHandler", (HasDisplayData) exceptionHandler);
  }
}
 

@Test
public void testWritingDisplayData() {
  final String table = tmpTable.getName();
  HBaseIO.Write write = HBaseIO.write().withTableId(table).withConfiguration(conf);
  DisplayData displayData = DisplayData.from(write);
  assertThat(displayData, hasDisplayItem("tableId", table));
}
 

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

/**
 * A test that demonstrates that the replacement transform has the Display Data of the {@link
 * ParDo.SingleOutput} it replaces.
 */
@Test
public void getReplacementTransformPopulateDisplayData() {
  ParDo.SingleOutput<Integer, Long> originalTransform = ParDo.of(new ToLongFn());
  DisplayData originalDisplayData = DisplayData.from(originalTransform);
  PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3));
  AppliedPTransform<
          PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>>
      application =
          AppliedPTransform.of(
              "original",
              input.expand(),
              input.apply(originalTransform).expand(),
              originalTransform,
              pipeline);

  PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacement =
      factory.getReplacementTransform(application);
  DisplayData replacementDisplayData = DisplayData.from(replacement.getTransform());

  assertThat(replacementDisplayData, equalTo(originalDisplayData));

  DisplayData primitiveDisplayData =
      Iterables.getOnlyElement(
          DisplayDataEvaluator.create()
              .displayDataForPrimitiveTransforms(replacement.getTransform(), VarIntCoder.of()));
  assertThat(primitiveDisplayData, equalTo(replacementDisplayData));
}
 

private PerKey(
    GlobalCombineFn<? super InputT, ?, OutputT> fn,
    DisplayData.ItemSpec<? extends Class<?>> fnDisplayData,
    boolean fewKeys) {
  this.fn = fn;
  this.fnDisplayData = fnDisplayData;
  this.fewKeys = fewKeys;
  this.sideInputs = ImmutableList.of();
}
 

private static void populateGlobalDisplayData(
    DisplayData.Builder builder, int fanout, boolean insertDefault) {
  builder
      .addIfNotDefault(DisplayData.item("fanout", fanout).withLabel("Key Fanout Size"), 0)
      .add(
          DisplayData.item("emitDefaultOnEmptyInput", insertDefault)
              .withLabel("Emit Default On Empty Input"));
}
 

void populateDisplayData(DisplayData.Builder builder) {
  if (getDataSource() != null) {
    builder.addIfNotNull(DisplayData.item("dataSource", getDataSource().getClass().getName()));
  } else {
    builder.addIfNotNull(DisplayData.item("jdbcUrl", getUrl()));
    builder.addIfNotNull(DisplayData.item("username", getUsername()));
  }
}
 

@Test
public void testDirectRunnerApiSurface() throws Exception {
  // The DirectRunner can expose the Core SDK, anything exposed by the Core SDK, and itself
  @SuppressWarnings("unchecked")
  final Set<String> allowed =
      ImmutableSet.of(
          "org.apache.beam.sdk",
          "org.apache.beam.runners.direct",
          "org.joda.time",
          "javax.annotation",
          "java.math");

  final Package thisPackage = getClass().getPackage();
  final ClassLoader thisClassLoader = getClass().getClassLoader();
  ApiSurface apiSurface =
      ApiSurface.ofPackage(thisPackage, thisClassLoader)
          // Do not include dependencies that are required based on the known exposures. This
          // could alternatively prune everything exposed by the public parts of the Core SDK
          .pruningClass(Pipeline.class)
          .pruningClass(PipelineRunner.class)
          .pruningClass(PipelineOptions.class)
          .pruningClass(PipelineOptionsRegistrar.class)
          .pruningClass(PipelineOptions.DirectRunner.class)
          .pruningClass(DisplayData.Builder.class)
          .pruningClass(MetricResults.class)
          .pruningClass(DirectGraphs.class)
          .pruningClass(
              WatermarkManager.class /* TODO: BEAM-4237 Consider moving to local-java */)
          .pruningPattern(
              "org[.]apache[.]beam[.]runners[.]direct[.]portable.*"
              /* TODO: BEAM-4237 reconsider package layout with the ReferenceRunner */ )
          .pruningPattern("org[.]apache[.]beam[.].*Test.*")
          .pruningPattern("org[.]apache[.]beam[.].*IT")
          .pruningPattern("java[.]io.*")
          .pruningPattern("java[.]lang.*")
          .pruningPattern("java[.]util.*");

  assertThat(apiSurface, containsOnlyPackages(allowed));
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);

  builder.addIfNotNull(
      DisplayData.item("filePattern", getFilepattern()).withLabel("File Pattern"));
}
 

@Test
public void testParseFilesDisplayData() {
  TikaIO.ParseFiles parseFiles =
      TikaIO.parseFiles()
          .withTikaConfigPath("/tikaConfigPath")
          .withContentTypeHint("application/pdf");

  DisplayData displayData = DisplayData.from(parseFiles);

  assertThat(displayData, hasDisplayItem("tikaConfigPath", "/tikaConfigPath"));
  assertThat(displayData, hasDisplayItem("contentTypeHint", "application/pdf"));
}
 

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

  DisplayData displayData = DisplayData.from(write);

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