original.expand(), PCollectionTuple.of(strsTag, replacementStrs) .and(moreIntsTag, moreReplacementInts)
private static AppliedPTransform<?, ?, ?> multiMultiParDo(Pipeline pipeline) { PCollectionView<String> view = pipeline.apply(Create.of("foo")).apply(View.asSingleton()); PCollection<Long> input = pipeline.apply(GenerateSequence.from(0)); ParDo.MultiOutput<Long, KV<Long, String>> parDo = ParDo.of(new TestDoFn()) .withSideInputs(view) .withOutputTags( new TupleTag<KV<Long, String>>() {}, TupleTagList.of(new TupleTag<KV<String, Long>>() {})); PCollectionTuple output = input.apply(parDo); Map<TupleTag<?>, PValue> inputs = new HashMap<>(); inputs.putAll(parDo.getAdditionalInputs()); inputs.putAll(input.expand()); return AppliedPTransform .<PCollection<Long>, PCollectionTuple, ParDo.MultiOutput<Long, KV<Long, String>>>of( "MultiParDoInAndOut", inputs, output.expand(), parDo, pipeline); } }
assertThat(tuple.getAll(), equalTo(pcsByTag)); PCollectionTuple reconstructed = PCollectionTuple.empty(p); for (Entry<TupleTag<?>, PValue> taggedValue : tuple.expand().entrySet()) { TupleTag<?> tag = taggedValue.getKey(); PValue value = taggedValue.getValue();
private PCollection<String> applySplittableParDo( String name, PCollection<Integer> input, DoFn<Integer, String> fn) { ParDo.MultiOutput<Integer, String> multiOutput = ParDo.of(fn).withOutputTags(MAIN_OUTPUT_TAG, TupleTagList.empty()); PCollectionTuple output = multiOutput.expand(input); output.get(MAIN_OUTPUT_TAG).setName("main"); AppliedPTransform<PCollection<Integer>, PCollectionTuple, ?> transform = AppliedPTransform.of("ParDo", input.expand(), output.expand(), multiOutput, pipeline); return input.apply(name, SplittableParDo.forAppliedParDo(transform)).get(MAIN_OUTPUT_TAG); }
Iterables.getOnlyElement(replacementOutput.expand().entrySet()); hierarchy.replaceOutputs( Collections.singletonMap(
@Test public void taggedMissingReplacementThrows() { PCollectionTuple original = PCollectionTuple.of(intsTag, ints).and(strsTag, strs).and(moreIntsTag, moreInts); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Missing replacement"); thrown.expectMessage(intsTag.toString()); thrown.expectMessage(ints.toString()); ReplacementOutputs.tagged( original.expand(), PCollectionTuple.of(strsTag, replacementStrs).and(moreIntsTag, moreReplacementInts)); }
@Test public void taggedExtraReplacementThrows() { PCollectionTuple original = PCollectionTuple.of(intsTag, ints).and(strsTag, strs); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Missing original output"); thrown.expectMessage(moreIntsTag.toString()); thrown.expectMessage(moreReplacementInts.toString()); ReplacementOutputs.tagged( original.expand(), PCollectionTuple.of(strsTag, replacementStrs) .and(moreIntsTag, moreReplacementInts) .and(intsTag, replacementInts)); } }
Iterables.getOnlyElement(replacementOutput.expand().entrySet()); hierarchy.replaceOutputs( Collections.singletonMap(
PTransformTranslation.toProto( AppliedPTransform.<PCollection<KV<Long, String>>, PCollection<Void>, MultiOutput>of( "foo", inputs, output.expand(), parDo, p), sdkComponents); RunnerApi.Components components = sdkComponents.toComponents();