private AppliedPTransform<?, ?, ?> getAppliedTransform(TransformTreeNode node) {
@SuppressWarnings({"rawtypes", "unchecked"})
AppliedPTransform<?, ?, ?> application =
AppliedPTransform.of(
node.getFullName(),
node.getInput(),
node.getOutput(),
(PTransform) node.getTransform());
return application;
}
@Override
public void leaveCompositeTransform(TransformTreeNode node) {
checkState(
!finalized,
"Attempting to traverse a pipeline (node %s) with a %s which is already finalized",
node.getFullName(),
ConsumerTrackingPipelineVisitor.class.getSimpleName());
if (node.isRootNode()) {
finalized = true;
}
}
@Override
public void visitTransform(TransformTreeNode node) {
toFinalize.removeAll(node.getInput().expand());
AppliedPTransform<?, ?, ?> appliedTransform = getAppliedTransform(node);
stepNames.put(appliedTransform, genStepName());
if (node.getInput().expand().isEmpty()) {
rootTransforms.add(appliedTransform);
} else {
for (PValue value : node.getInput().expand()) {
valueToConsumers.get(value).add(appliedTransform);
}
}
}
@Override
public void enterCompositeTransform(TransformTreeNode node) {
checkState(
!finalized,
"Attempting to traverse a pipeline (node %s) with a %s "
+ "which has already visited a Pipeline and is finalized",
node.getFullName(),
ConsumerTrackingPipelineVisitor.class.getSimpleName());
}
