@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()); } } }