org.apache.beam.runners.core.construction.graph.PipelineNode Java Examples

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

private static ExecutableStage createExecutableStage(Collection<SideInputReference> sideInputs) {
  Components components = Components.getDefaultInstance();
  Environment environment = Environment.getDefaultInstance();
  PCollectionNode inputCollection =
      PipelineNode.pCollection("collection-id", RunnerApi.PCollection.getDefaultInstance());
  return ImmutableExecutableStage.of(
      components,
      environment,
      inputCollection,
      sideInputs,
      Collections.emptyList(),
      Collections.emptyList(),
      Collections.emptyList(),
      Collections.emptyList(),
      DEFAULT_WIRE_CODER_SETTINGS);
}
 

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

public static void translate(RunnerApi.Pipeline pipeline, PortableTranslationContext ctx) {
  QueryablePipeline queryablePipeline =
      QueryablePipeline.forTransforms(
          pipeline.getRootTransformIdsList(), pipeline.getComponents());

  for (PipelineNode.PTransformNode transform :
      queryablePipeline.getTopologicallyOrderedTransforms()) {
    ctx.setCurrentTransform(transform);

    LOG.info("Translating transform urn: {}", transform.getTransform().getSpec().getUrn());
    TRANSLATORS
        .get(transform.getTransform().getSpec().getUrn())
        .translatePortable(transform, queryablePipeline, ctx);

    ctx.clearCurrentTransform();
  }
}
 

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

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

public WindowedValue.WindowedValueCoder instantiateCoder(
    String collectionId, RunnerApi.Components components) {
  PipelineNode.PCollectionNode collectionNode =
      PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId));
  try {
    return (WindowedValue.WindowedValueCoder)
        WireCoders.instantiateRunnerWireCoder(collectionNode, components);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}
 

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

public static WindowedValue.WindowedValueCoder instantiateCoder(
    String collectionId, RunnerApi.Components components) {
  PipelineNode.PCollectionNode collectionNode =
      PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId));
  try {
    return (WindowedValue.WindowedValueCoder)
        WireCoders.instantiateRunnerWireCoder(collectionNode, components);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}
 

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

/**
 * Returns an artificial PCollectionView that can be used to fulfill API requirements of a {@link
 * SideInputReader} when used inside the Dataflow runner harness.
 *
 * <p>Generates length prefixed coder variants suitable to be used within the Dataflow Runner
 * harness so that encoding and decoding values matches the length prefixing that occurred when
 * materializing the side input.
 */
public static final PCollectionView<?> transformSideInputForRunner(
    RunnerApi.Pipeline pipeline,
    RunnerApi.PTransform parDoPTransform,
    String sideInputTag,
    RunnerApi.SideInput sideInput) {
  checkArgument(
      Materializations.MULTIMAP_MATERIALIZATION_URN.equals(sideInput.getAccessPattern().getUrn()),
      "This handler is only capable of dealing with %s materializations "
          + "but was asked to handle %s for PCollectionView with tag %s.",
      Materializations.MULTIMAP_MATERIALIZATION_URN,
      sideInput.getAccessPattern().getUrn(),
      sideInputTag);
  String sideInputPCollectionId = parDoPTransform.getInputsOrThrow(sideInputTag);
  RunnerApi.PCollection sideInputPCollection =
      pipeline.getComponents().getPcollectionsOrThrow(sideInputPCollectionId);
  try {
    FullWindowedValueCoder<KV<Object, Object>> runnerSideInputCoder =
        (FullWindowedValueCoder)
            WireCoders.instantiateRunnerWireCoder(
                PipelineNode.pCollection(sideInputPCollectionId, sideInputPCollection),
                pipeline.getComponents());

    return DataflowPortabilityPCollectionView.with(
        new TupleTag<>(sideInputTag), runnerSideInputCoder);
  } catch (IOException e) {
    throw new IllegalStateException("Unable to translate proto to coder", e);
  }
}
 

/**
 * Translates the portable {@link org.apache.beam.model.pipeline.v1.RunnerApi.PTransform} into
 * Samza API.
 */
default void translatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  throw new UnsupportedOperationException(
      "Portable translation is not supported for " + this.getClass().getSimpleName());
}
 

/** Creates a coder for a given PCollection id from the Proto definition. */
public static <T> Coder<WindowedValue<T>> instantiateCoder(
    String collectionId, RunnerApi.Components components) {
  PipelineNode.PCollectionNode collectionNode =
      PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId));
  try {
    return WireCoders.instantiateRunnerWireCoder(collectionNode, components);
  } catch (IOException e) {
    throw new RuntimeException("Could not instantiate Coder", e);
  }
}
 

public static void createConfig(
    RunnerApi.Pipeline pipeline, ConfigBuilder configBuilder, SamzaPipelineOptions options) {
  QueryablePipeline queryablePipeline =
      QueryablePipeline.forTransforms(
          pipeline.getRootTransformIdsList(), pipeline.getComponents());
  for (PipelineNode.PTransformNode transform :
      queryablePipeline.getTopologicallyOrderedTransforms()) {
    TransformTranslator<?> translator =
        TRANSLATORS.get(transform.getTransform().getSpec().getUrn());
    if (translator instanceof TransformConfigGenerator) {
      TransformConfigGenerator configGenerator = (TransformConfigGenerator) translator;
      configBuilder.putAll(configGenerator.createPortableConfig(transform, options));
    }
  }
}
 

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

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

private static <T> WindowedValueCoder<T> getWindowedValueCoder(
    String pCollectionId, RunnerApi.Components components) {
  PCollection pCollection = components.getPcollectionsOrThrow(pCollectionId);
  PCollectionNode pCollectionNode = PipelineNode.pCollection(pCollectionId, pCollection);
  WindowedValueCoder<T> coder;
  try {
    coder =
        (WindowedValueCoder) WireCoders.instantiateRunnerWireCoder(pCollectionNode, components);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  return coder;
}
 

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
@SuppressWarnings("unchecked")
public ExecutableStage createExecutableStage() throws Exception {
	RunnerApi.Components components =
		RunnerApi.Components.newBuilder()
			.putPcollections(
				INPUT_ID,
				RunnerApi.PCollection.newBuilder()
					.setWindowingStrategyId(WINDOW_STRATEGY)
					.setCoderId(INPUT_CODER_ID)
					.build())
			.putPcollections(
				OUTPUT_ID,
				RunnerApi.PCollection.newBuilder()
					.setWindowingStrategyId(WINDOW_STRATEGY)
					.setCoderId(OUTPUT_CODER_ID)
					.build())
			.putTransforms(
				TRANSFORM_ID,
				RunnerApi.PTransform.newBuilder()
					.setUniqueName(TRANSFORM_ID)
					.setSpec(RunnerApi.FunctionSpec.newBuilder()
						.setUrn(functionUrn)
						.setPayload(
							org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString.copyFrom(
								getUserDefinedFunctionsProto().toByteArray()))
						.build())
					.putInputs(MAIN_INPUT_NAME, INPUT_ID)
					.putOutputs(MAIN_OUTPUT_NAME, OUTPUT_ID)
					.build())
			.putWindowingStrategies(
				WINDOW_STRATEGY,
				RunnerApi.WindowingStrategy.newBuilder()
					.setWindowCoderId(WINDOW_CODER_ID)
					.build())
			.putCoders(
				INPUT_CODER_ID,
				getInputCoderProto())
			.putCoders(
				OUTPUT_CODER_ID,
				getOutputCoderProto())
			.putCoders(
				WINDOW_CODER_ID,
				getWindowCoderProto())
			.build();

	PipelineNode.PCollectionNode input =
		PipelineNode.pCollection(INPUT_ID, components.getPcollectionsOrThrow(INPUT_ID));
	List<SideInputReference> sideInputs = Collections.EMPTY_LIST;
	List<UserStateReference> userStates = Collections.EMPTY_LIST;
	List<TimerReference> timers = Collections.EMPTY_LIST;
	List<PipelineNode.PTransformNode> transforms =
		Collections.singletonList(
			PipelineNode.pTransform(TRANSFORM_ID, components.getTransformsOrThrow(TRANSFORM_ID)));
	List<PipelineNode.PCollectionNode> outputs =
		Collections.singletonList(
			PipelineNode.pCollection(OUTPUT_ID, components.getPcollectionsOrThrow(OUTPUT_ID)));
	return ImmutableExecutableStage.of(
		components, createPythonExecutionEnvironment(), input, sideInputs, userStates, timers, transforms, outputs, createValueOnlyWireCoderSetting());
}
 

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

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

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

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

/** Generate config for portable api PTransform. */
default Map<String, String> createPortableConfig(
    PipelineNode.PTransformNode transform, SamzaPipelineOptions options) {
  return Collections.emptyMap();
}
 

private static <K, InputT, OutputT> void doTranslatePortable(
    PipelineNode.PTransformNode transform,
    QueryablePipeline pipeline,
    PortableTranslationContext ctx) {
  final MessageStream<OpMessage<KV<K, InputT>>> inputStream =
      ctx.getOneInputMessageStream(transform);
  final boolean needRepartition = ctx.getSamzaPipelineOptions().getMaxSourceParallelism() > 1;
  final WindowingStrategy<?, BoundedWindow> windowingStrategy =
      ctx.getPortableWindowStrategy(transform, pipeline);
  final Coder<BoundedWindow> windowCoder = windowingStrategy.getWindowFn().windowCoder();

  final String inputId = ctx.getInputId(transform);
  final WindowedValue.WindowedValueCoder<KV<K, InputT>> windowedInputCoder =
      ctx.instantiateCoder(inputId, pipeline.getComponents());
  final KvCoder<K, InputT> kvInputCoder = (KvCoder<K, InputT>) windowedInputCoder.getValueCoder();
  final Coder<WindowedValue<KV<K, InputT>>> elementCoder =
      WindowedValue.FullWindowedValueCoder.of(kvInputCoder, windowCoder);

  final TupleTag<KV<K, OutputT>> outputTag =
      new TupleTag<>(Iterables.getOnlyElement(transform.getTransform().getOutputsMap().keySet()));

  @SuppressWarnings("unchecked")
  final SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> reduceFn =
      (SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow>)
          SystemReduceFn.buffering(kvInputCoder.getValueCoder());

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

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