| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.runners.core.construction.graph; |
| |
| import com.google.auto.value.AutoValue; |
| import java.io.IOException; |
| import java.util.Arrays; |
| import java.util.Map; |
| import java.util.function.Predicate; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.Coder; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ComponentsOrBuilder; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload; |
| import org.apache.beam.runners.core.construction.ModelCoders; |
| import org.apache.beam.runners.core.construction.PTransformTranslation; |
| import org.apache.beam.runners.core.construction.ParDoTranslation; |
| import org.apache.beam.runners.core.construction.graph.ProtoOverrides.TransformReplacement; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps; |
| |
| /** |
| * A set of transform replacements for expanding a splittable ParDo into various sub components. |
| * |
| * <p>Further details about the expansion can be found at <a |
| * href="https://github.com/apache/beam/blob/cb15994d5228f729dda922419b08520c8be8804e/model/pipeline/src/main/proto/beam_runner_api.proto#L279" |
| * /> |
| */ |
| public class SplittableParDoExpander { |
| |
| /** |
| * Returns a transform replacement which expands a splittable ParDo from: |
| * |
| * <pre>{@code |
| * sideInputA ---------\ |
| * sideInputB ---------V |
| * mainInput ---> SplittableParDo --> outputA |
| * \-> outputB |
| * }</pre> |
| * |
| * into: |
| * |
| * <pre>{@code |
| * sideInputA ---------\---------------------\--------------------------\ |
| * sideInputB ---------V---------------------V--------------------------V |
| * mainInput ---> PairWithRestricton --> SplitAndSize --> ProcessSizedElementsAndRestriction --> outputA |
| * \-> outputB |
| * }</pre> |
| * |
| * <p>Specifically this transform ensures that initial splitting is performed and that the sizing |
| * information is available to the runner if it chooses to inspect it. |
| */ |
| public static TransformReplacement createSizedReplacement() { |
| return SizedReplacement.builder().setDrain(false).build(); |
| } |
| |
| /** |
| * Returns a transform replacement in drain mode which expands a splittable ParDo from: |
| * |
| * <pre>{@code |
| * sideInputA ---------\ |
| * sideInputB ---------V |
| * mainInput ---> SplittableParDo --> outputA |
| * \-> outputB |
| * }</pre> |
| * |
| * into: |
| * |
| * <pre>{@code |
| * sideInputA ---------\---------------------\----------------------\--------------------------\ |
| * sideInputB ---------V---------------------V----------------------V--------------------------V |
| * mainInput ---> PairWithRestriction --> SplitAndSize --> TruncateAndSize --> ProcessSizedElementsAndRestriction --> outputA |
| * \-> outputB |
| * }</pre> |
| * |
| * . |
| */ |
| public static TransformReplacement createTruncateReplacement() { |
| return SizedReplacement.builder().setDrain(true).build(); |
| } |
| |
| /** See {@link #createSizedReplacement()} for details. */ |
| @AutoValue |
| abstract static class SizedReplacement implements TransformReplacement { |
| |
| static Builder builder() { |
| return new AutoValue_SplittableParDoExpander_SizedReplacement.Builder(); |
| } |
| |
| abstract boolean isDrain(); |
| |
| @AutoValue.Builder |
| abstract static class Builder { |
| abstract Builder setDrain(boolean isDrain); |
| |
| abstract SizedReplacement build(); |
| } |
| |
| @Override |
| public MessageWithComponents getReplacement( |
| String transformId, ComponentsOrBuilder existingComponents) { |
| try { |
| MessageWithComponents.Builder rval = MessageWithComponents.newBuilder(); |
| |
| PTransform splittableParDo = existingComponents.getTransformsOrThrow(transformId); |
| ParDoPayload payload = ParDoPayload.parseFrom(splittableParDo.getSpec().getPayload()); |
| // Only perform the expansion if this is a splittable DoFn. |
| if (payload.getRestrictionCoderId() == null || payload.getRestrictionCoderId().isEmpty()) { |
| return null; |
| } |
| |
| String mainInputName = ParDoTranslation.getMainInputName(splittableParDo); |
| String mainInputPCollectionId = splittableParDo.getInputsOrThrow(mainInputName); |
| PCollection mainInputPCollection = |
| existingComponents.getPcollectionsOrThrow(mainInputPCollectionId); |
| Map<String, String> sideInputs = |
| Maps.filterKeys( |
| splittableParDo.getInputsMap(), input -> payload.containsSideInputs(input)); |
| |
| String pairWithRestrictionOutCoderId = |
| generateUniqueId( |
| mainInputPCollection.getCoderId() + "/PairWithRestriction", |
| existingComponents::containsCoders); |
| rval.getComponentsBuilder() |
| .putCoders( |
| pairWithRestrictionOutCoderId, |
| ModelCoders.kvCoder( |
| mainInputPCollection.getCoderId(), payload.getRestrictionCoderId())); |
| |
| String pairWithRestrictionOutId = |
| generateUniqueId( |
| mainInputPCollectionId + "/PairWithRestriction", |
| existingComponents::containsPcollections); |
| rval.getComponentsBuilder() |
| .putPcollections( |
| pairWithRestrictionOutId, |
| PCollection.newBuilder() |
| .setCoderId(pairWithRestrictionOutCoderId) |
| .setIsBounded(mainInputPCollection.getIsBounded()) |
| .setWindowingStrategyId(mainInputPCollection.getWindowingStrategyId()) |
| .setUniqueName( |
| generateUniquePCollectonName( |
| mainInputPCollection.getUniqueName() + "/PairWithRestriction", |
| existingComponents)) |
| .build()); |
| |
| String splitAndSizeOutCoderId = |
| generateUniqueId( |
| mainInputPCollection.getCoderId() + "/SplitAndSize", |
| existingComponents::containsCoders); |
| rval.getComponentsBuilder() |
| .putCoders( |
| splitAndSizeOutCoderId, |
| ModelCoders.kvCoder( |
| pairWithRestrictionOutCoderId, getOrAddDoubleCoder(existingComponents, rval))); |
| |
| String splitAndSizeOutId = |
| generateUniqueId( |
| mainInputPCollectionId + "/SplitAndSize", existingComponents::containsPcollections); |
| rval.getComponentsBuilder() |
| .putPcollections( |
| splitAndSizeOutId, |
| PCollection.newBuilder() |
| .setCoderId(splitAndSizeOutCoderId) |
| .setIsBounded(mainInputPCollection.getIsBounded()) |
| .setWindowingStrategyId(mainInputPCollection.getWindowingStrategyId()) |
| .setUniqueName( |
| generateUniquePCollectonName( |
| mainInputPCollection.getUniqueName() + "/SplitAndSize", |
| existingComponents)) |
| .build()); |
| |
| String pairWithRestrictionId = |
| generateUniqueId( |
| transformId + "/PairWithRestriction", existingComponents::containsTransforms); |
| { |
| PTransform.Builder pairWithRestriction = PTransform.newBuilder(); |
| pairWithRestriction.putAllInputs(splittableParDo.getInputsMap()); |
| pairWithRestriction.putOutputs("out", pairWithRestrictionOutId); |
| pairWithRestriction.setUniqueName( |
| generateUniquePCollectonName( |
| splittableParDo.getUniqueName() + "/PairWithRestriction", existingComponents)); |
| pairWithRestriction.setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN) |
| .setPayload(splittableParDo.getSpec().getPayload())); |
| pairWithRestriction.setEnvironmentId(splittableParDo.getEnvironmentId()); |
| rval.getComponentsBuilder() |
| .putTransforms(pairWithRestrictionId, pairWithRestriction.build()); |
| } |
| |
| String splitAndSizeId = |
| generateUniqueId(transformId + "/SplitAndSize", existingComponents::containsTransforms); |
| { |
| PTransform.Builder splitAndSize = PTransform.newBuilder(); |
| splitAndSize.putInputs(mainInputName, pairWithRestrictionOutId); |
| splitAndSize.putAllInputs(sideInputs); |
| splitAndSize.putOutputs("out", splitAndSizeOutId); |
| splitAndSize.setUniqueName( |
| generateUniquePCollectonName( |
| splittableParDo.getUniqueName() + "/SplitAndSize", existingComponents)); |
| splitAndSize.setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN) |
| .setPayload(splittableParDo.getSpec().getPayload())); |
| splitAndSize.setEnvironmentId(splittableParDo.getEnvironmentId()); |
| rval.getComponentsBuilder().putTransforms(splitAndSizeId, splitAndSize.build()); |
| } |
| PTransform.Builder newCompositeRoot = |
| splittableParDo |
| .toBuilder() |
| // Clear the original splittable ParDo spec and add all the new transforms as |
| // children. |
| .clearSpec() |
| .addAllSubtransforms(Arrays.asList(pairWithRestrictionId, splitAndSizeId)); |
| |
| String processSizedElementsAndRestrictionsId = |
| generateUniqueId( |
| transformId + "/ProcessSizedElementsAndRestrictions", |
| existingComponents::containsTransforms); |
| String processSizedElementsInputPCollectionId = splitAndSizeOutId; |
| if (isDrain()) { |
| String truncateAndSizeCoderId = |
| generateUniqueId( |
| mainInputPCollection.getCoderId() + "/TruncateAndSize", |
| existingComponents::containsCoders); |
| rval.getComponentsBuilder() |
| .putCoders( |
| truncateAndSizeCoderId, |
| ModelCoders.kvCoder( |
| splitAndSizeOutCoderId, getOrAddDoubleCoder(existingComponents, rval))); |
| String truncateAndSizeOutId = |
| generateUniqueId( |
| mainInputPCollectionId + "/TruncateAndSize", |
| existingComponents::containsPcollections); |
| |
| rval.getComponentsBuilder() |
| .putPcollections( |
| truncateAndSizeOutId, |
| PCollection.newBuilder() |
| .setCoderId(truncateAndSizeCoderId) |
| .setIsBounded(mainInputPCollection.getIsBounded()) |
| .setWindowingStrategyId(mainInputPCollection.getWindowingStrategyId()) |
| .setUniqueName( |
| generateUniquePCollectonName( |
| mainInputPCollection.getUniqueName() + "/TruncateAndSize", |
| existingComponents)) |
| .build()); |
| String truncateAndSizeId = |
| generateUniqueId( |
| transformId + "/TruncateAndSize", existingComponents::containsTransforms); |
| { |
| PTransform.Builder truncateAndSize = PTransform.newBuilder(); |
| truncateAndSize.putInputs(mainInputName, splitAndSizeOutId); |
| truncateAndSize.putAllInputs(sideInputs); |
| truncateAndSize.putOutputs("out", truncateAndSizeOutId); |
| truncateAndSize.setUniqueName( |
| generateUniquePCollectonName( |
| splittableParDo.getUniqueName() + "/TruncateAndSize", existingComponents)); |
| truncateAndSize.setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(PTransformTranslation.SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN) |
| .setPayload(splittableParDo.getSpec().getPayload())); |
| truncateAndSize.setEnvironmentId(splittableParDo.getEnvironmentId()); |
| rval.getComponentsBuilder().putTransforms(truncateAndSizeId, truncateAndSize.build()); |
| } |
| newCompositeRoot.addSubtransforms(truncateAndSizeId); |
| processSizedElementsInputPCollectionId = truncateAndSizeOutId; |
| } |
| { |
| PTransform.Builder processSizedElementsAndRestrictions = PTransform.newBuilder(); |
| processSizedElementsAndRestrictions.putInputs( |
| mainInputName, processSizedElementsInputPCollectionId); |
| processSizedElementsAndRestrictions.putAllInputs(sideInputs); |
| processSizedElementsAndRestrictions.putAllOutputs(splittableParDo.getOutputsMap()); |
| processSizedElementsAndRestrictions.setUniqueName( |
| generateUniquePCollectonName( |
| splittableParDo.getUniqueName() + "/ProcessSizedElementsAndRestrictions", |
| existingComponents)); |
| processSizedElementsAndRestrictions.setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn( |
| PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN) |
| .setPayload(splittableParDo.getSpec().getPayload())); |
| processSizedElementsAndRestrictions.setEnvironmentId(splittableParDo.getEnvironmentId()); |
| rval.getComponentsBuilder() |
| .putTransforms( |
| processSizedElementsAndRestrictionsId, |
| processSizedElementsAndRestrictions.build()); |
| } |
| newCompositeRoot.addSubtransforms(processSizedElementsAndRestrictionsId); |
| rval.setPtransform(newCompositeRoot); |
| return rval.build(); |
| } catch (IOException e) { |
| throw new RuntimeException("Unable to perform expansion for transform " + transformId, e); |
| } |
| } |
| } |
| |
| private static String getOrAddDoubleCoder( |
| ComponentsOrBuilder existingComponents, MessageWithComponents.Builder out) { |
| for (Map.Entry<String, Coder> coder : existingComponents.getCodersMap().entrySet()) { |
| if (ModelCoders.DOUBLE_CODER_URN.equals(coder.getValue().getSpec().getUrn())) { |
| return coder.getKey(); |
| } |
| } |
| String doubleCoderId = generateUniqueId("DoubleCoder", existingComponents::containsCoders); |
| out.getComponentsBuilder() |
| .putCoders( |
| doubleCoderId, |
| Coder.newBuilder() |
| .setSpec(FunctionSpec.newBuilder().setUrn(ModelCoders.DOUBLE_CODER_URN)) |
| .build()); |
| return doubleCoderId; |
| } |
| |
| /** |
| * Returns a PCollection name that uses the supplied prefix that does not exist in {@code |
| * existingComponents}. |
| */ |
| private static String generateUniquePCollectonName( |
| String prefix, ComponentsOrBuilder existingComponents) { |
| return generateUniqueId( |
| prefix, |
| input -> { |
| for (PCollection pc : existingComponents.getPcollectionsMap().values()) { |
| if (input.equals(pc.getUniqueName())) { |
| return true; |
| } |
| } |
| return false; |
| }); |
| } |
| |
| /** Generates a unique id given a prefix and a predicate to compare if the id is already used. */ |
| private static String generateUniqueId(String prefix, Predicate<String> isExistingId) { |
| int i = 0; |
| while (isExistingId.test(prefix + i)) { |
| i += 1; |
| } |
| return prefix + i; |
| } |
| } |
