Java Code Examples for org.apache.beam.sdk.testing.TestStream#TestStreamCoder

@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 static <T> byte[] getEncodedPayload(
    TestStream<T> testStream, TestStream.TestStreamCoder<T> coder) {
  try {
    return CoderUtils.encodeToByteArray(coder, testStream);
  } catch (CoderException e) {
    throw ExceptionUtil.rethrow(e);
  }
}
 

@SuppressWarnings("unchecked")
private TestStreamP(byte[] payload, TestStream.TestStreamCoder payloadCoder, Coder outputCoder) {
  List events = decodePayload(payload, payloadCoder).getEvents();
  traverser =
      Traversers.traverseStream(
          events.stream()
              .flatMap(
                  event -> {
                    if (event instanceof TestStream.WatermarkEvent) {
                      Instant watermark = ((TestStream.WatermarkEvent) event).getWatermark();
                      if (BoundedWindow.TIMESTAMP_MAX_VALUE.equals(watermark)) {
                        // this is an element added by advanceWatermarkToInfinity(), we ignore it,
                        // it's always at the end
                        return null;
                      }
                      return Stream.of(new Watermark(watermark.getMillis()));
                    } else if (event instanceof TestStream.ElementEvent) {
                      return StreamSupport.stream(
                              ((TestStream.ElementEvent<?>) event).getElements().spliterator(),
                              false)
                          .map(
                              tv ->
                                  WindowedValue.timestampedValueInGlobalWindow(
                                      tv.getValue(), tv.getTimestamp()))
                          .map(wV -> Utils.encode(wV, outputCoder));
                    } else {
                      throw new UnsupportedOperationException(
                          "Event type not supported in TestStream: "
                              + event.getClass()
                              + ", event: "
                              + event);
                    }
                  }));
}
 

private static TestStream decodePayload(byte[] payload, TestStream.TestStreamCoder coder) {
  try {
    return (TestStream) CoderUtils.decodeFromByteArray(coder, payload);
  } catch (CoderException e) {
    throw ExceptionUtil.rethrow(e);
  }
}
 

public static <T> ProcessorMetaSupplier supplier(
    byte[] payload, TestStream.TestStreamCoder payloadCoder, Coder outputCoder) {
  return ProcessorMetaSupplier.forceTotalParallelismOne(
      ProcessorSupplier.of(
          () -> new ThrottleWrappedP(new TestStreamP(payload, payloadCoder, outputCoder), 4)));
}