@Override
public void visitValue(PValue value, TransformHierarchy.Node producer) {
boolean inputsAreKeyed = true;
for (TaggedPValue input : producer.getInputs()) {
inputsAreKeyed = inputsAreKeyed && keyedValues.contains(input.getValue());
}
if (PRODUCES_KEYED_OUTPUTS.contains(producer.getTransform().getClass())
|| (isKeyPreserving(producer.getTransform()) && inputsAreKeyed)) {
keyedValues.add(value);
}
}
@Override
public void leaveCompositeTransform(TransformHierarchy.Node node) {
checkState(
!finalized,
"Attempted to use a %s that has already been finalized on a pipeline (visiting node %s)",
KeyedPValueTrackingVisitor.class.getSimpleName(),
node);
if (node.isRootNode()) {
finalized = true;
} else if (PRODUCES_KEYED_OUTPUTS.contains(node.getTransform().getClass())) {
List<TaggedPValue> outputs = node.getOutputs();
for (TaggedPValue output : outputs) {
keyedValues.add(output.getValue());
}
}
}