Java Code Examples for org.apache.beam.runners.core.construction.ParDoTranslation#getMainOutputTag()

The following examples show how to use org.apache.beam.runners.core.construction.ParDoTranslation#getMainOutputTag() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: ParDoTranslatorBatch.java    From beam with Apache License 2.0 5 votes vote down vote up 
private TupleTag<?> getTupleTag(TranslationContext context) {
  TupleTag<?> mainOutputTag;
  try {
    mainOutputTag = ParDoTranslation.getMainOutputTag(context.getCurrentTransform());
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  return mainOutputTag;
}
Example 2
Source File: ParDoMultiOverrideFactory.java    From beam with Apache License 2.0 5 votes vote down vote up 
@SuppressWarnings("unchecked")
private PTransform<PCollection<? extends InputT>, PCollectionTuple> getReplacementForApplication(
    AppliedPTransform<
            PCollection<? extends InputT>,
            PCollectionTuple,
            PTransform<PCollection<? extends InputT>, PCollectionTuple>>
        application)
    throws IOException {

  DoFn<InputT, OutputT> fn = (DoFn<InputT, OutputT>) ParDoTranslation.getDoFn(application);

  DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());

  if (signature.processElement().isSplittable()) {
    return SplittableParDo.forAppliedParDo((AppliedPTransform) application);
  } else if (signature.stateDeclarations().size() > 0
      || signature.timerDeclarations().size() > 0
      || signature.timerFamilyDeclarations().size() > 0) {
    return new GbkThenStatefulParDo(
        fn,
        ParDoTranslation.getMainOutputTag(application),
        ParDoTranslation.getAdditionalOutputTags(application),
        ParDoTranslation.getSideInputs(application),
        ParDoTranslation.getSchemaInformation(application),
        ParDoTranslation.getSideInputMapping(application));
  } else {
    return application.getTransform();
  }
}
Example 3
Source File: PipelineTranslator.java    From incubator-nemo with Apache License 2.0 4 votes vote down vote up 
/**
 * @param ctx          provides translation context.
 * @param beamNode     the beam node to be translated.
 * @param sideInputMap side inputs.
 * @return the created DoFnTransform.
 */
private static AbstractDoFnTransform createDoFnTransform(final PipelineTranslationContext ctx,
                                                         final TransformHierarchy.Node beamNode,
                                                         final Map<Integer, PCollectionView<?>> sideInputMap) {
  try {
    final AppliedPTransform pTransform = beamNode.toAppliedPTransform(ctx.getPipeline());
    final DoFn doFn = ParDoTranslation.getDoFn(pTransform);
    final TupleTag mainOutputTag = ParDoTranslation.getMainOutputTag(pTransform);
    final TupleTagList additionalOutputTags = ParDoTranslation.getAdditionalOutputTags(pTransform);

    final PCollection<?> mainInput = (PCollection<?>)
      Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(pTransform));

    final HasDisplayData displayData = (builder) -> builder.add(DisplayData.item("name", beamNode.getFullName()));
    final DoFnSchemaInformation doFnSchemaInformation =
      ParDoTranslation.getSchemaInformation(beamNode.toAppliedPTransform(ctx.getPipeline()));

    if (sideInputMap.isEmpty()) {
      return new DoFnTransform(
        doFn,
        mainInput.getCoder(),
        getOutputCoders(pTransform),
        mainOutputTag,
        additionalOutputTags.getAll(),
        mainInput.getWindowingStrategy(),
        ctx.getPipelineOptions(),
        DisplayData.from(displayData),
        doFnSchemaInformation,
        Collections.emptyMap());
    } else {
      return new PushBackDoFnTransform(
        doFn,
        mainInput.getCoder(),
        getOutputCoders(pTransform),
        mainOutputTag,
        additionalOutputTags.getAll(),
        mainInput.getWindowingStrategy(),
        sideInputMap,
        ctx.getPipelineOptions(),
        DisplayData.from(displayData),
        doFnSchemaInformation,
        Collections.emptyMap());
    }
  } catch (final IOException e) {
    throw new RuntimeException(e);
  }
}
Example 4
Source File: ParDoMultiOutputTranslatorBatch.java    From twister2 with Apache License 2.0 4 votes vote down vote up 
@Override
public void translateNode(
    ParDo.MultiOutput<IT, OT> transform, Twister2BatchTranslationContext context) {
  DoFn<IT, OT> doFn;
  doFn = transform.getFn();
  BatchTSetImpl<WindowedValue<IT>> inputTTSet =
      context.getInputDataSet(context.getInput(transform));

  WindowingStrategy<?, ?> windowingStrategy = context.getInput(transform).getWindowingStrategy();
  Coder<IT> inputCoder = (Coder<IT>) context.getInput(transform).getCoder();

  Map<TupleTag<?>, PValue> outputs = context.getOutputs();
  Map<TupleTag<?>, Coder<?>> outputCoders = context.getOutputCoders();

  DoFnSignature signature = DoFnSignatures.getSignature(transform.getFn().getClass());
  DoFnSchemaInformation doFnSchemaInformation;
  doFnSchemaInformation = ParDoTranslation.getSchemaInformation(context.getCurrentTransform());
  TupleTag<OT> mainOutput = transform.getMainOutputTag();
  List<TupleTag<?>> additionalOutputTags = new ArrayList<>(outputs.size() - 1);
  Collection<PCollectionView<?>> sideInputs = transform.getSideInputs();

  // construct a map from side input to WindowingStrategy so that
  // the DoFn runner can map main-input windows to side input windows
  Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputStrategies = new HashMap<>();
  for (PCollectionView<?> sideInput : sideInputs) {
    sideInputStrategies.put(sideInput, sideInput.getWindowingStrategyInternal());
  }

  TupleTag<?> mainOutputTag;
  try {
    mainOutputTag = ParDoTranslation.getMainOutputTag(context.getCurrentTransform());
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  Map<TupleTag<?>, Integer> outputMap = Maps.newHashMap();
  // put the main output at index 0, FlinkMultiOutputDoFnFunction  expects this
  outputMap.put(mainOutputTag, 0);
  int count = 1;
  for (TupleTag<?> tag : outputs.keySet()) {
    if (!outputMap.containsKey(tag)) {
      outputMap.put(tag, count++);
    }
  }

  ComputeTSet<RawUnionValue, Iterator<WindowedValue<IT>>> outputTSet =
      inputTTSet
          .direct()
          .<RawUnionValue>compute(
              new DoFnFunction<OT, IT>(
                  context,
                  doFn,
                  inputCoder,
                  outputCoders,
                  additionalOutputTags,
                  windowingStrategy,
                  sideInputStrategies,
                  mainOutput,
                  doFnSchemaInformation,
                  outputMap));

  for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
    ComputeTSet<WindowedValue<OT>, Iterator<RawUnionValue>> tempTSet =
        outputTSet.direct().compute(new OutputTagFilter(outputMap.get(output.getKey())));
    context.setOutputDataSet((PCollection) output.getValue(), tempTSet);
  }
}
Example 5
Source File: ParDoMultiOutputTranslatorBatch.java    From beam with Apache License 2.0 4 votes vote down vote up 
@Override
public void translateNode(
    ParDo.MultiOutput<InputT, OutputT> transform, Twister2BatchTranslationContext context) {
  DoFn<InputT, OutputT> doFn;
  doFn = transform.getFn();
  if (DoFnSignatures.signatureForDoFn(doFn).processElement().isSplittable()) {
    throw new UnsupportedOperationException(
        String.format(
            "Not expected to directly translate splittable DoFn, should have been overridden: %s",
            doFn));
  }
  BatchTSetImpl<WindowedValue<InputT>> inputTTSet =
      context.getInputDataSet(context.getInput(transform));

  WindowingStrategy<?, ?> windowingStrategy = context.getInput(transform).getWindowingStrategy();
  Coder<InputT> inputCoder = (Coder<InputT>) context.getInput(transform).getCoder();
  Map<String, PCollectionView<?>> sideInputMapping;

  Map<TupleTag<?>, PValue> outputs = context.getOutputs();
  Map<TupleTag<?>, Coder<?>> outputCoders = context.getOutputCoders();

  // DoFnSignature signature = DoFnSignatures.getSignature(transform.getFn().getClass());
  DoFnSchemaInformation doFnSchemaInformation;
  doFnSchemaInformation = ParDoTranslation.getSchemaInformation(context.getCurrentTransform());
  sideInputMapping = ParDoTranslation.getSideInputMapping(context.getCurrentTransform());

  TupleTag<OutputT> mainOutput = transform.getMainOutputTag();
  List<TupleTag<?>> additionalOutputTags =
      new ArrayList<>(transform.getAdditionalOutputTags().getAll());
  Map<String, PCollectionView<?>> sideInputs = transform.getSideInputs();
  // TODO : note change from List to map in sideinputs

  // construct a map from side input to WindowingStrategy so that
  // the DoFn runner can map main-input windows to side input windows
  Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputStrategies = new HashMap<>();
  for (PCollectionView<?> sideInput : sideInputs.values()) {
    sideInputStrategies.put(sideInput, sideInput.getWindowingStrategyInternal());
  }

  TupleTag<?> mainOutputTag;
  try {
    mainOutputTag = ParDoTranslation.getMainOutputTag(context.getCurrentTransform());
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  Map<TupleTag<?>, Integer> outputMap = Maps.newHashMap();
  outputMap.put(mainOutputTag, 0);
  int count = 1;
  for (TupleTag<?> tag : outputs.keySet()) {
    if (!outputMap.containsKey(tag)) {
      outputMap.put(tag, count++);
    }
  }

  ComputeTSet<RawUnionValue, Iterator<WindowedValue<InputT>>> outputTSet =
      inputTTSet
          .direct()
          .<RawUnionValue>compute(
              new DoFnFunction<OutputT, InputT>(
                  context,
                  doFn,
                  inputCoder,
                  outputCoders,
                  additionalOutputTags,
                  windowingStrategy,
                  sideInputStrategies,
                  mainOutput,
                  doFnSchemaInformation,
                  outputMap,
                  sideInputMapping));

  for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
    ComputeTSet<WindowedValue<OutputT>, Iterator<RawUnionValue>> tempTSet =
        outputTSet.direct().compute(new OutputTagFilter(outputMap.get(output.getKey())));
    context.setOutputDataSet((PCollection) output.getValue(), tempTSet);
  }
}