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

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

@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.add(DisplayData.item("class", originalFnForDisplayData.getClass()));
  if (originalFnForDisplayData instanceof HasDisplayData) {
    builder.include("fn", (HasDisplayData) originalFnForDisplayData);
  }
}
 

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

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

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

/**
 * Populate display data for the {@code combineFns} that make up a composed combine transform.
 *
 * <p>The same combineFn class may be used multiple times, in which case we must take special care
 * to register display data with unique namespaces.
 */
private static void populateDisplayData(
    DisplayData.Builder builder, List<? extends HasDisplayData> combineFns) {
  for (int i = 0; i < combineFns.size(); i++) {
    HasDisplayData combineFn = combineFns.get(i);
    String token = "combineFn" + (i + 1);
    builder.add(DisplayData.item(token, combineFn.getClass()).withLabel("Combine Function"));
    builder.include(token, combineFn);
  }
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  builder.add(DisplayData.item("query", getQuery()));
  if (getDataSourceProviderFn() instanceof HasDisplayData) {
    ((HasDisplayData) getDataSourceProviderFn()).populateDisplayData(builder);
  }
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  builder.add(DisplayData.item("query", getQuery()));
  builder.add(DisplayData.item("rowMapper", getRowMapper().getClass().getName()));
  builder.add(DisplayData.item("coder", getCoder().getClass().getName()));
  if (getDataSourceProviderFn() instanceof HasDisplayData) {
    ((HasDisplayData) getDataSourceProviderFn()).populateDisplayData(builder);
  }
}
 

@Override
public void populateDisplayData(DisplayData.Builder builder) {
  super.populateDisplayData(builder);
  builder.add(DisplayData.item("query", getQuery()));
  builder.add(DisplayData.item("rowMapper", getRowMapper().getClass().getName()));
  builder.add(DisplayData.item("coder", getCoder().getClass().getName()));
  if (getDataSourceProviderFn() instanceof HasDisplayData) {
    ((HasDisplayData) getDataSourceProviderFn()).populateDisplayData(builder);
  }
}
 

@Test
public void testTranslation() {
  assertThat(
      DisplayDataTranslation.toProto(
          DisplayData.from(
              new HasDisplayData() {

                @Override
                public void populateDisplayData(DisplayData.Builder builder) {
                  builder.add(DisplayData.item("foo", "value"));
                  builder.add(DisplayData.item("foo2", "value2").withLabel("label"));
                }
              })),
      containsInAnyOrder(
          RunnerApi.DisplayData.newBuilder()
              .setUrn(DisplayDataTranslation.LABELLED_STRING)
              .setPayload(
                  RunnerApi.LabelledStringPayload.newBuilder()
                      .setLabel("foo")
                      .setValue("value")
                      .build()
                      .toByteString())
              .build(),
          RunnerApi.DisplayData.newBuilder()
              .setUrn(DisplayDataTranslation.LABELLED_STRING)
              .setPayload(
                  RunnerApi.LabelledStringPayload.newBuilder()
                      .setLabel("label")
                      .setValue("value2")
                      .build()
                      .toByteString())
              .build()));
}
 

/**
 * @param ctx          provides translation context.
 * @param beamNode     the beam node to be translated.
 * @param sideInputMap side inputs.
 * @return the created DoFnTransform.
 */
private static AbstractDoFnTransform createDoFnTransform(final PipelineTranslationContext ctx,
                                                         final TransformHierarchy.Node beamNode,
                                                         final Map<Integer, PCollectionView<?>> sideInputMap) {
  try {
    final AppliedPTransform pTransform = beamNode.toAppliedPTransform(ctx.getPipeline());
    final DoFn doFn = ParDoTranslation.getDoFn(pTransform);
    final TupleTag mainOutputTag = ParDoTranslation.getMainOutputTag(pTransform);
    final TupleTagList additionalOutputTags = ParDoTranslation.getAdditionalOutputTags(pTransform);

    final PCollection<?> mainInput = (PCollection<?>)
      Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(pTransform));

    final HasDisplayData displayData = (builder) -> builder.add(DisplayData.item("name", beamNode.getFullName()));
    final DoFnSchemaInformation doFnSchemaInformation =
      ParDoTranslation.getSchemaInformation(beamNode.toAppliedPTransform(ctx.getPipeline()));

    if (sideInputMap.isEmpty()) {
      return new DoFnTransform(
        doFn,
        mainInput.getCoder(),
        getOutputCoders(pTransform),
        mainOutputTag,
        additionalOutputTags.getAll(),
        mainInput.getWindowingStrategy(),
        ctx.getPipelineOptions(),
        DisplayData.from(displayData),
        doFnSchemaInformation,
        Collections.emptyMap());
    } else {
      return new PushBackDoFnTransform(
        doFn,
        mainInput.getCoder(),
        getOutputCoders(pTransform),
        mainOutputTag,
        additionalOutputTags.getAll(),
        mainInput.getWindowingStrategy(),
        sideInputMap,
        ctx.getPipelineOptions(),
        DisplayData.from(displayData),
        doFnSchemaInformation,
        Collections.emptyMap());
    }
  } catch (final IOException e) {
    throw new RuntimeException(e);
  }
}
 

private static void populateDisplayData(
    DisplayData.Builder builder,
    HasDisplayData fn,
    DisplayData.ItemSpec<? extends Class<?>> fnDisplayItem) {
  builder.include("combineFn", fn).add(fnDisplayItem);
}
 

private void addDisplayData(Step step, String stepName, HasDisplayData hasDisplayData) {
  DisplayData displayData = DisplayData.from(hasDisplayData);
  List<Map<String, Object>> list = MAPPER.convertValue(displayData, List.class);
  addList(getProperties(), PropertyNames.DISPLAY_DATA, list);
}
 

private static void evaluateDisplayData(HasDisplayData component) {
  DisplayData.from(component);
}