org.apache.beam.sdk.runners.TransformHierarchy Java Examples

@Override
  public void visitPrimitiveTransform(final TransformHierarchy.Node beamNode) {
//    Print if needed for development
//    LOG.info("visitp " + beamNode.getTransform());
    if (beamNode.getOutputs().size() > 1) {
      throw new UnsupportedOperationException(beamNode.toString());
    }

    final IRVertex irVertex = convertToVertex(beamNode, builder, pValueToVertex, pValueToCoder, options,
        loopVertexStack);
    beamNode.getOutputs().values().stream().filter(v -> v instanceof PCollection).map(v -> (PCollection) v)
        .forEach(output -> pValueToCoder.put(output, new BeamCoder(output.getCoder())));

    beamNode.getOutputs().values().forEach(output -> pValueToVertex.put(output, irVertex));

    beamNode.getInputs().values().stream().filter(pValueToVertex::containsKey)
        .forEach(pValue -> {
          final IRVertex src = pValueToVertex.get(pValue);
          final BeamCoder coder = pValueToCoder.get(pValue);
          final IREdge edge = new IREdge(getEdgeCommunicationPattern(src, irVertex), src, irVertex, coder);
          edge.setProperty(KeyExtractorProperty.of(new BeamKeyExtractor()));
          this.builder.connectVertices(edge);
        });
  }
 

@Override
public void visitValue(PValue value, TransformHierarchy.Node producer) {
  LOG.debug("Checking translation of {}", value);
  // Primitive transforms are the only ones assigned step names.
  if (producer.getTransform() instanceof CreateDataflowView
      && !hasExperiment(options, "beam_fn_api")) {
    // CreateDataflowView produces a dummy output (as it must be a primitive transform)
    // but in the Dataflow Job graph produces only the view and not the output PCollection.
    asOutputReference(
        ((CreateDataflowView) producer.getTransform()).getView(),
        producer.toAppliedPTransform(getPipeline()));
    return;
  } else if (producer.getTransform() instanceof View.CreatePCollectionView
      && hasExperiment(options, "beam_fn_api")) {
    // View.CreatePCollectionView produces a dummy output (as it must be a primitive transform)
    // but in the Dataflow Job graph produces only the view and not the output PCollection.
    asOutputReference(
        ((View.CreatePCollectionView) producer.getTransform()).getView(),
        producer.toAppliedPTransform(getPipeline()));
    return;
  }
  asOutputReference(value, producer.toAppliedPTransform(getPipeline()));
}
 

private void verifyPipelineExecution() {
  if (!isEmptyPipeline(pipeline)) {
    if (!runAttempted && !enableAutoRunIfMissing) {
      throw new PipelineRunMissingException("The pipeline has not been run.");

    } else {
      final List<TransformHierarchy.Node> pipelineNodes = recordPipelineNodes(pipeline);
      if (pipelineRunSucceeded() && !visitedAll(pipelineNodes)) {
        final boolean hasDanglingPAssert =
            FluentIterable.from(pipelineNodes)
                .filter(Predicates.not(Predicates.in(runVisitedNodes)))
                .anyMatch(isPAssertNode);
        if (hasDanglingPAssert) {
          throw new AbandonedNodeException("The pipeline contains abandoned PAssert(s).");
        } else {
          throw new AbandonedNodeException("The pipeline contains abandoned PTransform(s).");
        }
      }
    }
  }
}
 

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

@Override
public void leaveCompositeTransform(TransformHierarchy.Node node) {
  checkState(
      !finalized,
      "Attempting to traverse a pipeline (node %s) with a %s which is already finalized",
      node.getFullName(),
      getClass().getSimpleName());
  if (node.isRootNode()) {
    finalized = true;
    checkState(
        viewWriters.keySet().containsAll(consumedViews),
        "All %ss that are consumed must be written by some %s %s: Missing %s",
        PCollectionView.class.getSimpleName(),
        WriteView.class.getSimpleName(),
        PTransform.class.getSimpleName(),
        Sets.difference(consumedViews, viewWriters.keySet()));
  }
}
 

/** Traverse the pipeline and return the first {@link Combine.GroupedValues} found. */
private static AppliedPTransform<?, ?, ?> getCombineGroupedValuesFrom(TestPipeline pipeline) {
  final AppliedPTransform<?, ?, ?>[] transform = new AppliedPTransform<?, ?, ?>[1];
  pipeline.traverseTopologically(
      new Pipeline.PipelineVisitor.Defaults() {
        @Override
        public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
          if (!node.isRootNode()
              && node.toAppliedPTransform(getPipeline())
                  .getTransform()
                  .getClass()
                  .equals(Combine.GroupedValues.class)) {
            transform[0] = node.toAppliedPTransform(getPipeline());
            return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
          }
          return CompositeBehavior.ENTER_TRANSFORM;
        }
      });
  return transform[0];
}
 

/**
 * With {@link #testWindowIntoWindowFnAssign()}, demonstrates that the expansions of the {@link
 * Window} transform depends on if it actually assigns elements to windows.
 */
@Test
public void testWindowIntoNullWindowFnNoAssign() {
  pipeline
      .apply(Create.of(1, 2, 3))
      .apply(
          Window.<Integer>configure()
              .triggering(AfterWatermark.pastEndOfWindow())
              .withAllowedLateness(Duration.ZERO)
              .accumulatingFiredPanes());

  pipeline.traverseTopologically(
      new PipelineVisitor.Defaults() {
        @Override
        public void visitPrimitiveTransform(TransformHierarchy.Node node) {
          assertThat(node.getTransform(), not(instanceOf(Window.Assign.class)));
        }
      });
}
 

/**
 * @param ctx       provides translation context
 * @param beamNode  the beam node to be translated
 * @param transform transform which can be obtained from {@code beamNode}
 */
@PrimitiveTransformTranslator(ParDo.MultiOutput.class)
private static void parDoMultiOutputTranslator(final PipelineTranslationContext ctx,
                                               final TransformHierarchy.Node beamNode,
                                               final ParDo.MultiOutput<?, ?> transform) {
  final Map<Integer, PCollectionView<?>> sideInputMap = getSideInputMap(transform.getSideInputs().values());
  final AbstractDoFnTransform doFnTransform = createDoFnTransform(ctx, beamNode, sideInputMap);
  final IRVertex vertex = new OperatorVertex(doFnTransform);
  ctx.addVertex(vertex);
  beamNode.getInputs().values().stream()
    .filter(input -> !transform.getAdditionalInputs().values().contains(input))
    .forEach(input -> ctx.addEdgeTo(vertex, input));
  ctx.addSideInputEdges(vertex, sideInputMap);
  beamNode.getOutputs().entrySet().stream()
    .filter(pValueWithTupleTag -> pValueWithTupleTag.getKey().equals(transform.getMainOutputTag()))
    .forEach(pValueWithTupleTag -> ctx.registerMainOutputFrom(beamNode, vertex, pValueWithTupleTag.getValue()));
  beamNode.getOutputs().entrySet().stream()
    .filter(pValueWithTupleTag -> !pValueWithTupleTag.getKey().equals(transform.getMainOutputTag()))
    .forEach(pValueWithTupleTag -> ctx.registerAdditionalOutputFrom(beamNode, vertex, pValueWithTupleTag.getValue(),
      pValueWithTupleTag.getKey()));
}
 

/**
 * @param ctx       provides translation context
 * @param beamNode  the beam node to be translated
 * @param transform transform which can be obtained from {@code beamNode}
 */
@PrimitiveTransformTranslator(ParDo.SingleOutput.class)
private static void parDoSingleOutputTranslator(final PipelineTranslationContext ctx,
                                                final TransformHierarchy.Node beamNode,
                                                final ParDo.SingleOutput<?, ?> transform) {
  final Map<Integer, PCollectionView<?>> sideInputMap = getSideInputMap(transform.getSideInputs().values());
  final AbstractDoFnTransform doFnTransform = createDoFnTransform(ctx, beamNode, sideInputMap);
  final IRVertex vertex = new OperatorVertex(doFnTransform);

  ctx.addVertex(vertex);
  beamNode.getInputs().values().stream()
    .filter(input -> !transform.getAdditionalInputs().values().contains(input))
    .forEach(input -> ctx.addEdgeTo(vertex, input));
  ctx.addSideInputEdges(vertex, sideInputMap);
  beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output));
}
 

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  if (finalized) {
    throw new IllegalStateException("Attempting to traverse an already finalized pipeline!");
  }

  PTransform<?, ?> transform = node.getTransform();
  if (transform != null) {
    JetTransformTranslator<?> translator = translatorProvider.apply(transform);
    if (translator != null) {
      translate(node, translator);
      return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
    }
  }
  return CompositeBehavior.ENTER_TRANSFORM;
}
 

@Override
public void visitPrimitiveTransform(TransformHierarchy.Node node) {
  final int nodeId = nextNodeId++;
  writeLine("%d [label=\"%s\"]", nodeId, escapeString(node.getTransform().getName()));

  node.getOutputs().values().forEach(x -> valueToProducerNodeId.put(x, nodeId));

  node.getInputs()
      .forEach(
          (key, value) -> {
            final int producerId = valueToProducerNodeId.get(value);
            String style = "solid";
            if (node.getTransform().getAdditionalInputs().containsKey(key)) {
              style = "dashed";
            }
            writeLine("%d -> %d [style=%s label=\"%s\"]", producerId, nodeId, style, "");
          });
}
 

@Override
public void translate(
    SamzaPublishView<ElemT, ViewT> transform,
    TransformHierarchy.Node node,
    TranslationContext ctx) {
  doTranslate(transform, node, ctx);
}
 

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

/** Returns a translator for the given node, if it is possible, otherwise null. */
@Override
protected TransformTranslator<?> getTransformTranslator(TransformHierarchy.Node node) {
  @Nullable PTransform<?, ?> transform = node.getTransform();
  // Root of the graph is null
  if (transform == null) {
    return null;
  }
  return TRANSFORM_TRANSLATORS.get(transform.getClass());
}
 

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  if (!node.isRootNode()) {
    parents.addFirst(node);
  }
  return CompositeBehavior.ENTER_TRANSFORM;
}
 

@Override
public void visitValue(PValue value, TransformHierarchy.Node producer) {
  final String valueDesc = value.toString().replaceFirst(".*:([a-zA-Z#0-9]+).*", "$1");

  final String samzaSafeValueDesc = valueDesc.replaceAll("[^A-Za-z0-9_-]", "_");

  idMap.put(value, String.format("%d-%s", nextId++, samzaSafeValueDesc));
}
 

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  CompositeBehavior compositeBehavior = super.enterCompositeTransform(node);
  PTransform<?, ?> transform = node.getTransform();
  if (transform != null) {
    @SuppressWarnings("unchecked")
    final Class<PTransform<?, ?>> transformClass = (Class<PTransform<?, ?>>) transform.getClass();
    if (compositeBehavior == CompositeBehavior.ENTER_TRANSFORM
        && !knownComposite(transformClass)
        && shouldDebug(node)) {
      transforms.add(new NativeTransform(node, null, transform, true));
    }
  }
  return compositeBehavior;
}
 

/** Apply the given TransformTranslator to the given node. */
private <T extends PTransform<?, ?>> void applyTransformTranslator(
    TransformHierarchy.Node node, TransformTranslator<?> transformTranslator) {
  // create the applied PTransform on the translationContext
  translationContext.setCurrentTransform(node.toAppliedPTransform(getPipeline()));

  // avoid type capture
  @SuppressWarnings("unchecked")
  T typedTransform = (T) node.getTransform();
  @SuppressWarnings("unchecked")
  TransformTranslator<T> typedTransformTranslator = (TransformTranslator<T>) transformTranslator;

  // apply the transformTranslator
  typedTransformTranslator.translateTransform(typedTransform, translationContext);
}
 

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

private static <K, InputT, OutputT> void doTranslate(
    PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, OutputT>>> transform,
    TransformHierarchy.Node node,
    TranslationContext ctx) {
  final PCollection<KV<K, InputT>> input = ctx.getInput(transform);

  final PCollection<KV<K, OutputT>> output = ctx.getOutput(transform);
  final TupleTag<KV<K, OutputT>> outputTag = ctx.getOutputTag(transform);

  @SuppressWarnings("unchecked")
  final WindowingStrategy<?, BoundedWindow> windowingStrategy =
      (WindowingStrategy<?, BoundedWindow>) input.getWindowingStrategy();

  final MessageStream<OpMessage<KV<K, InputT>>> inputStream = ctx.getMessageStream(input);

  final KvCoder<K, InputT> kvInputCoder = (KvCoder<K, InputT>) input.getCoder();
  final Coder<WindowedValue<KV<K, InputT>>> elementCoder = SamzaCoders.of(input);

  final SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> reduceFn =
      getSystemReduceFn(transform, input.getPipeline(), kvInputCoder);

  final MessageStream<OpMessage<KV<K, OutputT>>> outputStream =
      doTranslateGBK(
          inputStream,
          needRepartition(node, ctx),
          reduceFn,
          windowingStrategy,
          kvInputCoder,
          elementCoder,
          ctx.getTransformFullName(),
          ctx.getTransformId(),
          outputTag,
          input.isBounded());

  ctx.registerMessageStream(output, outputStream);
}
 

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

@Override
public Vertex translate(
    Pipeline pipeline,
    AppliedPTransform<?, ?, ?> appliedTransform,
    TransformHierarchy.Node node,
    JetTranslationContext context) {
  String transformName = appliedTransform.getFullName();
  DAGBuilder dagBuilder = context.getDagBuilder();
  String vertexId = dagBuilder.newVertexId(transformName);

  TestStream<T> testStream = (TestStream<T>) appliedTransform.getTransform();

  // events in the transform are not serializable, we have to translate them. We'll also flatten
  // the collection.
  Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
  Coder outputCoder = Utils.getCoder((PCollection) output.getValue());
  TestStream.TestStreamCoder<T> payloadCoder =
      TestStream.TestStreamCoder.of(testStream.getValueCoder());
  byte[] encodedPayload = getEncodedPayload(testStream, payloadCoder);
  Vertex vertex =
      dagBuilder.addVertex(
          vertexId, TestStreamP.supplier(encodedPayload, payloadCoder, outputCoder));

  String outputEdgeId = Utils.getTupleTagId(output.getValue());
  dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
  dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
  return vertex;
}
 

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

  @SuppressWarnings("unchecked")
  final TransformTranslator<T> typedTranslator = (TransformTranslator<T>) translator;
  visitorFn.apply(transform, node, getPipeline(), typedTranslator);
}
 

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  checkState(
      !finalized,
      "Attempted to use a %s that has already been finalized on a pipeline (visiting node %s)",
      KeyedPValueTrackingVisitor.class.getSimpleName(),
      node);
  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) {
  if (node.getTransform() instanceof SamzaPublishView) {
    final PCollectionView view = ((SamzaPublishView) node.getTransform()).getView();
    visitValue(view, node);
  }
}
 

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

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  if ((node.getTransform()) != null) {
    final AppliedPTransform<?, ?, ?> applied = node.toAppliedPTransform(getPipeline());
    lookupTable.put(applied.getTransform(), applied);
  }
  return CompositeBehavior.ENTER_TRANSFORM;
}
 

@Override
public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
  final PTransform<?, ?> transform = node.getTransform();
  final String urn = getUrnForTransform(transform);
  if (canTranslate(urn, transform)) {
    applyTransform(transform, node, TRANSLATORS.get(urn));
    return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
  }
  return CompositeBehavior.ENTER_TRANSFORM;
}
 

@Override
public Map<String, String> createConfig(
    ParDo.MultiOutput<InT, OutT> transform, TransformHierarchy.Node node, ConfigContext ctx) {
  final Map<String, String> config = new HashMap<>();
  final DoFnSignature signature = DoFnSignatures.getSignature(transform.getFn().getClass());
  final SamzaPipelineOptions options = ctx.getPipelineOptions();

  if (signature.usesState()) {
    // set up user state configs
    for (DoFnSignature.StateDeclaration state : signature.stateDeclarations().values()) {
      final String storeId = state.id();
      config.put(
          "stores." + storeId + ".factory",
          "org.apache.samza.storage.kv.RocksDbKeyValueStorageEngineFactory");
      config.put("stores." + storeId + ".key.serde", "byteArraySerde");
      config.put("stores." + storeId + ".msg.serde", "byteSerde");

      if (options.getStateDurable()) {
        config.put(
            "stores." + storeId + ".changelog",
            ConfigBuilder.getChangelogTopic(options, storeId));
      }
    }
  }

  if (doFnInvokerRegistrar != null) {
    config.putAll(doFnInvokerRegistrar.configFor(transform.getFn()));
  }

  return config;
}