use of org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode in project beam by apache.
the class ExecutableStage method fromPayload.
/**
* Return an {@link ExecutableStage} constructed from the provided {@link FunctionSpec}
* representation.
*
* <p>See {@link #toPTransform} for how the payload is constructed.
*
* <p>Note: The payload contains some information redundant with the {@link PTransform} it is the
* payload of. The {@link ExecutableStagePayload} should be sufficiently rich to construct a
* {@code ProcessBundleDescriptor} using only the payload.
*/
static ExecutableStage fromPayload(ExecutableStagePayload payload) {
Components components = payload.getComponents();
Environment environment = payload.getEnvironment();
Collection<WireCoderSetting> wireCoderSettings = payload.getWireCoderSettingsList();
PCollectionNode input = PipelineNode.pCollection(payload.getInput(), components.getPcollectionsOrThrow(payload.getInput()));
List<SideInputReference> sideInputs = payload.getSideInputsList().stream().map(sideInputId -> SideInputReference.fromSideInputId(sideInputId, components)).collect(Collectors.toList());
List<UserStateReference> userStates = payload.getUserStatesList().stream().map(userStateId -> UserStateReference.fromUserStateId(userStateId, components)).collect(Collectors.toList());
List<TimerReference> timers = payload.getTimersList().stream().map(timerId -> TimerReference.fromTimerId(timerId, components)).collect(Collectors.toList());
List<PTransformNode> transforms = payload.getTransformsList().stream().map(id -> PipelineNode.pTransform(id, components.getTransformsOrThrow(id))).collect(Collectors.toList());
List<PCollectionNode> outputs = payload.getOutputsList().stream().map(id -> PipelineNode.pCollection(id, components.getPcollectionsOrThrow(id))).collect(Collectors.toList());
return ImmutableExecutableStage.of(components, environment, input, sideInputs, userStates, timers, transforms, outputs, wireCoderSettings);
}
use of org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode in project beam by apache.
the class QueryablePipeline method buildNetwork.
private MutableNetwork<PipelineNode, PipelineEdge> buildNetwork(Collection<String> transformIds, Components components) {
MutableNetwork<PipelineNode, PipelineEdge> network = NetworkBuilder.directed().allowsParallelEdges(true).allowsSelfLoops(false).build();
Set<PCollectionNode> unproducedCollections = new HashSet<>();
for (String transformId : transformIds) {
PTransform transform = components.getTransformsOrThrow(transformId);
PTransformNode transformNode = PipelineNode.pTransform(transformId, this.components.getTransformsOrThrow(transformId));
network.addNode(transformNode);
for (String produced : transform.getOutputsMap().values()) {
PCollectionNode producedNode = PipelineNode.pCollection(produced, components.getPcollectionsOrThrow(produced));
network.addNode(producedNode);
network.addEdge(transformNode, producedNode, new PerElementEdge());
checkArgument(network.inDegree(producedNode) == 1, "A %s should have exactly one producing %s, but found %s:\nPCollection:\n%s\nProducers:\n%s", PCollectionNode.class.getSimpleName(), PTransformNode.class.getSimpleName(), network.predecessors(producedNode).size(), producedNode, network.predecessors(producedNode));
unproducedCollections.remove(producedNode);
}
for (Map.Entry<String, String> consumed : transform.getInputsMap().entrySet()) {
// This loop may add an edge between the consumed PCollection and the current PTransform.
// The local name of the transform must be used to determine the type of edge.
String pcollectionId = consumed.getValue();
PCollectionNode consumedNode = PipelineNode.pCollection(pcollectionId, this.components.getPcollectionsOrThrow(pcollectionId));
if (network.addNode(consumedNode)) {
// This node has been added to the network for the first time, so it has no producer.
unproducedCollections.add(consumedNode);
}
if (getLocalSideInputNames(transform).contains(consumed.getKey())) {
network.addEdge(consumedNode, transformNode, new SingletonEdge());
} else {
network.addEdge(consumedNode, transformNode, new PerElementEdge());
}
}
}
checkArgument(unproducedCollections.isEmpty(), "%ss %s were consumed but never produced", PCollectionNode.class.getSimpleName(), unproducedCollections);
return network;
}
use of org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode in project beam by apache.
the class OutputDeduplicator method deduplicateStageOutput.
/**
* Returns an {@link ExecutableStage} where all of the {@link PCollectionNode PCollections}
* matching the original are replaced with the introduced partial {@link PCollection} in all
* references made within the {@link ExecutableStage}.
*/
private static ExecutableStage deduplicateStageOutput(ExecutableStage stage, Map<String, PCollectionNode> originalToPartial) {
Collection<PTransformNode> updatedTransforms = new ArrayList<>();
for (PTransformNode transform : stage.getTransforms()) {
PTransform updatedTransform = updateOutputs(transform.getTransform(), originalToPartial);
updatedTransforms.add(PipelineNode.pTransform(transform.getId(), updatedTransform));
}
Collection<PCollectionNode> updatedOutputs = new ArrayList<>();
for (PCollectionNode output : stage.getOutputPCollections()) {
updatedOutputs.add(originalToPartial.getOrDefault(output.getId(), output));
}
RunnerApi.Components updatedStageComponents = stage.getComponents().toBuilder().clearTransforms().putAllTransforms(updatedTransforms.stream().collect(Collectors.toMap(PTransformNode::getId, PTransformNode::getTransform))).putAllPcollections(originalToPartial.values().stream().collect(Collectors.toMap(PCollectionNode::getId, PCollectionNode::getPCollection))).build();
return ImmutableExecutableStage.of(updatedStageComponents, stage.getEnvironment(), stage.getInputPCollection(), stage.getSideInputs(), stage.getUserStates(), stage.getTimers(), updatedTransforms, updatedOutputs, stage.getWireCoderSettings());
}
use of org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode in project beam by apache.
the class OutputDeduplicator method ensureSingleProducer.
/**
* Ensure that no {@link PCollection} output by any of the {@code stages} or {@code
* unfusedTransforms} is produced by more than one of those stages or transforms.
*
* <p>For each {@link PCollection} output by multiple stages and/or transforms, each producer is
* rewritten to produce a partial {@link PCollection}, which are then flattened together via an
* introduced Flatten node which produces the original output.
*/
static DeduplicationResult ensureSingleProducer(QueryablePipeline pipeline, Collection<ExecutableStage> stages, Collection<PTransformNode> unfusedTransforms) {
RunnerApi.Components.Builder unzippedComponents = pipeline.getComponents().toBuilder();
Multimap<PCollectionNode, StageOrTransform> pcollectionProducers = getProducers(pipeline, stages, unfusedTransforms);
Multimap<StageOrTransform, PCollectionNode> requiresNewOutput = HashMultimap.create();
// ExecutableStage must also be rewritten to have updated outputs and transforms.
for (Map.Entry<PCollectionNode, Collection<StageOrTransform>> collectionProducer : pcollectionProducers.asMap().entrySet()) {
if (collectionProducer.getValue().size() > 1) {
for (StageOrTransform producer : collectionProducer.getValue()) {
requiresNewOutput.put(producer, collectionProducer.getKey());
}
}
}
Map<ExecutableStage, ExecutableStage> updatedStages = new LinkedHashMap<>();
Map<String, PTransformNode> updatedTransforms = new LinkedHashMap<>();
Multimap<String, PCollectionNode> originalToPartial = HashMultimap.create();
for (Map.Entry<StageOrTransform, Collection<PCollectionNode>> deduplicationTargets : requiresNewOutput.asMap().entrySet()) {
if (deduplicationTargets.getKey().getStage() != null) {
StageDeduplication deduplication = deduplicatePCollections(deduplicationTargets.getKey().getStage(), deduplicationTargets.getValue(), unzippedComponents::containsPcollections);
for (Entry<String, PCollectionNode> originalToPartialReplacement : deduplication.getOriginalToPartialPCollections().entrySet()) {
originalToPartial.put(originalToPartialReplacement.getKey(), originalToPartialReplacement.getValue());
unzippedComponents.putPcollections(originalToPartialReplacement.getValue().getId(), originalToPartialReplacement.getValue().getPCollection());
}
updatedStages.put(deduplicationTargets.getKey().getStage(), deduplication.getUpdatedStage());
} else if (deduplicationTargets.getKey().getTransform() != null) {
PTransformDeduplication deduplication = deduplicatePCollections(deduplicationTargets.getKey().getTransform(), deduplicationTargets.getValue(), unzippedComponents::containsPcollections);
for (Entry<String, PCollectionNode> originalToPartialReplacement : deduplication.getOriginalToPartialPCollections().entrySet()) {
originalToPartial.put(originalToPartialReplacement.getKey(), originalToPartialReplacement.getValue());
unzippedComponents.putPcollections(originalToPartialReplacement.getValue().getId(), originalToPartialReplacement.getValue().getPCollection());
}
updatedTransforms.put(deduplicationTargets.getKey().getTransform().getId(), deduplication.getUpdatedTransform());
} else {
throw new IllegalStateException(String.format("%s with no %s or %s", StageOrTransform.class.getSimpleName(), ExecutableStage.class.getSimpleName(), PTransformNode.class.getSimpleName()));
}
}
Set<PTransformNode> introducedFlattens = new LinkedHashSet<>();
for (Map.Entry<String, Collection<PCollectionNode>> partialFlattenTargets : originalToPartial.asMap().entrySet()) {
String flattenId = SyntheticComponents.uniqueId("unzipped_flatten", unzippedComponents::containsTransforms);
PTransform flattenPartialPCollections = createFlattenOfPartials(flattenId, partialFlattenTargets.getKey(), partialFlattenTargets.getValue());
unzippedComponents.putTransforms(flattenId, flattenPartialPCollections);
introducedFlattens.add(PipelineNode.pTransform(flattenId, flattenPartialPCollections));
}
Components components = unzippedComponents.build();
return DeduplicationResult.of(components, introducedFlattens, updatedStages, updatedTransforms);
}
use of org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode in project beam by apache.
the class OutputDeduplicator method deduplicatePCollections.
private static PTransformDeduplication deduplicatePCollections(PTransformNode transform, Collection<PCollectionNode> duplicates, Predicate<String> existingPCollectionIds) {
Map<String, PCollectionNode> unzippedOutputs = createPartialPCollections(duplicates, existingPCollectionIds);
PTransform pTransform = updateOutputs(transform.getTransform(), unzippedOutputs);
return PTransformDeduplication.of(PipelineNode.pTransform(transform.getId(), pTransform), unzippedOutputs);
}
Aggregations