@Test public void getInputEmptySucceeds() { PTransformReplacement<PCollectionList<Long>, PCollection<Long>> replacement = factory.getReplacementTransform( AppliedPTransform.of( "nonEmptyInput", Collections.emptyMap(), Collections.emptyMap(), Flatten.pCollections(), pipeline)); assertThat(replacement.getInput().getAll(), emptyIterable()); }
private < InputT extends PInput, OutputT extends POutput, TransformT extends PTransform<? super InputT, OutputT>> void applyReplacement( Node original, PTransformOverrideFactory<InputT, OutputT, TransformT> replacementFactory) { PTransformReplacement<InputT, OutputT> replacement = replacementFactory.getReplacementTransform( (AppliedPTransform<InputT, OutputT, TransformT>) original.toAppliedPTransform(this)); if (replacement.getTransform() == original.getTransform()) { return; } InputT originalInput = replacement.getInput(); LOG.debug("Replacing {} with {}", original, replacement); transforms.replaceNode(original, originalInput, replacement.getTransform()); try { OutputT newOutput = replacement.getTransform().expand(originalInput); Map<PValue, ReplacementOutput> originalToReplacement = replacementFactory.mapOutputs(original.getOutputs(), newOutput); // Ensure the internal TransformHierarchy data structures are consistent. transforms.setOutput(newOutput); transforms.replaceOutputs(originalToReplacement); } finally { transforms.popNode(); } }
@Test public void getInputSucceeds() { TestPipeline p = TestPipeline.create(); PCollection<KV<String, Integer>> input = p.apply( Create.of(KV.of("foo", 1)) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); PCollection<KV<String, Iterable<Integer>>> grouped = input.apply(GroupByKey.create()); AppliedPTransform<?, ?, ?> producer = DirectGraphs.getProducer(grouped); PTransformReplacement< PCollection<KV<String, Integer>>, PCollection<KV<String, Iterable<Integer>>>> replacement = factory.getReplacementTransform((AppliedPTransform) producer); assertThat(replacement.getInput(), Matchers.<PCollection<?>>equalTo(input)); } }