OutputDeduplicator { 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(); 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); }    }
@Test public void unchangedWithNoDuplicates() { PCollection redOut = PCollection.newBuilder().setUniqueName("red.out").build(); PTransform red = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putOutputs("out", redOut.getUniqueName()) .build(); PCollection oneOut = PCollection.newBuilder().setUniqueName("one.out").build(); PTransform one = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", oneOut.getUniqueName()) .build(); PCollection twoOut = PCollection.newBuilder().setUniqueName("two.out").build(); PTransform two = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", twoOut.getUniqueName()) .build(); PCollection blueOut = PCollection.newBuilder().setUniqueName("blue.out").build(); PTransform blue = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("one", oneOut.getUniqueName()) .putInputs("two", twoOut.getUniqueName()) .putOutputs("out", blueOut.getUniqueName()) .build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut) .putTransforms("two", two) .putPcollections(twoOut.getUniqueName(), twoOut) .putTransforms("red", red) .putPcollections(redOut.getUniqueName(), redOut) .putTransforms("blue", blue) .putPcollections(blueOut.getUniqueName(), blueOut) .build(); ExecutableStage oneStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(PipelineNode.pTransform("one", one)), ImmutableList.of(PipelineNode.pCollection(oneOut.getUniqueName(), oneOut)), DEFAULT_WIRE_CODER_SETTINGS); ExecutableStage twoStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(PipelineNode.pTransform("two", two)), ImmutableList.of(PipelineNode.pCollection(twoOut.getUniqueName(), twoOut)), DEFAULT_WIRE_CODER_SETTINGS); PTransformNode redTransform = PipelineNode.pTransform("red", red); PTransformNode blueTransform = PipelineNode.pTransform("blue", blue); QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components); DeduplicationResult result = OutputDeduplicator.ensureSingleProducer( pipeline, ImmutableList.of(oneStage, twoStage), ImmutableList.of(redTransform, blueTransform)); assertThat(result.getDeduplicatedComponents(), equalTo(components)); assertThat(result.getDeduplicatedStages().keySet(), empty()); assertThat(result.getDeduplicatedTransforms().keySet(), empty()); assertThat(result.getIntroducedTransforms(), empty()); }
@Test public void duplicateOverStages() { PCollection redOut = PCollection.newBuilder().setUniqueName("red.out").build(); PTransform red = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putOutputs("out", redOut.getUniqueName()) .build(); PCollection oneOut = PCollection.newBuilder().setUniqueName("one.out").build(); PTransform one = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", oneOut.getUniqueName()) .build(); PCollection twoOut = PCollection.newBuilder().setUniqueName("two.out").build(); PTransform two = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", twoOut.getUniqueName()) .build(); PCollection sharedOut = PCollection.newBuilder().setUniqueName("shared.out").build(); PTransform shared = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("one", oneOut.getUniqueName()) .putInputs("two", twoOut.getUniqueName()) .putOutputs("shared", sharedOut.getUniqueName()) .build(); PCollection blueOut = PCollection.newBuilder().setUniqueName("blue.out").build(); PTransform blue = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", sharedOut.getUniqueName()) .putOutputs("out", blueOut.getUniqueName()) .build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut) .putTransforms("two", two) .putPcollections(twoOut.getUniqueName(), twoOut) .putTransforms("shared", shared) .putPcollections(sharedOut.getUniqueName(), sharedOut) .putTransforms("red", red) .putPcollections(redOut.getUniqueName(), redOut) .putTransforms("blue", blue) .putPcollections(blueOut.getUniqueName(), blueOut) .build(); ExecutableStage oneStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of( PipelineNode.pTransform("one", one), PipelineNode.pTransform("shared", shared)), ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)), DEFAULT_WIRE_CODER_SETTINGS); ExecutableStage twoStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of( PipelineNode.pTransform("two", two), PipelineNode.pTransform("shared", shared)), ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)), DEFAULT_WIRE_CODER_SETTINGS); PTransformNode redTransform = PipelineNode.pTransform("red", red); PTransformNode blueTransform = PipelineNode.pTransform("blue", blue); QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components); DeduplicationResult result = OutputDeduplicator.ensureSingleProducer( pipeline, ImmutableList.of(oneStage, twoStage), ImmutableList.of(redTransform, blueTransform)); assertThat(result.getIntroducedTransforms(), hasSize(1)); PTransformNode introduced = getOnlyElement(result.getIntroducedTransforms()); assertThat(introduced.getTransform().getOutputsMap().size(), equalTo(1)); assertThat( getOnlyElement(introduced.getTransform().getOutputsMap().values()), equalTo(sharedOut.getUniqueName())); assertThat( result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introduced.getTransform().getInputsMap().values().toArray(new String[0]))); assertThat(result.getDeduplicatedStages().keySet(), hasSize(2)); List<String> stageOutputs = result.getDeduplicatedStages().values().stream() .flatMap(stage -> stage.getOutputPCollections().stream().map(PCollectionNode::getId)) .collect(Collectors.toList()); assertThat( stageOutputs, containsInAnyOrder(introduced.getTransform().getInputsMap().values().toArray())); assertThat(result.getDeduplicatedTransforms().keySet(), empty()); assertThat( result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(stageOutputs.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getTransformsMap(), hasEntry(introduced.getId(), introduced.getTransform())); }
@Test public void duplicateOverStagesAndTransforms() { PCollection redOut = PCollection.newBuilder().setUniqueName("red.out").build(); PTransform red = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putOutputs("out", redOut.getUniqueName()) .build(); PCollection oneOut = PCollection.newBuilder().setUniqueName("one.out").build(); PTransform one = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", oneOut.getUniqueName()) .build(); PCollection sharedOut = PCollection.newBuilder().setUniqueName("shared.out").build(); PTransform shared = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("one", oneOut.getUniqueName()) .putInputs("red", redOut.getUniqueName()) .putOutputs("shared", sharedOut.getUniqueName()) .build(); PCollection blueOut = PCollection.newBuilder().setUniqueName("blue.out").build(); PTransform blue = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", sharedOut.getUniqueName()) .putOutputs("out", blueOut.getUniqueName()) .build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut) .putTransforms("red", red) .putPcollections(redOut.getUniqueName(), redOut) .putTransforms("shared", shared) .putPcollections(sharedOut.getUniqueName(), sharedOut) .putTransforms("blue", blue) .putPcollections(blueOut.getUniqueName(), blueOut) .build(); PTransformNode sharedTransform = PipelineNode.pTransform("shared", shared); ExecutableStage oneStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(PipelineNode.pTransform("one", one), sharedTransform), ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)), DEFAULT_WIRE_CODER_SETTINGS); PTransformNode redTransform = PipelineNode.pTransform("red", red); PTransformNode blueTransform = PipelineNode.pTransform("blue", blue); QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components); DeduplicationResult result = OutputDeduplicator.ensureSingleProducer( pipeline, ImmutableList.of(oneStage), ImmutableList.of(redTransform, blueTransform, sharedTransform)); assertThat(result.getIntroducedTransforms(), hasSize(1)); PTransformNode introduced = getOnlyElement(result.getIntroducedTransforms()); assertThat(introduced.getTransform().getOutputsMap().size(), equalTo(1)); assertThat( getOnlyElement(introduced.getTransform().getOutputsMap().values()), equalTo(sharedOut.getUniqueName())); assertThat( result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introduced.getTransform().getInputsMap().values().toArray(new String[0]))); assertThat(result.getDeduplicatedStages().keySet(), hasSize(1)); assertThat(result.getDeduplicatedTransforms().keySet(), containsInAnyOrder("shared")); List<String> introducedOutputs = new ArrayList<>(); introducedOutputs.addAll( result.getDeduplicatedTransforms().get("shared").getTransform().getOutputsMap().values()); introducedOutputs.addAll( result.getDeduplicatedStages().get(oneStage).getOutputPCollections().stream() .map(PCollectionNode::getId) .collect(Collectors.toList())); assertThat( introduced.getTransform().getInputsMap().values(), containsInAnyOrder(introducedOutputs.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introducedOutputs.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getTransformsMap(), hasEntry(introduced.getId(), introduced.getTransform())); }
@Test public void multipleDuplicatesInStages() { PCollection redOut = PCollection.newBuilder().setUniqueName("red.out").build(); PTransform red = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putOutputs("out", redOut.getUniqueName()) .build(); PCollection threeOut = PCollection.newBuilder().setUniqueName("three.out").build(); PTransform three = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", threeOut.getUniqueName()) .build(); PCollection oneOut = PCollection.newBuilder().setUniqueName("one.out").build(); PTransform one = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", oneOut.getUniqueName()) .build(); PCollection twoOut = PCollection.newBuilder().setUniqueName("two.out").build(); PTransform two = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", redOut.getUniqueName()) .putOutputs("out", twoOut.getUniqueName()) .build(); PCollection sharedOut = PCollection.newBuilder().setUniqueName("shared.out").build(); PTransform shared = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("one", oneOut.getUniqueName()) .putInputs("two", twoOut.getUniqueName()) .putOutputs("shared", sharedOut.getUniqueName()) .build(); PCollection otherSharedOut = PCollection.newBuilder().setUniqueName("shared.out2").build(); PTransform otherShared = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("multi", threeOut.getUniqueName()) .putInputs("two", twoOut.getUniqueName()) .putOutputs("out", otherSharedOut.getUniqueName()) .build(); PCollection blueOut = PCollection.newBuilder().setUniqueName("blue.out").build(); PTransform blue = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("in", sharedOut.getUniqueName()) .putOutputs("out", blueOut.getUniqueName()) .build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut) .putTransforms("two", two) .putPcollections(twoOut.getUniqueName(), twoOut) .putTransforms("multi", three) .putPcollections(threeOut.getUniqueName(), threeOut) .putTransforms("shared", shared) .putPcollections(sharedOut.getUniqueName(), sharedOut) .putTransforms("otherShared", otherShared) .putPcollections(otherSharedOut.getUniqueName(), otherSharedOut) .putTransforms("red", red) .putPcollections(redOut.getUniqueName(), redOut) .putTransforms("blue", blue) .putPcollections(blueOut.getUniqueName(), blueOut) .build(); ExecutableStage multiStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of( PipelineNode.pTransform("multi", three), PipelineNode.pTransform("shared", shared), PipelineNode.pTransform("otherShared", otherShared)), ImmutableList.of( PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut), PipelineNode.pCollection(otherSharedOut.getUniqueName(), otherSharedOut)), DEFAULT_WIRE_CODER_SETTINGS); ExecutableStage oneStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of( PipelineNode.pTransform("one", one), PipelineNode.pTransform("shared", shared)), ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)), DEFAULT_WIRE_CODER_SETTINGS); ExecutableStage twoStage = ImmutableExecutableStage.of( components, Environment.getDefaultInstance(), PipelineNode.pCollection(redOut.getUniqueName(), redOut), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of( PipelineNode.pTransform("two", two), PipelineNode.pTransform("otherShared", otherShared)), ImmutableList.of( PipelineNode.pCollection(otherSharedOut.getUniqueName(), otherSharedOut)), DEFAULT_WIRE_CODER_SETTINGS); PTransformNode redTransform = PipelineNode.pTransform("red", red); PTransformNode blueTransform = PipelineNode.pTransform("blue", blue); QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components); DeduplicationResult result = OutputDeduplicator.ensureSingleProducer( pipeline, ImmutableList.of(oneStage, twoStage, multiStage), ImmutableList.of(redTransform, blueTransform)); assertThat(result.getIntroducedTransforms(), hasSize(2)); assertThat( result.getDeduplicatedStages().keySet(), containsInAnyOrder(multiStage, oneStage, twoStage)); assertThat(result.getDeduplicatedTransforms().keySet(), empty()); Collection<String> introducedIds = result.getIntroducedTransforms().stream() .flatMap(pt -> pt.getTransform().getInputsMap().values().stream()) .collect(Collectors.toList()); String[] stageOutputs = result.getDeduplicatedStages().values().stream() .flatMap(s -> s.getOutputPCollections().stream().map(PCollectionNode::getId)) .toArray(String[]::new); assertThat(introducedIds, containsInAnyOrder(stageOutputs)); assertThat( result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introducedIds.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getTransformsMap().entrySet(), hasItems( result.getIntroducedTransforms().stream() .collect(Collectors.toMap(PTransformNode::getId, PTransformNode::getTransform)) .entrySet() .toArray(new Map.Entry[0]))); }