You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by tg...@apache.org on 2017/02/14 21:13:28 UTC
[1/2] beam git commit: Add Pipeline#replaceTransform
Repository: beam
Updated Branches:
refs/heads/master 938a66ec3 -> 345c2acca
Add Pipeline#replaceTransform
This is the base method for Pipeline Surgery. It takes a
PTransformMatcher and a PTransformOverrideFactory, and replaces all
matching PTransforms with the result of the Override Factory.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/16c051ac
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/16c051ac
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/16c051ac
Branch: refs/heads/master
Commit: 16c051aca7fadfe3935988488f29723b932e47c7
Parents: 938a66e
Author: Thomas Groh <tg...@google.com>
Authored: Thu Feb 9 11:40:50 2017 -0800
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Feb 14 13:12:56 2017 -0800
----------------------------------------------------------------------
.../main/java/org/apache/beam/sdk/Pipeline.java | 63 ++++-
.../beam/sdk/runners/TransformHierarchy.java | 157 +++++++++---
.../apache/beam/sdk/values/TupleTagList.java | 6 +
.../sdk/runners/TransformHierarchyTest.java | 237 ++++++++++++++++++-
4 files changed, 431 insertions(+), 32 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/16c051ac/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java
index eb0b199..109424d 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java
@@ -27,9 +27,14 @@ import java.util.Set;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptions.CheckEnabled;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.runners.PTransformMatcher;
+import org.apache.beam.sdk.runners.PTransformOverrideFactory;
+import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
import org.apache.beam.sdk.runners.PipelineRunner;
import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.runners.TransformHierarchy.Node;
import org.apache.beam.sdk.transforms.Aggregator;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.PTransform;
@@ -167,6 +172,33 @@ public class Pipeline {
return begin().apply(name, root);
}
+ public void replace(
+ final PTransformMatcher matcher, PTransformOverrideFactory replacementFactory) {
+ final Collection<Node> matches = new ArrayList<>();
+ transforms.visit(
+ new PipelineVisitor.Defaults() {
+ @Override
+ public CompositeBehavior enterCompositeTransform(Node node) {
+ if (!node.isRootNode() && matcher.matches(node.toAppliedPTransform())) {
+ matches.add(node);
+ // This node will be replaced. It should not be visited.
+ return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
+ }
+ return CompositeBehavior.ENTER_TRANSFORM;
+ }
+
+ @Override
+ public void visitPrimitiveTransform(Node node) {
+ if (matcher.matches(node.toAppliedPTransform())) {
+ matches.add(node);
+ }
+ }
+ });
+ for (Node match : matches) {
+ applyReplacement(match, replacementFactory);
+ }
+ }
+
/**
* Runs the {@link Pipeline} using its {@link PipelineRunner}.
*/
@@ -323,7 +355,7 @@ public class Pipeline {
private final PipelineRunner<?> runner;
private final PipelineOptions options;
- private final TransformHierarchy transforms = new TransformHierarchy();
+ private final TransformHierarchy transforms = new TransformHierarchy(this);
private Collection<PValue> values = new ArrayList<>();
private Set<String> usedFullNames = new HashSet<>();
private CoderRegistry coderRegistry;
@@ -394,6 +426,35 @@ public class Pipeline {
}
}
+ private <InputT extends PInput, OutputT extends POutput,
+ TransformT extends PTransform<? super InputT, OutputT>>
+ void applyReplacement(
+ Node original,
+ PTransformOverrideFactory<InputT, OutputT, TransformT> replacementFactory) {
+ // Names for top-level transforms have been assigned. Any new collisions are within a node
+ // and its replacement.
+ getOptions().setStableUniqueNames(CheckEnabled.OFF);
+ PTransform<InputT, OutputT> replacement =
+ replacementFactory.getReplacementTransform((TransformT) original.getTransform());
+ if (replacement == original.getTransform()) {
+ return;
+ }
+ InputT originalInput = replacementFactory.getInput(original.getInputs(), this);
+
+ LOG.debug("Replacing {} with {}", original, replacement);
+ transforms.replaceNode(original, originalInput, replacement);
+ try {
+ OutputT newOutput = runner.apply(replacement, 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();
+ }
+ }
+
/**
* Returns the configured {@link PipelineRunner}.
*/
http://git-wip-us.apache.org/repos/asf/beam/blob/16c051ac/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
index f2e3097..d626543 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
@@ -17,13 +17,13 @@
*/
package org.apache.beam.sdk.runners;
+import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
import java.util.ArrayList;
-import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
@@ -32,32 +32,42 @@ import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import javax.annotation.Nullable;
+import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.Pipeline.PipelineVisitor.CompositeBehavior;
+import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
import org.apache.beam.sdk.transforms.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TaggedPValue;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Captures information about a collection of transformations and their
* associated {@link PValue}s.
*/
public class TransformHierarchy {
+ private static final Logger LOG = LoggerFactory.getLogger(TransformHierarchy.class);
+
+ private final Pipeline pipeline;
private final Node root;
+ private final Map<Node, PInput> unexpandedInputs;
private final Map<POutput, Node> producers;
// A map of PValue to the PInput the producing PTransform is applied to
private final Map<PValue, PInput> producerInput;
// Maintain a stack based on the enclosing nodes
private Node current;
- public TransformHierarchy() {
- root = new Node(null, null, "", null);
- current = root;
+ public TransformHierarchy(Pipeline pipeline) {
+ this.pipeline = pipeline;
producers = new HashMap<>();
producerInput = new HashMap<>();
+ unexpandedInputs = new HashMap<>();
+ root = new Node(null, null, "", null);
+ current = root;
}
/**
@@ -78,11 +88,27 @@ public class TransformHierarchy {
checkNotNull(
input, "An input must be provided for all %s Nodes", PTransform.class.getSimpleName());
Node node = new Node(current, transform, name, input);
+ unexpandedInputs.put(node, input);
current.addComposite(node);
current = node;
return current;
}
+ public Node replaceNode(Node existing, PInput input, PTransform<?, ?> transform) {
+ checkNotNull(existing);
+ checkNotNull(input);
+ checkNotNull(transform);
+ checkState(
+ unexpandedInputs.isEmpty(),
+ "Replacing a node when the graph has an unexpanded input. This is an SDK bug.");
+ Node replacement =
+ new Node(existing.getEnclosingNode(), transform, existing.getFullName(), input);
+ existing.getEnclosingNode().replaceChild(existing, replacement);
+ unexpandedInputs.remove(existing);
+ current = replacement;
+ return replacement;
+ }
+
/**
* Finish specifying all of the input {@link PValue PValues} of the current {@link
* Node}. Ensures that all of the inputs to the current node have been fully
@@ -120,21 +146,31 @@ public class TransformHierarchy {
if (!producers.containsKey(value.getValue())) {
producers.put(value.getValue(), current);
}
- value.getValue().finishSpecifyingOutput(current.input, current.transform);
- producerInput.put(value.getValue(), current.input);
+ value.getValue().finishSpecifyingOutput(unexpandedInputs.get(current), current.transform);
+ producerInput.put(value.getValue(), unexpandedInputs.get(current));
}
- output.finishSpecifyingOutput(current.input, current.transform);
+ output.finishSpecifyingOutput(unexpandedInputs.get(current), current.transform);
current.setOutput(output);
// TODO: Replace with a "generateDefaultNames" method.
output.recordAsOutput(current.toAppliedPTransform());
}
/**
+ * Recursively replace the outputs of the current {@link Node} with the original outputs of the
+ * node it is replacing. No value that is a key in {@code originalToReplacement} may be present
+ * within the {@link TransformHierarchy} after this method completes.
+ */
+ public void replaceOutputs(Map<PValue, ReplacementOutput> originalToReplacement) {
+ current.replaceOutputs(originalToReplacement);
+ }
+
+ /**
* Pops the current node off the top of the stack, finishing it. Outputs of the node are finished
* once they are consumed as input.
*/
public void popNode() {
current.finishSpecifying();
+ unexpandedInputs.remove(current);
current = current.getEnclosingNode();
checkState(current != null, "Can't pop the root node of a TransformHierarchy");
}
@@ -181,14 +217,14 @@ public class TransformHierarchy {
private final String fullName;
// Nodes for sub-transforms of a composite transform.
- private final Collection<Node> parts = new ArrayList<>();
+ private final List<Node> parts = new ArrayList<>();
- // Input to the transform, in unexpanded form.
- private final PInput input;
+ // Input to the transform, in expanded form.
+ private final List<TaggedPValue> inputs;
// TODO: track which outputs need to be exported to parent.
- // Output of the transform, in unexpanded form.
- private POutput output;
+ // Output of the transform, in expanded form.
+ private List<TaggedPValue> outputs;
@VisibleForTesting
boolean finishedSpecifying = false;
@@ -212,7 +248,7 @@ public class TransformHierarchy {
this.enclosingNode = enclosingNode;
this.transform = transform;
this.fullName = fullName;
- this.input = input;
+ this.inputs = input == null ? Collections.<TaggedPValue>emptyList() : input.expand();
}
/**
@@ -240,6 +276,29 @@ public class TransformHierarchy {
}
/**
+ * Replace a sub-node of this transform node.
+ *
+ * <p>The existing {@link Node} must be a direct child of this {@link Node}. The replacement
+ * node need not be fully specified.
+ */
+ public void replaceChild(Node existing, Node replacement) {
+ checkNotNull(existing);
+ checkNotNull(replacement);
+ int existingIndex = parts.indexOf(existing);
+ checkArgument(
+ existingIndex >= 0,
+ "Tried to replace a node %s that doesn't exist as a component of node %s",
+ existing.getFullName(),
+ getFullName());
+ LOG.info(
+ "Replaced original node {} with replacement {} at index {}",
+ existing,
+ replacement,
+ existingIndex);
+ parts.set(existingIndex, replacement);
+ }
+
+ /**
* Returns true if this node represents a composite transform that does not perform processing
* of its own, but merely encapsulates a sub-pipeline (which may be empty).
*
@@ -252,8 +311,8 @@ public class TransformHierarchy {
private boolean returnsOthersOutput() {
PTransform<?, ?> transform = getTransform();
- if (output != null) {
- for (TaggedPValue outputValue : output.expand()) {
+ if (outputs != null) {
+ for (TaggedPValue outputValue : outputs) {
if (!getProducer(outputValue.getValue()).getTransform().equals(transform)) {
return true;
}
@@ -272,7 +331,7 @@ public class TransformHierarchy {
/** Returns the transform input, in unexpanded form. */
public List<TaggedPValue> getInputs() {
- return input == null ? Collections.<TaggedPValue>emptyList() : input.expand();
+ return inputs == null ? Collections.<TaggedPValue>emptyList() : inputs;
}
/**
@@ -281,9 +340,9 @@ public class TransformHierarchy {
private void setOutput(POutput output) {
checkState(!finishedSpecifying);
checkState(
- this.output == null, "Tried to specify more than one output for %s", getFullName());
+ this.outputs == null, "Tried to specify more than one output for %s", getFullName());
checkNotNull(output, "Tried to set the output of %s to null", getFullName());
- this.output = output;
+ this.outputs = output.expand();
// Validate that a primitive transform produces only primitive output, and a composite
// transform does not produce primitive output.
@@ -315,9 +374,53 @@ public class TransformHierarchy {
}
}
- /** Returns the transform output, in unexpanded form. */
+ /**
+ * Replaces each value in {@code originalToReplacement} present in this {@link Node Node's}
+ * outputs with the key that maps to that value.
+ *
+ * <p>Outputs produced by this node that do not have an associated key in this map remain
+ * as-is. Generally, these are nodes which are internal to a composite replacement; they should
+ * not appear as outputs of the replacement composite.
+ *
+ * @param originalToReplacement A map from the outputs of the replacement {@link Node} to the
+ * original output.
+ */
+ void replaceOutputs(Map<PValue, ReplacementOutput> originalToReplacement) {
+ checkNotNull(this.outputs, "Outputs haven't been specified for node %s yet", getFullName());
+ for (Node component : this.parts) {
+ // Replace the outputs of the component nodes
+ component.replaceOutputs(originalToReplacement);
+ }
+ List<TaggedPValue> newOutputs = new ArrayList<>(outputs.size());
+ for (TaggedPValue output : outputs) {
+ ReplacementOutput mapping = originalToReplacement.get(output.getValue());
+ if (mapping != null) {
+ if (this.equals(producers.get(mapping.getReplacement().getValue()))) {
+ // This Node produced the replacement PCollection. The structure of this if statement
+ // requires the replacement transform to produce only new outputs; otherwise the
+ // producers map will not be appropriately updated. TODO: investigate alternatives
+ producers.put(mapping.getOriginal().getValue(), this);
+ producers.remove(mapping.getReplacement().getValue());
+ producerInput.remove(mapping.getReplacement().getValue());
+ }
+ LOG.debug(
+ "Replacing output {} with original {}",
+ mapping.getReplacement(),
+ mapping.getOriginal());
+ newOutputs.add(TaggedPValue.of(output.getTag(), mapping.getOriginal().getValue()));
+ } else {
+ newOutputs.add(output);
+ }
+ }
+ checkState(
+ outputs.size() == newOutputs.size(),
+ "Number of outputs must be stable across replacement");
+ this.outputs = newOutputs;
+ }
+
+ /** Returns the transform output, in expanded form. */
public List<TaggedPValue> getOutputs() {
- return output == null ? Collections.<TaggedPValue>emptyList() : output.expand();
+ return outputs == null ? Collections.<TaggedPValue>emptyList() : outputs;
}
/**
@@ -325,11 +428,7 @@ public class TransformHierarchy {
*/
public AppliedPTransform<?, ?, ?> toAppliedPTransform() {
return AppliedPTransform.of(
- getFullName(),
- input.expand(),
- output.expand(),
- (PTransform) getTransform(),
- input.getPipeline());
+ getFullName(), inputs, outputs, (PTransform) getTransform(), pipeline);
}
/**
@@ -345,7 +444,7 @@ public class TransformHierarchy {
if (!isRootNode()) {
// Visit inputs.
- for (TaggedPValue inputValue : input.expand()) {
+ for (TaggedPValue inputValue : inputs) {
if (visitedValues.add(inputValue.getValue())) {
visitor.visitValue(inputValue.getValue(), getProducer(inputValue.getValue()));
}
@@ -366,8 +465,9 @@ public class TransformHierarchy {
}
if (!isRootNode()) {
+ checkNotNull(outputs, "Outputs for non-root node %s are null", getFullName());
// Visit outputs.
- for (TaggedPValue pValue : output.expand()) {
+ for (TaggedPValue pValue : outputs) {
if (visitedValues.add(pValue.getValue())) {
visitor.visitValue(pValue.getValue(), this);
}
@@ -389,6 +489,9 @@ public class TransformHierarchy {
@Override
public String toString() {
+ if (isRootNode()) {
+ return "RootNode";
+ }
return MoreObjects.toStringHelper(getClass()).add("fullName", fullName).toString();
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/16c051ac/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TupleTagList.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TupleTagList.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TupleTagList.java
index 412678b..b4ce941 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TupleTagList.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TupleTagList.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.values;
+import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableList;
import java.io.Serializable;
import java.util.ArrayList;
@@ -145,4 +146,9 @@ public class TupleTagList implements Serializable {
TupleTagList(List<TupleTag<?>> tupleTags) {
this.tupleTags = Collections.unmodifiableList(tupleTags);
}
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(TupleTagList.class).add("tupleTags", tupleTags).toString();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/16c051ac/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
index 9a77b9b..426fd8b 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
@@ -17,32 +17,50 @@
*/
package org.apache.beam.sdk.runners;
+import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
import com.google.common.base.Function;
+import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import java.io.Serializable;
+import java.util.Collections;
import java.util.HashSet;
import java.util.List;
+import java.util.Map;
import java.util.Set;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
+import org.apache.beam.sdk.Pipeline.PipelineVisitor.Defaults;
+import org.apache.beam.sdk.io.CountingInput;
+import org.apache.beam.sdk.io.CountingInput.UnboundedCountingInput;
import org.apache.beam.sdk.io.CountingSource;
import org.apache.beam.sdk.io.Read;
+import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
+import org.apache.beam.sdk.runners.TransformHierarchy.Node;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.ParDo.Bound;
+import org.apache.beam.sdk.transforms.ParDo.BoundMulti;
import org.apache.beam.sdk.util.WindowingStrategy;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PCollectionList;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.PInput;
+import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
import org.hamcrest.Matchers;
import org.junit.Before;
import org.junit.Rule;
@@ -56,14 +74,16 @@ import org.junit.runners.JUnit4;
*/
@RunWith(JUnit4.class)
public class TransformHierarchyTest implements Serializable {
- @Rule public final transient TestPipeline pipeline = TestPipeline.create();
- @Rule public transient ExpectedException thrown = ExpectedException.none();
+ @Rule
+ public final transient TestPipeline pipeline =
+ TestPipeline.create().enableAbandonedNodeEnforcement(false);
+ @Rule public transient ExpectedException thrown = ExpectedException.none();
private transient TransformHierarchy hierarchy;
@Before
public void setup() {
- hierarchy = new TransformHierarchy();
+ hierarchy = new TransformHierarchy(pipeline);
}
@Test
@@ -111,7 +131,7 @@ public class TransformHierarchyTest implements Serializable {
hierarchy.popNode();
assertThat(hierarchy.getProducer(created), equalTo(node));
assertThat(
- "A Transform that produces non-primtive output should be composite",
+ "A Transform that produces non-primitive output should be composite",
emptyTransform.isCompositeNode(),
is(true));
}
@@ -174,6 +194,125 @@ public class TransformHierarchyTest implements Serializable {
}
@Test
+ public void replaceSucceeds() {
+ PTransform<?, ?> enclosingPT =
+ new PTransform<PInput, POutput>() {
+ @Override
+ public POutput expand(PInput input) {
+ return PDone.in(input.getPipeline());
+ }
+ };
+
+ TransformHierarchy.Node enclosing =
+ hierarchy.pushNode("Enclosing", PBegin.in(pipeline), enclosingPT);
+
+ Create.Values<Long> originalTransform = Create.of(1L);
+ TransformHierarchy.Node original =
+ hierarchy.pushNode("Create", PBegin.in(pipeline), originalTransform);
+ assertThat(hierarchy.getCurrent(), equalTo(original));
+ PCollection<Long> originalOutput = pipeline.apply(originalTransform);
+ hierarchy.setOutput(originalOutput);
+ hierarchy.popNode();
+ assertThat(original.finishedSpecifying, is(true));
+ hierarchy.setOutput(PDone.in(pipeline));
+ hierarchy.popNode();
+
+ assertThat(hierarchy.getCurrent(), not(equalTo(enclosing)));
+ 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);
+
+ TaggedPValue taggedOriginal = Iterables.getOnlyElement(originalOutput.expand());
+ TaggedPValue taggedReplacement = Iterables.getOnlyElement(replacementOutput.expand());
+ Map<PValue, ReplacementOutput> replacementOutputs =
+ Collections.<PValue, ReplacementOutput>singletonMap(
+ replacementOutput, ReplacementOutput.of(taggedOriginal, taggedReplacement));
+ hierarchy.replaceOutputs(replacementOutputs);
+
+ assertThat(replacement.getInputs(), equalTo(original.getInputs()));
+ assertThat(replacement.getEnclosingNode(), equalTo(original.getEnclosingNode()));
+ assertThat(replacement.getEnclosingNode(), equalTo(enclosing));
+ assertThat(
+ replacement.getTransform(), Matchers.<PTransform<?, ?>>equalTo(replacementTransform));
+ // THe tags of the replacement transform are matched to the appropriate PValues of the original
+ assertThat(
+ replacement.getOutputs(),
+ contains(TaggedPValue.of(taggedReplacement.getTag(), taggedOriginal.getValue())));
+ hierarchy.popNode();
+ }
+
+ @Test
+ public void replaceWithCompositeSucceeds() {
+ final ParDo.Bound<Long, Long> originalParDo =
+ ParDo.of(
+ new DoFn<Long, Long>() {
+ @ProcessElement
+ public void processElement(ProcessContext ctxt) {
+ ctxt.output(ctxt.element() + 1L);
+ }
+ });
+
+ UnboundedCountingInput genUpstream = CountingInput.unbounded();
+ PCollection<Long> upstream = pipeline.apply(genUpstream);
+ PCollection<Long> output = upstream.apply("Original", originalParDo);
+ hierarchy.pushNode("Upstream", pipeline.begin(), genUpstream);
+ hierarchy.finishSpecifyingInput();
+ hierarchy.setOutput(upstream);
+ hierarchy.popNode();
+
+ TransformHierarchy.Node original = hierarchy.pushNode("Original", upstream, originalParDo);
+ hierarchy.finishSpecifyingInput();
+ hierarchy.setOutput(output);
+ hierarchy.popNode();
+
+ final TupleTag<Long> longs = new TupleTag<>();
+ final ParDo.BoundMulti<Long, Long> replacementParDo =
+ ParDo.of(
+ new DoFn<Long, Long>() {
+ @ProcessElement
+ public void processElement(ProcessContext ctxt) {
+ ctxt.output(ctxt.element() + 1L);
+ }
+ })
+ .withOutputTags(longs, TupleTagList.empty());
+ PTransform<PCollection<Long>, PCollection<Long>> replacementComposite =
+ new PTransform<PCollection<Long>, PCollection<Long>>() {
+ @Override
+ public PCollection<Long> expand(PCollection<Long> input) {
+ return input.apply("Contained", replacementParDo).get(longs);
+ }
+ };
+
+ PCollectionTuple replacementOutput = upstream.apply("Contained", replacementParDo);
+
+ Node compositeNode = hierarchy.replaceNode(original, upstream, replacementComposite);
+ Node replacementParNode = hierarchy.pushNode("Original/Contained", upstream, replacementParDo);
+ hierarchy.finishSpecifyingInput();
+ hierarchy.setOutput(replacementOutput);
+ hierarchy.popNode();
+ hierarchy.setOutput(replacementOutput.get(longs));
+
+ TaggedPValue originalLongs = Iterables.getOnlyElement(output.expand());
+ TaggedPValue replacementLongs = Iterables.getOnlyElement(replacementOutput.expand());
+ hierarchy.replaceOutputs(
+ Collections.<PValue, ReplacementOutput>singletonMap(
+ replacementOutput.get(longs), ReplacementOutput.of(originalLongs, replacementLongs)));
+
+ assertThat(
+ replacementParNode.getOutputs(),
+ contains(TaggedPValue.of(replacementLongs.getTag(), originalLongs.getValue())));
+ assertThat(
+ compositeNode.getOutputs(),
+ contains(
+ TaggedPValue.of(
+ Iterables.getOnlyElement(replacementOutput.get(longs).expand()).getTag(),
+ originalLongs.getValue())));
+ hierarchy.popNode();
+ }
+
+ @Test
public void visitVisitsAllPushed() {
TransformHierarchy.Node root = hierarchy.getCurrent();
PBegin begin = PBegin.in(pipeline);
@@ -261,6 +400,96 @@ public class TransformHierarchyTest implements Serializable {
assertThat(visitedValuesInVisitor, equalTo(visitedValues));
}
+ /**
+ * Tests that visiting the {@link TransformHierarchy} after replacing nodes does not visit any
+ * of the original nodes or inaccessible values but does visit all of the replacement nodes,
+ * new inaccessible replacement values, and the original output values.
+ */
+ @Test
+ public void visitAfterReplace() {
+ Node root = hierarchy.getCurrent();
+ final Bound<Long, Long> originalParDo =
+ ParDo.of(
+ new DoFn<Long, Long>() {
+ @ProcessElement
+ public void processElement(ProcessContext ctxt) {
+ ctxt.output(ctxt.element() + 1L);
+ }
+ });
+
+ UnboundedCountingInput genUpstream = CountingInput.unbounded();
+ PCollection<Long> upstream = pipeline.apply(genUpstream);
+ PCollection<Long> output = upstream.apply("Original", originalParDo);
+ Node upstreamNode = hierarchy.pushNode("Upstream", pipeline.begin(), genUpstream);
+ hierarchy.finishSpecifyingInput();
+ hierarchy.setOutput(upstream);
+ hierarchy.popNode();
+
+ Node original = hierarchy.pushNode("Original", upstream, originalParDo);
+ hierarchy.finishSpecifyingInput();
+ hierarchy.setOutput(output);
+ hierarchy.popNode();
+
+ final TupleTag<Long> longs = new TupleTag<>();
+ final BoundMulti<Long, Long> replacementParDo =
+ ParDo.of(
+ new DoFn<Long, Long>() {
+ @ProcessElement
+ public void processElement(ProcessContext ctxt) {
+ ctxt.output(ctxt.element() + 1L);
+ }
+ })
+ .withOutputTags(longs, TupleTagList.empty());
+ PTransform<PCollection<Long>, PCollection<Long>> replacementComposite =
+ new PTransform<PCollection<Long>, PCollection<Long>>() {
+ @Override
+ public PCollection<Long> expand(PCollection<Long> input) {
+ return input.apply("Contained", replacementParDo).get(longs);
+ }
+ };
+
+ PCollectionTuple replacementOutput = upstream.apply("Contained", replacementParDo);
+
+ Node compositeNode = hierarchy.replaceNode(original, upstream, replacementComposite);
+ Node replacementParNode = hierarchy.pushNode("Original/Contained", upstream, replacementParDo);
+ hierarchy.finishSpecifyingInput();
+ hierarchy.setOutput(replacementOutput);
+ hierarchy.popNode();
+ hierarchy.setOutput(replacementOutput.get(longs));
+
+ TaggedPValue originalLongs = Iterables.getOnlyElement(output.expand());
+ TaggedPValue replacementLongs = Iterables.getOnlyElement(replacementOutput.expand());
+ hierarchy.replaceOutputs(
+ Collections.<PValue, ReplacementOutput>singletonMap(
+ replacementOutput.get(longs), ReplacementOutput.of(originalLongs, replacementLongs)));
+ hierarchy.popNode();
+
+ final Set<Node> visitedCompositeNodes = new HashSet<>();
+ final Set<Node> visitedPrimitiveNodes = new HashSet<>();
+ Set<PValue> visitedValues =
+ hierarchy.visit(
+ new Defaults() {
+ @Override
+ public CompositeBehavior enterCompositeTransform(Node node) {
+ visitedCompositeNodes.add(node);
+ return CompositeBehavior.ENTER_TRANSFORM;
+ }
+
+ @Override
+ public void visitPrimitiveTransform(Node node) {
+ visitedPrimitiveNodes.add(node);
+ }
+ });
+
+ /*
+ Final Graph:
+ Upstream -> Upstream.out -> Composite -> (ReplacementParDo -> OriginalParDo.out)
+ */
+ assertThat(visitedCompositeNodes, containsInAnyOrder(root, compositeNode));
+ assertThat(visitedPrimitiveNodes, containsInAnyOrder(upstreamNode, replacementParNode));
+ assertThat(visitedValues, Matchers.<PValue>containsInAnyOrder(upstream, output));
+ }
+
private static List<PValue> fromTaggedValues(List<TaggedPValue> taggedValues) {
return Lists.transform(
taggedValues,
[2/2] beam git commit: This closes #1998
Posted by tg...@apache.org.
This closes #1998
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/345c2acc
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/345c2acc
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/345c2acc
Branch: refs/heads/master
Commit: 345c2acca9b502fa40209de76d845118282d58c7
Parents: 938a66e 16c051a
Author: Thomas Groh <tg...@google.com>
Authored: Tue Feb 14 13:12:57 2017 -0800
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Feb 14 13:12:57 2017 -0800
----------------------------------------------------------------------
.../main/java/org/apache/beam/sdk/Pipeline.java | 63 ++++-
.../beam/sdk/runners/TransformHierarchy.java | 157 +++++++++---
.../apache/beam/sdk/values/TupleTagList.java | 6 +
.../sdk/runners/TransformHierarchyTest.java | 237 ++++++++++++++++++-
4 files changed, 431 insertions(+), 32 deletions(-)
----------------------------------------------------------------------