Java Code Examples for org.apache.beam.sdk.runners.TransformHierarchy#Node

@Override
public void visitPrimitiveTransform(TransformHierarchy.Node node) {
  LOG.info("{} visitPrimitiveTransform- {}", genSpaces(this.depth), node.getFullName());
  // get the transformation corresponding to hte node we are
  // currently visiting and translate it into its Flink alternative.

  PTransform<?, ?> transform = node.getTransform();
  StreamTransformTranslator<?> translator =
      FlinkStreamingTransformTranslators.getTranslator(transform);

  if (translator == null || !applyCanTranslate(transform, node, translator)) {
    String transformUrn = PTransformTranslation.urnForTransform(transform);
    LOG.info(transformUrn);
    throw new UnsupportedOperationException(
        "The transform " + transformUrn + " is currently not supported.");
  }
  applyStreamingTransform(transform, node, translator);
}
 

@Override
public void visitPrimitiveTransform(TransformHierarchy.Node node) {
  AppliedPTransform<?, ?, ?> appliedPTransform = node.toAppliedPTransform(getPipeline());
  if (hasUnboundedOutput(appliedPTransform)) {
    Class<? extends PTransform> transformClass = node.getTransform().getClass();
    LOG.debug("Found unbounded PCollection for transform %s", transformClass);
    hasUnboundedCollections = true;
  }
}
 

@Override
public void visitValue(PValue value, TransformHierarchy.Node producer) {
  if (translationMode.equals(TranslationMode.BATCH)) {
    if (value instanceof PCollection
        && ((PCollection) value).isBounded() == PCollection.IsBounded.UNBOUNDED) {
      LOG.info(
          "Found unbounded PCollection {}. Switching to streaming execution.", value.getName());
      translationMode = TranslationMode.STREAMING;
    }
  }
}
 

/**
 * @param ctx       provides translation context
 * @param beamNode  the beam node to be translated
 * @param transform transform which can be obtained from {@code beamNode}
 * @return behavior controls whether or not child transforms are visited.
 */
@CompositeTransformTranslator(LoopCompositeTransform.class)
private static Pipeline.PipelineVisitor.CompositeBehavior loopTranslator(
  final PipelineTranslationContext ctx,
  final TransformHierarchy.Node beamNode,
  final LoopCompositeTransform<?, ?> transform) {
  // Do nothing here, as the context handles the loop vertex stack.
  // We just keep this method to signal that the loop vertex is acknowledged.
  return Pipeline.PipelineVisitor.CompositeBehavior.ENTER_TRANSFORM;
}
 

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

private <T extends PTransform<?, ?>> void applyBatchTransform(
    PTransform<?, ?> transform,
    TransformHierarchy.Node node,
    BatchTransformTranslator<?> translator) {

  @SuppressWarnings("unchecked")
  T typedTransform = (T) transform;

  @SuppressWarnings("unchecked")
  BatchTransformTranslator<T> typedTranslator = (BatchTransformTranslator<T>) translator;

  // create the applied PTransform on the batchContext
  batchContext.setCurrentTransform(node.toAppliedPTransform(getPipeline()));
  typedTranslator.translateNode(typedTransform, batchContext);
}
 

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  LOG.debug("{} enterCompositeTransform- {}", genSpaces(depth), node.getFullName());
  depth++;

  TransformTranslator<?> transformTranslator = getTransformTranslator(node);

  if (transformTranslator != null) {
    applyTransformTranslator(node, transformTranslator);
    LOG.debug("{} translated- {}", genSpaces(depth), node.getFullName());
    return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
  } else {
    return CompositeBehavior.ENTER_TRANSFORM;
  }
}
 

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  PTransform<?, ?> transform = node.getTransform();
  if (transform != null) {
    if (translator.hasTranslation(transform) && !shouldDefer(node)) {
      LOG.info("Entering directly-translatable composite transform: '{}'", node.getFullName());
      LOG.debug("Composite transform class: '{}'", transform);
      doVisitTransform(node);
      return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
    }
  }
  return CompositeBehavior.ENTER_TRANSFORM;
}
 

/** Returns a translator for the given node, if it is possible, otherwise null. */
private static BatchTransformTranslator<?> getTranslator(
    TransformHierarchy.Node node, FlinkBatchTranslationContext context) {
  @Nullable PTransform<?, ?> transform = node.getTransform();

  // Root of the graph is null
  if (transform == null) {
    return null;
  }

  return FlinkBatchTransformTranslators.getTranslator(transform, context);
}
 

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

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  if (Objects.equals(root, node.getTransform())) {
    inCompositeRoot = true;
  }

  return CompositeBehavior.ENTER_TRANSFORM;
}
 

@Override
public void visitPrimitiveTransform(TransformHierarchy.Node node) {
  PTransform<?, ?> transform = node.getTransform();
  JetTransformTranslator<?> translator = translatorProvider.apply(transform);
  if (translator == null) {
    String transformUrn = PTransformTranslation.urnForTransform(transform);
    throw new UnsupportedOperationException(
        "The transform " + transformUrn + " is currently not supported.");
  }
  translate(node, translator);
}
 

@Override
public void visitPrimitiveTransform(TransformHierarchy.Node node) {
  for (PValue input : node.getInputs().values()) {
    numConsumers.merge(input, 1, Integer::sum);
  }
}
 

<T extends PTransform<?, ?>> void apply(
T transform,
TransformHierarchy.Node node,
Pipeline pipeline,
TransformTranslator<T> translator);
 

static Collection<PValue> getMainInputs(Pipeline pipeline, TransformHierarchy.Node node) {
  if (node.getTransform() == null) {
    return null;
  }
  return TransformInputs.nonAdditionalInputs(node.toAppliedPTransform(pipeline));
}
 

@Override
public void leaveCompositeTransform(final TransformHierarchy.Node node) {
  visited.add(node);
}
 

@Override
public void leaveCompositeTransform(TransformHierarchy.Node node) {
  this.depth--;
  LOG.info("{} leaveCompositeTransform- {}", genSpaces(this.depth), node.getFullName());
}
 

@Override
public void leaveCompositeTransform(TransformHierarchy.Node node) {
  if (Objects.equals(root, node.getTransform())) {
    inCompositeRoot = false;
  }
}
 

/**
 * @param pValue {@link PValue}
 * @return the producer beam node.
 */
TransformHierarchy.Node getProducerBeamNodeOf(final PValue pValue) {
  return pValueToProducerBeamNode.get(pValue);
}
 

/**
 * Called for each composite transform after all topological predecessors have been visited but
 * before any of its component transforms.
 *
 * <p>The return value controls whether or not child transforms are visited.
 */
CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node);