content: /*
   * 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.flink;
  
  import static org.apache.beam.sdk.util.construction.PTransformTranslation.WRITE_FILES_TRANSFORM_URN;
  
  import java.io.IOException;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.concurrent.ThreadLocalRandom;
  import java.util.concurrent.TimeUnit;
  import org.apache.beam.runners.flink.adapter.FlinkKey;
  import org.apache.beam.sdk.Pipeline;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.coders.ShardedKeyCoder;
  import org.apache.beam.sdk.coders.VarIntCoder;
  import org.apache.beam.sdk.io.FileBasedSink;
  import org.apache.beam.sdk.io.ShardingFunction;
  import org.apache.beam.sdk.io.WriteFiles;
  import org.apache.beam.sdk.io.WriteFilesResult;
  import org.apache.beam.sdk.options.PipelineOptions;
  import org.apache.beam.sdk.runners.AppliedPTransform;
  import org.apache.beam.sdk.runners.PTransformMatcher;
  import org.apache.beam.sdk.runners.PTransformOverrideFactory;
  import org.apache.beam.sdk.runners.TransformHierarchy;
  import org.apache.beam.sdk.transforms.PTransform;
  import org.apache.beam.sdk.util.CoderUtils;
  import org.apache.beam.sdk.util.construction.PTransformReplacements;
  import org.apache.beam.sdk.util.construction.PTransformTranslation;
  import org.apache.beam.sdk.util.construction.ReplacementOutputs;
  import org.apache.beam.sdk.util.construction.UnconsumedReads;
  import org.apache.beam.sdk.util.construction.WriteFilesTranslation;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.PCollectionView;
  import org.apache.beam.sdk.values.PValue;
  import org.apache.beam.sdk.values.ShardedKey;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
  import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
  import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
  import org.apache.flink.util.Preconditions;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * This is a {@link FlinkPipelineTranslator} for streaming jobs. Its role is to translate the
   * user-provided {@link org.apache.beam.sdk.values.PCollection}-based job into a {@link
   * org.apache.flink.streaming.api.datastream.DataStream} one.
   */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  class FlinkStreamingPipelineTranslator extends FlinkPipelineTranslator {
  
    private static final Logger LOG = LoggerFactory.getLogger(FlinkStreamingPipelineTranslator.class);
  
    /** The necessary context in the case of a straming job. */
    private final FlinkStreamingTranslationContext streamingContext;
  
    private int depth = 0;
  
    public FlinkStreamingPipelineTranslator(
        StreamExecutionEnvironment env, PipelineOptions options, boolean isStreaming) {
      this.streamingContext = new FlinkStreamingTranslationContext(env, options, isStreaming);
    }
  
    @Override
    public void translate(Pipeline pipeline) {
      // Ensure all outputs of all reads are consumed.
      UnconsumedReads.ensureAllReadsConsumed(pipeline);
      super.translate(pipeline);
    }
  
    // --------------------------------------------------------------------------------------------
    //  Pipeline Visitor Methods
    // --------------------------------------------------------------------------------------------
  
    @Override
    public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
      LOG.info("{} enterCompositeTransform- {}", genSpaces(this.depth), node.getFullName());
      this.depth++;
  
      PTransform<?, ?> transform = node.getTransform();
      if (transform != null) {
        StreamTransformTranslator<?> translator =
            FlinkStreamingTransformTranslators.getTranslator(transform);
  
        if (translator != null && applyCanTranslate(transform, node, translator)) {
          applyStreamingTransform(transform, node, translator);
          LOG.info("{} translated- {}", genSpaces(this.depth), node.getFullName());
          return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
        }
      }
      return CompositeBehavior.ENTER_TRANSFORM;
    }
  
    @Override
    public void leaveCompositeTransform(TransformHierarchy.Node node) {
      this.depth--;
      LOG.info("{} leaveCompositeTransform- {}", genSpaces(this.depth), node.getFullName());
    }
  
    @Override
    public void visitPrimitiveTransform(TransformHierarchy.Node node) {
      LOG.info("{} visitPrimitiveTransform- {}", genSpaces(this.depth), node.getFullName());
      // get the transformation corresponding to hte node we are
      // currently visiting and translate it into its Flink alternative.
  
      PTransform<?, ?> transform = node.getTransform();
      StreamTransformTranslator<?> translator =
          FlinkStreamingTransformTranslators.getTranslator(transform);
  
      if (translator == null || !applyCanTranslate(transform, node, translator)) {
        String transformUrn = PTransformTranslation.urnForTransform(transform);
        LOG.info(transformUrn);
        throw new UnsupportedOperationException(
            "The transform " + transformUrn + " is currently not supported.");
      }
      applyStreamingTransform(transform, node, translator);
    }
  
    @Override
    public void visitValue(PValue value, TransformHierarchy.Node producer) {
      // do nothing here
    }
  
    private <T extends PTransform<?, ?>> void applyStreamingTransform(
        PTransform<?, ?> transform,
        TransformHierarchy.Node node,
        StreamTransformTranslator<?> translator) {
  
      @SuppressWarnings("unchecked")
      T typedTransform = (T) transform;
  
      @SuppressWarnings("unchecked")
      StreamTransformTranslator<T> typedTranslator = (StreamTransformTranslator<T>) translator;
  
      // create the applied PTransform on the streamingContext
      streamingContext.setCurrentTransform(node.toAppliedPTransform(getPipeline()));
      typedTranslator.translateNode(typedTransform, streamingContext);
    }
  
    private <T extends PTransform<?, ?>> boolean applyCanTranslate(
        PTransform<?, ?> transform,
        TransformHierarchy.Node node,
        StreamTransformTranslator<?> translator) {
  
      @SuppressWarnings("unchecked")
      T typedTransform = (T) transform;
  
      @SuppressWarnings("unchecked")
      StreamTransformTranslator<T> typedTranslator = (StreamTransformTranslator<T>) translator;
  
      streamingContext.setCurrentTransform(node.toAppliedPTransform(getPipeline()));
  
      return typedTranslator.canTranslate(typedTransform, streamingContext);
    }
  
    /**
     * The interface that every Flink translator of a Beam operator should implement. This interface
     * is for <b>streaming</b> jobs. For examples of such translators see {@link
     * FlinkStreamingTransformTranslators}.
     */
    abstract static class StreamTransformTranslator<T extends PTransform> {
  
      /** Translate the given transform. */
      abstract void translateNode(T transform, FlinkStreamingTranslationContext context);
  
      /** Returns true iff this translator can translate the given transform. */
      boolean canTranslate(T transform, FlinkStreamingTranslationContext context) {
        return true;
      }
    }
  
    @VisibleForTesting
    static class StreamingShardedWriteFactory<UserT, DestinationT, OutputT>
        implements PTransformOverrideFactory<
            PCollection<UserT>,
            WriteFilesResult<DestinationT>,
            WriteFiles<UserT, DestinationT, OutputT>> {
  
      /**
       * {@link PTransformMatcher} which decides if {@link StreamingShardedWriteFactory} should be
       * applied.
       */
      static PTransformMatcher writeFilesNeedsOverrides() {
        return application -> {
          if (WRITE_FILES_TRANSFORM_URN.equals(
              PTransformTranslation.urnForTransformOrNull(application.getTransform()))) {
            try {
              FlinkPipelineOptions options =
                  application.getPipeline().getOptions().as(FlinkPipelineOptions.class);
              ShardingFunction shardingFn =
                  ((WriteFiles<?, ?, ?>) application.getTransform()).getShardingFunction();
              return WriteFilesTranslation.isRunnerDeterminedSharding((AppliedPTransform) application)
                  || (options.isAutoBalanceWriteFilesShardingEnabled() && shardingFn == null);
            } catch (IOException exc) {
              throw new RuntimeException(
                  String.format(
                      "Transform with URN %s failed to parse: %s",
                      WRITE_FILES_TRANSFORM_URN, application.getTransform()),
                  exc);
            }
          }
          return false;
        };
      }
  
      FlinkPipelineOptions options;
  
      StreamingShardedWriteFactory(PipelineOptions options) {
        this.options = options.as(FlinkPipelineOptions.class);
      }
  
      @Override
      public PTransformReplacement<PCollection<UserT>, WriteFilesResult<DestinationT>>
          getReplacementTransform(
              AppliedPTransform<
                      PCollection<UserT>,
                      WriteFilesResult<DestinationT>,
                      WriteFiles<UserT, DestinationT, OutputT>>
                  transform) {
        // By default, if numShards is not set WriteFiles will produce one file per bundle. In
        // streaming, there are large numbers of small bundles, resulting in many tiny files.
        // Instead we pick parallelism * 2 to ensure full parallelism, but prevent too-many files.
        Integer jobParallelism = options.getParallelism();
  
        Preconditions.checkArgument(
            jobParallelism > 0,
            "Parallelism of a job should be greater than 0. Currently set: %s",
            jobParallelism);
        int numShards = jobParallelism * 2;
  
        try {
          List<PCollectionView<?>> sideInputs =
              WriteFilesTranslation.getDynamicDestinationSideInputs(transform);
          FileBasedSink sink = WriteFilesTranslation.getSink(transform);
  
          @SuppressWarnings("unchecked")
          WriteFiles<UserT, DestinationT, OutputT> replacement =
              WriteFiles.to(sink).withSideInputs(sideInputs);
          if (WriteFilesTranslation.isWindowedWrites(transform)) {
            replacement = replacement.withWindowedWrites();
          }
  
          if (WriteFilesTranslation.isRunnerDeterminedSharding(transform)) {
            replacement = replacement.withNumShards(numShards);
          } else {
            if (transform.getTransform().getNumShardsProvider() != null) {
              replacement =
                  replacement.withNumShards(transform.getTransform().getNumShardsProvider());
            }
            if (transform.getTransform().getComputeNumShards() != null) {
              replacement = replacement.withSharding(transform.getTransform().getComputeNumShards());
            }
          }
  
          if (options.isAutoBalanceWriteFilesShardingEnabled()) {
  
            replacement =
                replacement.withShardingFunction(
                    new FlinkAutoBalancedShardKeyShardingFunction<>(
                        jobParallelism,
                        options.getMaxParallelism(),
                        sink.getDynamicDestinations().getDestinationCoder()));
          }
  
          return PTransformReplacement.of(
              PTransformReplacements.getSingletonMainInput(transform), replacement);
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      }
  
      @Override
      public Map<PCollection<?>, ReplacementOutput> mapOutputs(
          Map<TupleTag<?>, PCollection<?>> outputs, WriteFilesResult<DestinationT> newOutput) {
        return ReplacementOutputs.tagged(outputs, newOutput);
      }
    }
  
    /**
     * Flink assigns elements to parallel operators (workers) based their key group. Key group is
     * determined based on a Murmur hash of element's key. This will have skew distribution in case
     * when number of keys is not >> key groups, which is typical scenario when writing files, where
     * one do not want to end up with too many small files. This {@link ShardingFunction} implements
     * what was suggested on Flink's <a
     * href="http://mail-archives.apache.org/mod_mbox/flink-dev/201901.mbox/%3CCAOUjMkygmFMDbOJNCmndrZ0bYug=iQJmVz6QvMW7C87n=pw8SQ@mail.gmail.com%3E">mailing
     * list</a> and properly chooses shard keys in a way that they are distributed evenly among
     * workers by Flink.
     */
    @VisibleForTesting
    static class FlinkAutoBalancedShardKeyShardingFunction<UserT, DestinationT>
        implements ShardingFunction<UserT, DestinationT> {
  
      @VisibleForTesting static final int CACHE_MAX_SIZE = 100;
      private static final long CACHE_EXPIRE_SECONDS = 600;
  
      private final int parallelism;
      private final int maxParallelism;
      private final Coder<DestinationT> destinationCoder;
      private final ShardedKeyCoder<Integer> shardedKeyCoder = ShardedKeyCoder.of(VarIntCoder.of());
      private transient Cache<Integer, Map<Integer, ShardedKey<Integer>>> cache;
  
      private int shardNumber = -1;
  
      @VisibleForTesting
      Map<Integer, Map<Integer, ShardedKey<Integer>>> getCache() {
        return cache == null ? null : cache.asMap();
      }
  
      @VisibleForTesting
      int getMaxParallelism() {
        return maxParallelism;
      }
  
      FlinkAutoBalancedShardKeyShardingFunction(
          int parallelism, int maxParallelism, Coder<DestinationT> destinationCoder) {
        this.parallelism = parallelism;
        // keep resolution of maxParallelism to sharding functions, as it relies on Flink's
        //  state API at KeyGroupRangeAssignment
        this.maxParallelism =
            maxParallelism > 0
                ? maxParallelism
                : KeyGroupRangeAssignment.computeDefaultMaxParallelism(parallelism);
  
        this.destinationCoder = destinationCoder;
      }
  
      @Override
      public ShardedKey<Integer> assignShardKey(
          DestinationT destination, UserT element, int shardCount) throws Exception {
  
        // Same as in WriteFiles ...
        // We want to desynchronize the first record sharding key for each instance of
        // ApplyShardingKey, so records in a small PCollection will be statistically balanced.
        if (shardNumber == -1) {
          shardNumber = ThreadLocalRandom.current().nextInt(shardCount);
        } else {
          shardNumber = (shardNumber + 1) % shardCount;
        }
  
        int destinationKey =
            Arrays.hashCode(CoderUtils.encodeToByteArray(destinationCoder, destination));
  
        if (cache == null) {
          cache =
              CacheBuilder.newBuilder()
                  .maximumSize(CACHE_MAX_SIZE)
                  .expireAfterAccess(CACHE_EXPIRE_SECONDS, TimeUnit.SECONDS)
                  .build();
        }
  
        // we need to ensure that keys are always stable no matter at which worker they
        // are created and what is an order of observed shard numbers
        if (cache.getIfPresent(destinationKey) == null) {
          cache.put(destinationKey, generateShardedKeys(destinationKey, shardCount));
        }
  
        return cache.getIfPresent(destinationKey).get(shardNumber);
      }
  
      private Map<Integer, ShardedKey<Integer>> generateShardedKeys(int key, int shardCount) {
  
        Map<Integer, ShardedKey<Integer>> shardedKeys = new HashMap<>();
  
        for (int shard = 0; shard < shardCount; shard++) {
  
          int salt = -1;
          while (true) {
            if (salt++ == Integer.MAX_VALUE) {
              throw new RuntimeException(
                  "Failed to find sharded key in [ " + Integer.MAX_VALUE + " ] iterations");
            }
            ShardedKey<Integer> shk = ShardedKey.of(Objects.hash(key, salt), shard);
            int targetPartition = shard % parallelism;
  
            // create effective key in the same way Beam/Flink will do so we can see if it gets
            // allocated to the partition we want
            FlinkKey effectiveKey = FlinkKey.of(shk, shardedKeyCoder);
  
            int partition =
                KeyGroupRangeAssignment.assignKeyToParallelOperator(
                    effectiveKey, maxParallelism, parallelism);
  
            if (partition == targetPartition) {
              shardedKeys.put(shard, shk);
              break;
            }
          }
        }
  
        return shardedKeys;
      }
    }
  }
  