public static TupleTag<?> getMainOutputTag(AppliedPTransform<?, ?, ?> application) throws IOException { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.MultiOutput) { return ((ParDo.MultiOutput<?, ?>) transform).getMainOutputTag(); } return getMainOutputTag(getParDoPayload(application)); }
/** * Creates the transform for a {@link ParDo}-compatible {@link AppliedPTransform}. * * <p>The input may generally be a deserialized transform so it may not actually be a {@link * ParDo}. Instead {@link ParDoTranslation} will be used to extract fields. */ @SuppressWarnings({"unchecked", "rawtypes"}) public static <InputT, OutputT> SplittableParDo<InputT, OutputT, ?> forAppliedParDo( AppliedPTransform<PCollection<InputT>, PCollectionTuple, ?> parDo) { checkArgument(parDo != null, "parDo must not be null"); try { Map<TupleTag<?>, Coder<?>> outputTagsToCoders = Maps.newHashMap(); for (Map.Entry<TupleTag<?>, PValue> entry : parDo.getOutputs().entrySet()) { outputTagsToCoders.put(entry.getKey(), ((PCollection) entry.getValue()).getCoder()); } return new SplittableParDo( ParDoTranslation.getDoFn(parDo), ParDoTranslation.getSideInputs(parDo), ParDoTranslation.getMainOutputTag(parDo), ParDoTranslation.getAdditionalOutputTags(parDo), outputTagsToCoders); } catch (IOException exc) { throw new RuntimeException(exc); } }
public static TupleTagList getAdditionalOutputTags(AppliedPTransform<?, ?, ?> application) throws IOException { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.MultiOutput) { return ((ParDo.MultiOutput<?, ?>) transform).getAdditionalOutputTags(); } RunnerApi.PTransform protoTransform = PTransformTranslation.toProto( application, SdkComponents.create(application.getPipeline().getOptions())); ParDoPayload payload = ParDoPayload.parseFrom(protoTransform.getSpec().getPayload()); TupleTag<?> mainOutputTag = getMainOutputTag(payload); Set<String> outputTags = Sets.difference( protoTransform.getOutputsMap().keySet(), Collections.singleton(mainOutputTag.getId())); ArrayList<TupleTag<?>> additionalOutputTags = new ArrayList<>(); for (String outputTag : outputTags) { additionalOutputTags.add(new TupleTag<>(outputTag)); } return TupleTagList.of(additionalOutputTags); }
mainOutputTag = ParDoTranslation.getMainOutputTag(context.getCurrentTransform()); } catch (IOException e) { throw new RuntimeException(e);
mainOutputTag = ParDoTranslation.getMainOutputTag(context.getCurrentTransform()); } catch (IOException e) { throw new RuntimeException(e);
try { mainOutputTag = (TupleTag<OutputT>) ParDoTranslation.getMainOutputTag(context.getCurrentTransform()); } catch (IOException e) { throw new RuntimeException(e);
try { mainOutputTag = (TupleTag<OutputT>) ParDoTranslation.getMainOutputTag(context.getCurrentTransform()); } catch (IOException e) { throw new RuntimeException(e);
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);
@Test public void testToProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ParDoPayload payload = ParDoTranslation.translateParDo(parDo, p, components); assertThat(ParDoTranslation.getDoFn(payload), equalTo(parDo.getFn())); assertThat(ParDoTranslation.getMainOutputTag(payload), equalTo(parDo.getMainOutputTag())); for (PCollectionView<?> view : parDo.getSideInputs()) { payload.getSideInputsOrThrow(view.getTagInternal().getId()); } }