Java Code Examples for org.apache.beam.runners.core.construction.graph.PipelineNode#PTransformNode

private void visitTransform(PipelineNode.PTransformNode node) {
  final int nodeId = nextNodeId++;
  final RunnerApi.PTransform transform = node.getTransform();
  writeLine(
      "%d [label=\"%s\\n%s\"]",
      nodeId,
      escapeString(transform.getUniqueName()),
      escapeString(transform.getSpec().getUrn()));

  transform.getOutputsMap().values().forEach(x -> valueToProducerNodeId.put(x, nodeId));

  transform
      .getInputsMap()
      .forEach(
          (key, value) -> {
            final int producerId = valueToProducerNodeId.get(value);
            String style = "solid";
            writeLine(
                "%d -> %d [style=%s label=\"%s\"]",
                producerId,
                nodeId,
                style,
                escapeString(value.substring(value.lastIndexOf('_') + 1)));
          });
}
 

@Override
public FlinkPortablePipelineTranslator.Executor translate(
    BatchTranslationContext context, RunnerApi.Pipeline pipeline) {
  // Use a QueryablePipeline to traverse transforms topologically.
  QueryablePipeline p =
      QueryablePipeline.forTransforms(
          pipeline.getRootTransformIdsList(), pipeline.getComponents());
  for (PipelineNode.PTransformNode transform : p.getTopologicallyOrderedTransforms()) {
    urnToTransformTranslator
        .getOrDefault(
            transform.getTransform().getSpec().getUrn(),
            FlinkBatchPortablePipelineTranslator::urnNotFound)
        .translate(transform, pipeline, context);
  }

  // Ensure that side effects are performed for unconsumed DataSets.
  for (DataSet<?> dataSet : context.getDanglingDataSets()) {
    dataSet.output(new DiscardingOutputFormat<>()).name("DiscardingOutput");
  }

  return context;
}
 

@Override
public void translatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {

  final String outputId = ctx.getOutputId(transform);
  final GenericSystemDescriptor systemDescriptor =
      new GenericSystemDescriptor(outputId, SamzaImpulseSystemFactory.class.getName());

  // The KvCoder is needed here for Samza not to crop the key.
  final Serde<KV<?, OpMessage<byte[]>>> kvSerde = KVSerde.of(new NoOpSerde(), new NoOpSerde<>());
  final GenericInputDescriptor<KV<?, OpMessage<byte[]>>> inputDescriptor =
      systemDescriptor.getInputDescriptor(outputId, kvSerde);

  ctx.registerInputMessageStream(outputId, inputDescriptor);
}
 

@Override
public void translatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  final RunnerApi.WindowIntoPayload payload;
  try {
    payload =
        RunnerApi.WindowIntoPayload.parseFrom(transform.getTransform().getSpec().getPayload());
  } catch (InvalidProtocolBufferException e) {
    throw new IllegalArgumentException(
        String.format("failed to parse WindowIntoPayload: %s", transform.getId()), e);
  }

  @SuppressWarnings("unchecked")
  final WindowFn<T, ?> windowFn =
      (WindowFn<T, ?>) WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn());

  final MessageStream<OpMessage<T>> inputStream = ctx.getOneInputMessageStream(transform);

  final MessageStream<OpMessage<T>> outputStream =
      inputStream.flatMap(OpAdapter.adapt(new WindowAssignOp<>(windowFn)));

  ctx.registerMessageStream(ctx.getOutputId(transform), outputStream);
}
 

private String toDot(RunnerApi.Pipeline pipeline) {
  final QueryablePipeline p =
      QueryablePipeline.forTransforms(
          pipeline.getRootTransformIdsList(), pipeline.getComponents());

  begin();

  for (PipelineNode.PTransformNode transform : p.getTopologicallyOrderedTransforms()) {
    visitTransform(transform);
  }

  end();

  return dotBuilder.toString();
}
 

@Override
public FlinkPortablePipelineTranslator.Executor translate(
    StreamingTranslationContext context, RunnerApi.Pipeline pipeline) {
  QueryablePipeline p =
      QueryablePipeline.forTransforms(
          pipeline.getRootTransformIdsList(), pipeline.getComponents());
  for (PipelineNode.PTransformNode transform : p.getTopologicallyOrderedTransforms()) {
    urnToTransformTranslator
        .getOrDefault(transform.getTransform().getSpec().getUrn(), this::urnNotFound)
        .translate(transform.getId(), pipeline, context);
  }

  return context;
}
 

@Override
public void translatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  doTranslatePortable(transform, pipeline, ctx);
}
 

public static WindowingStrategy<?, BoundedWindow> getPortableWindowStrategy(
    PipelineNode.PTransformNode transform, QueryablePipeline pipeline) {
  String inputId = Iterables.getOnlyElement(transform.getTransform().getInputsMap().values());
  RehydratedComponents rehydratedComponents =
      RehydratedComponents.forComponents(pipeline.getComponents());

  RunnerApi.WindowingStrategy windowingStrategyProto =
      pipeline
          .getComponents()
          .getWindowingStrategiesOrThrow(
              pipeline.getComponents().getPcollectionsOrThrow(inputId).getWindowingStrategyId());

  WindowingStrategy<?, ?> windowingStrategy;
  try {
    windowingStrategy =
        WindowingStrategyTranslation.fromProto(windowingStrategyProto, rehydratedComponents);
  } catch (InvalidProtocolBufferException e) {
    throw new IllegalStateException(
        String.format(
            "Unable to hydrate GroupByKey windowing strategy %s.", windowingStrategyProto),
        e);
  }

  @SuppressWarnings("unchecked")
  WindowingStrategy<?, BoundedWindow> ret =
      (WindowingStrategy<?, BoundedWindow>) windowingStrategy;
  return ret;
}
 

private static <T> void doTranslatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  final List<MessageStream<OpMessage<T>>> inputStreams = ctx.getAllInputMessageStreams(transform);
  final String outputId = ctx.getOutputId(transform);

  // For portable api there should be at least the impulse as a dummy input
  // We will know once validateRunner tests are available for portable runners
  checkState(
      !inputStreams.isEmpty(), "no input streams defined for Flatten: %s", transform.getId());

  ctx.registerMessageStream(outputId, mergeInputStreams(inputStreams));
}
 

@Override
public void translatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  doTranslatePortable(transform, pipeline, ctx);
}
 

public WindowingStrategy<?, BoundedWindow> getPortableWindowStrategy(
    PipelineNode.PTransformNode transform, QueryablePipeline pipeline) {
  String inputId = Iterables.getOnlyElement(transform.getTransform().getInputsMap().values());
  RehydratedComponents rehydratedComponents =
      RehydratedComponents.forComponents(pipeline.getComponents());

  RunnerApi.WindowingStrategy windowingStrategyProto =
      pipeline
          .getComponents()
          .getWindowingStrategiesOrThrow(
              pipeline.getComponents().getPcollectionsOrThrow(inputId).getWindowingStrategyId());

  WindowingStrategy<?, ?> windowingStrategy;
  try {
    windowingStrategy =
        WindowingStrategyTranslation.fromProto(windowingStrategyProto, rehydratedComponents);
  } catch (Exception e) {
    throw new IllegalStateException(
        String.format(
            "Unable to hydrate GroupByKey windowing strategy %s.", windowingStrategyProto),
        e);
  }

  @SuppressWarnings("unchecked")
  WindowingStrategy<?, BoundedWindow> ret =
      (WindowingStrategy<?, BoundedWindow>) windowingStrategy;
  return ret;
}
 

@Override
public void translatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  doTranslatePortable(transform, pipeline, ctx);
}
 

public void setCurrentTransform(PipelineNode.PTransformNode currentTransform) {
  this.currentTransform = currentTransform;
}
 

public <T> MessageStream<OpMessage<T>> getOneInputMessageStream(
    PipelineNode.PTransformNode transform) {
  String id = Iterables.getOnlyElement(transform.getTransform().getInputsMap().values());
  return getMessageStreamById(id);
}
 

public <T> List<MessageStream<OpMessage<T>>> getAllInputMessageStreams(
    PipelineNode.PTransformNode transform) {
  final Collection<String> inputStreamIds = transform.getTransform().getInputsMap().values();
  return inputStreamIds.stream().map(this::<T>getMessageStreamById).collect(Collectors.toList());
}
 

public String getOutputId(PipelineNode.PTransformNode transform) {
  return Iterables.getOnlyElement(transform.getTransform().getOutputsMap().values());
}
 

public String getInputId(PipelineNode.PTransformNode transform) {
  return Iterables.getOnlyElement(transform.getTransform().getInputsMap().values());
}
 

/**
 * Translates transformNode from Beam into the Spark context.
 */
void translate(
    PipelineNode.PTransformNode transformNode,
    RunnerApi.Pipeline pipeline,
    Twister2BatchTranslationContext context);
 

private static <InT, OutT> void doTranslatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  Map<String, String> outputs = transform.getTransform().getOutputsMap();

  final RunnerApi.ExecutableStagePayload stagePayload;
  try {
    stagePayload =
        RunnerApi.ExecutableStagePayload.parseFrom(
            transform.getTransform().getSpec().getPayload());
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  String inputId = stagePayload.getInput();
  final MessageStream<OpMessage<InT>> inputStream = ctx.getMessageStreamById(inputId);
  // TODO: support side input
  final List<MessageStream<OpMessage<InT>>> sideInputStreams = Collections.emptyList();

  final Map<TupleTag<?>, Integer> tagToIndexMap = new HashMap<>();
  final Map<String, TupleTag<?>> idToTupleTagMap = new HashMap<>();

  // first output as the main output
  final TupleTag<OutT> mainOutputTag =
      outputs.isEmpty() ? null : new TupleTag(outputs.keySet().iterator().next());

  AtomicInteger index = new AtomicInteger(0);
  outputs
      .keySet()
      .iterator()
      .forEachRemaining(
          outputName -> {
            TupleTag<?> tupleTag = new TupleTag<>(outputName);
            tagToIndexMap.put(tupleTag, index.get());
            index.incrementAndGet();
            String collectionId = outputs.get(outputName);
            idToTupleTagMap.put(collectionId, tupleTag);
          });

  WindowedValue.WindowedValueCoder<InT> windowedInputCoder =
      ctx.instantiateCoder(inputId, pipeline.getComponents());

  final DoFnSchemaInformation doFnSchemaInformation;
  doFnSchemaInformation = ParDoTranslation.getSchemaInformation(transform.getTransform());

  Map<String, PCollectionView<?>> sideInputMapping =
      ParDoTranslation.getSideInputMapping(transform.getTransform());

  final RunnerApi.PCollection input = pipeline.getComponents().getPcollectionsOrThrow(inputId);
  final PCollection.IsBounded isBounded = SamzaPipelineTranslatorUtils.isBounded(input);

  final DoFnOp<InT, OutT, RawUnionValue> op =
      new DoFnOp<>(
          mainOutputTag,
          new NoOpDoFn<>(),
          null, // key coder not in use
          windowedInputCoder.getValueCoder(), // input coder not in use
          windowedInputCoder,
          Collections.emptyMap(), // output coders not in use
          Collections.emptyList(), // sideInputs not in use until side input support
          new ArrayList<>(idToTupleTagMap.values()), // used by java runner only
          SamzaPipelineTranslatorUtils.getPortableWindowStrategy(transform, pipeline),
          Collections.emptyMap(), // idToViewMap not in use until side input support
          new DoFnOp.MultiOutputManagerFactory(tagToIndexMap),
          ctx.getTransformFullName(),
          ctx.getTransformId(),
          isBounded,
          true,
          stagePayload,
          idToTupleTagMap,
          doFnSchemaInformation,
          sideInputMapping);

  final MessageStream<OpMessage<InT>> mergedStreams;
  if (sideInputStreams.isEmpty()) {
    mergedStreams = inputStream;
  } else {
    MessageStream<OpMessage<InT>> mergedSideInputStreams =
        MessageStream.mergeAll(sideInputStreams).flatMap(new SideInputWatermarkFn());
    mergedStreams = inputStream.merge(Collections.singletonList(mergedSideInputStreams));
  }

  final MessageStream<OpMessage<RawUnionValue>> taggedOutputStream =
      mergedStreams.flatMap(OpAdapter.adapt(op));

  for (int outputIndex : tagToIndexMap.values()) {
    final MessageStream<OpMessage<OutT>> outputStream =
        taggedOutputStream
            .filter(
                message ->
                    message.getType() != OpMessage.Type.ELEMENT
                        || message.getElement().getValue().getUnionTag() == outputIndex)
            .flatMap(OpAdapter.adapt(new RawUnionValueToValue()));

    ctx.registerMessageStream(ctx.getOutputId(transform), outputStream);
  }
}
 

/**
 * Translates transformNode from Beam into the Spark context.
 */
void translate(
    PipelineNode.PTransformNode transformNode,
    RunnerApi.Pipeline pipeline,
    Twister2StreamingTranslationContext context);