Java Code Examples for org.apache.beam.model.pipeline.v1.RunnerApi#SideInput

public static List<PCollectionView<?>> getSideInputs(AppliedPTransform<?, ?, ?> application)
    throws IOException {
  PTransform<?, ?> transform = application.getTransform();
  if (transform instanceof ParDo.MultiOutput) {
    return ((ParDo.MultiOutput<?, ?>) transform)
        .getSideInputs().values().stream().collect(Collectors.toList());
  }

  SdkComponents sdkComponents = SdkComponents.create(application.getPipeline().getOptions());
  RunnerApi.PTransform parDoProto = PTransformTranslation.toProto(application, sdkComponents);
  ParDoPayload payload = ParDoPayload.parseFrom(parDoProto.getSpec().getPayload());

  List<PCollectionView<?>> views = new ArrayList<>();
  RehydratedComponents components =
      RehydratedComponents.forComponents(sdkComponents.toComponents());
  for (Map.Entry<String, SideInput> sideInputEntry : payload.getSideInputsMap().entrySet()) {
    String sideInputTag = sideInputEntry.getKey();
    RunnerApi.SideInput sideInput = sideInputEntry.getValue();
    PCollection<?> originalPCollection =
        checkNotNull(
            (PCollection<?>) application.getInputs().get(new TupleTag<>(sideInputTag)),
            "no input with tag %s",
            sideInputTag);
    views.add(
        PCollectionViewTranslation.viewFromProto(
            sideInput, sideInputTag, originalPCollection, parDoProto, components));
  }
  return views;
}
 

/**
 * Create a {@link PCollectionView} from a side input spec and an already-deserialized {@link
 * PCollection} that should be wired up.
 */
public static PCollectionView<?> viewFromProto(
    RunnerApi.SideInput sideInput,
    String localName,
    PCollection<?> pCollection,
    RunnerApi.PTransform parDoTransform,
    RehydratedComponents components)
    throws IOException {
  checkArgument(
      localName != null,
      "%s.viewFromProto: localName must not be null",
      ParDoTranslation.class.getSimpleName());
  TupleTag<?> tag = new TupleTag<>(localName);
  WindowMappingFn<?> windowMappingFn = windowMappingFnFromProto(sideInput.getWindowMappingFn());
  ViewFn<?, ?> viewFn = viewFnFromProto(sideInput.getViewFn());
  WindowingStrategy<?, ?> windowingStrategy = pCollection.getWindowingStrategy().fixDefaults();

  PCollectionView<?> view =
      new RunnerPCollectionView<>(
          pCollection,
          (TupleTag) tag,
          (ViewFn) viewFn,
          windowMappingFn,
          windowingStrategy,
          (Coder) pCollection.getCoder());
  return view;
}
 

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

Map<String, RunnerApi.SideInput> translateSideInputs(SdkComponents components); 

Map<String, RunnerApi.SideInput> translateSideInputs(SdkComponents components);