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(); } }
Read.Bounded<Long> replacementTransform = Read.from(CountingSource.upTo(1L)); PCollection<Long> replacementOutput = pipeline.apply(replacementTransform); Node replacement = hierarchy.replaceNode(original, PBegin.in(pipeline), replacementTransform); assertThat(hierarchy.getCurrent(), equalTo(replacement)); hierarchy.setOutput(replacementOutput);
Node compositeNode = hierarchy.replaceNode(original, upstream, replacementComposite); Node replacementParNode = hierarchy.pushNode("Original/Contained", upstream, replacementParDo); hierarchy.finishSpecifyingInput();
Node compositeNode = hierarchy.replaceNode(original, upstream, replacementComposite); Node replacementParNode = hierarchy.pushNode("Original/Contained", upstream, replacementParDo); hierarchy.finishSpecifyingInput();