com.google.cloud.dataflow.sdk.runners.TransformTreeNode Maven / Gradle / Ivy
Show all versions of google-cloud-dataflow-java-sdk-all Show documentation
/*
* Copyright (C) 2015 Google Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
* use this file except in compliance with the License. You may obtain a copy of
* the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations under
* the License.
*/
package com.google.cloud.dataflow.sdk.runners;
import com.google.cloud.dataflow.sdk.Pipeline;
import com.google.cloud.dataflow.sdk.transforms.PTransform;
import com.google.cloud.dataflow.sdk.values.PInput;
import com.google.cloud.dataflow.sdk.values.POutput;
import com.google.cloud.dataflow.sdk.values.PValue;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import javax.annotation.Nullable;
/**
* Provides internal tracking of transform relationships with helper methods
* for initialization and ordered visitation.
*/
public class TransformTreeNode {
private final TransformTreeNode enclosingNode;
// The PTransform for this node, which may be a composite PTransform.
// The root of a TransformHierarchy is represented as a TransformTreeNode
// with a null transform field.
private final PTransform transform;
private final String fullName;
// Nodes for sub-transforms of a composite transform.
private final Collection parts = new ArrayList<>();
// Inputs to the transform, in expanded form and mapped to the producer
// of the input.
private final Map inputs = new HashMap<>();
// Input to the transform, in unexpanded form.
private final PInput input;
// TODO: track which outputs need to be exported to parent.
// Output of the transform, in unexpanded form.
private POutput output;
private boolean finishedSpecifying = false;
/**
* Creates a new TransformTreeNode with the given parent and transform.
*
* EnclosingNode and transform may both be null for
* a root-level node, which holds all other nodes.
*
* @param enclosingNode the composite node containing this node
* @param transform the PTransform tracked by this node
* @param fullName the fully qualified name of the transform
* @param input the unexpanded input to the transform
*/
public TransformTreeNode(@Nullable TransformTreeNode enclosingNode,
@Nullable PTransform transform,
String fullName,
@Nullable PInput input) {
this.enclosingNode = enclosingNode;
this.transform = transform;
Preconditions.checkArgument((enclosingNode == null && transform == null)
|| (enclosingNode != null && transform != null),
"EnclosingNode and transform must both be specified, or both be null");
this.fullName = fullName;
this.input = input;
}
/**
* Returns the transform associated with this transform node.
*/
public PTransform getTransform() {
return transform;
}
/**
* Returns the enclosing composite transform node, or null if there is none.
*/
public TransformTreeNode getEnclosingNode() {
return enclosingNode;
}
/**
* Adds a composite operation to the transform node.
*
*
As soon as a node is added, the transform node is considered a
* composite operation instead of a primitive transform.
*/
public void addComposite(TransformTreeNode node) {
parts.add(node);
}
/**
* 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).
*
*
Note that a node may be composite with no sub-transforms if it returns its input directly
* extracts a component of a tuple, or other operations that occur at pipeline assembly time.
*/
public boolean isCompositeNode() {
return !parts.isEmpty() || returnsOthersOutput() || isRootNode();
}
private boolean returnsOthersOutput() {
PTransform transform = getTransform();
for (PValue output : getExpandedOutputs()) {
if (!output.getProducingTransformInternal().getTransform().equals(transform)) {
return true;
}
}
return false;
}
public boolean isRootNode() {
return transform == null;
}
public String getFullName() {
return fullName;
}
/**
* Adds an input to the transform node.
*/
public void addInputProducer(PValue expandedInput, TransformTreeNode producer) {
Preconditions.checkState(!finishedSpecifying);
inputs.put(expandedInput, producer);
}
/**
* Returns the transform input, in unexpanded form.
*/
public PInput getInput() {
return input;
}
/**
* Returns a mapping of inputs to the producing nodes for all inputs to
* the transform.
*/
public Map getInputs() {
return Collections.unmodifiableMap(inputs);
}
/**
* Adds an output to the transform node.
*/
public void setOutput(POutput output) {
Preconditions.checkState(!finishedSpecifying);
Preconditions.checkState(this.output == null);
this.output = output;
}
/**
* Returns the transform output, in unexpanded form.
*/
public POutput getOutput() {
return output;
}
/**
* Returns the transform outputs, in expanded form.
*/
public Collection getExpandedOutputs() {
if (output != null) {
return output.expand();
} else {
return Collections.emptyList();
}
}
/**
* Visit the transform node.
*
* Provides an ordered visit of the input values, the primitive
* transform (or child nodes for composite transforms), then the
* output values.
*/
public void visit(Pipeline.PipelineVisitor visitor,
Set visitedValues) {
if (!finishedSpecifying) {
finishSpecifying();
}
// Visit inputs.
for (Map.Entry entry : inputs.entrySet()) {
if (visitedValues.add(entry.getKey())) {
visitor.visitValue(entry.getKey(), entry.getValue());
}
}
if (isCompositeNode()) {
visitor.enterCompositeTransform(this);
for (TransformTreeNode child : parts) {
child.visit(visitor, visitedValues);
}
visitor.leaveCompositeTransform(this);
} else {
visitor.visitTransform(this);
}
// Visit outputs.
for (PValue pValue : getExpandedOutputs()) {
if (visitedValues.add(pValue)) {
visitor.visitValue(pValue, this);
}
}
}
/**
* Finish specifying a transform.
*
* All inputs are finished first, then the transform, then
* all outputs.
*/
public void finishSpecifying() {
if (finishedSpecifying) {
return;
}
finishedSpecifying = true;
for (TransformTreeNode input : inputs.values()) {
if (input != null) {
input.finishSpecifying();
}
}
if (output != null) {
output.finishSpecifyingOutput();
}
}
}