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.sdk.transforms;
  
  import com.google.auto.service.AutoService;
  import java.util.Map;
  import java.util.concurrent.ThreadLocalRandom;
  import org.apache.beam.model.pipeline.v1.RunnerApi;
  import org.apache.beam.sdk.annotations.Internal;
  import org.apache.beam.sdk.runners.AppliedPTransform;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.transforms.windowing.ReshuffleTrigger;
  import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
  import org.apache.beam.sdk.transforms.windowing.Window;
  import org.apache.beam.sdk.util.IdentityWindowFn;
  import org.apache.beam.sdk.util.construction.PTransformTranslation;
  import org.apache.beam.sdk.util.construction.SdkComponents;
  import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.ValueInSingleWindow;
  import org.apache.beam.sdk.values.WindowingStrategy;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedInteger;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.joda.time.Duration;
  
  /**
   * A family of {@link PTransform PTransforms} that returns a {@link PCollection} equivalent to its
   * input but functions as an operational hint to a runner that redistributing the data in some way
   * is likely useful.
   */
  public class Redistribute {
    /** @return a {@link RedistributeArbitrarily} transform with default configuration. */
    public static <T> RedistributeArbitrarily<T> arbitrarily() {
      return new RedistributeArbitrarily<>(null, false);
    }
  
    /** @return a {@link RedistributeByKey} transform with default configuration. */
    public static <K, V> RedistributeByKey<K, V> byKey() {
      return new RedistributeByKey<>(false);
    }
  
    /**
     * @param <K> The type of key being reshuffled on.
     * @param <V> The type of value being reshuffled.
     */
    public static class RedistributeByKey<K, V>
        extends PTransform<PCollection<KV<K, V>>, PCollection<KV<K, V>>> {
  
      private final boolean allowDuplicates;
  
      private RedistributeByKey(boolean allowDuplicates) {
        this.allowDuplicates = allowDuplicates;
      }
  
      public RedistributeByKey<K, V> withAllowDuplicates(boolean newAllowDuplicates) {
        return new RedistributeByKey<>(newAllowDuplicates);
      }
  
      public boolean getAllowDuplicates() {
        return allowDuplicates;
      }
  
      @Override
      public PCollection<KV<K, V>> expand(PCollection<KV<K, V>> input) {
        WindowingStrategy<?, ?> originalStrategy = input.getWindowingStrategy();
        // If the input has already had its windows merged, then the GBK that performed the merge
        // will have set originalStrategy.getWindowFn() to InvalidWindows, causing the GBK contained
        // here to fail. Instead, we install a valid WindowFn that leaves all windows unchanged.
        // The TimestampCombiner is set to ensure the GroupByKey does not shift elements forwards in
        // time.
        // Because this outputs as fast as possible, this should not hold the watermark.
        Window<KV<K, V>> rewindow =
            Window.<KV<K, V>>into(
                    new IdentityWindowFn<>(originalStrategy.getWindowFn().windowCoder()))
                .triggering(new ReshuffleTrigger<>())
                .discardingFiredPanes()
                .withTimestampCombiner(TimestampCombiner.EARLIEST)
                .withAllowedLateness(Duration.millis(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
  
        PCollection<KV<K, ValueInSingleWindow<V>>> reified =
            input
                .apply("SetIdentityWindow", rewindow)
                .apply("ReifyOriginalMetadata", Reify.windowsInValue());
  
        PCollection<KV<K, Iterable<ValueInSingleWindow<V>>>> grouped =
            reified.apply(GroupByKey.create());
        return grouped
            .apply(
                "ExpandIterable",
                ParDo.of(
                    new DoFn<
                        KV<K, Iterable<ValueInSingleWindow<V>>>, KV<K, ValueInSingleWindow<V>>>() {
                      @ProcessElement
                      public void processElement(
                          @Element KV<K, Iterable<ValueInSingleWindow<V>>> element,
                          OutputReceiver<KV<K, ValueInSingleWindow<V>>> r) {
                        K key = element.getKey();
                        for (ValueInSingleWindow<V> value : element.getValue()) {
                          r.output(KV.of(key, value));
                        }
                      }
                    }))
            .apply("RestoreMetadata", new RestoreMetadata<>())
            // Set the windowing strategy directly, so that it doesn't get counted as the user having
            // set allowed lateness.
            .setWindowingStrategyInternal(originalStrategy);
      }
    }
  
    /**
     * Noop transform that hints to the runner to try to redistribute the work evenly, or via whatever
     * clever strategy the runner comes up with.
     */
    public static class RedistributeArbitrarily<T>
        extends PTransform<PCollection<T>, PCollection<T>> {
      // The number of buckets to shard into.
      // A runner is free to ignore this (a runner may ignore the transform
      // entirely!) This is a performance optimization to prevent having
      // unit sized bundles on the output. If unset, uses a random integer key.
      private @Nullable Integer numBuckets = null;
      private boolean allowDuplicates = false;
  
      private RedistributeArbitrarily(@Nullable Integer numBuckets, boolean allowDuplicates) {
        this.numBuckets = numBuckets;
        this.allowDuplicates = allowDuplicates;
      }
  
      public RedistributeArbitrarily<T> withNumBuckets(@Nullable Integer numBuckets) {
        return new RedistributeArbitrarily<>(numBuckets, this.allowDuplicates);
      }
  
      public RedistributeArbitrarily<T> withAllowDuplicates(boolean allowDuplicates) {
        return new RedistributeArbitrarily<>(this.numBuckets, allowDuplicates);
      }
  
      public boolean getAllowDuplicates() {
        return allowDuplicates;
      }
  
      @Override
      public PCollection<T> expand(PCollection<T> input) {
        return input
            .apply("Pair with random key", ParDo.of(new AssignShardFn<>(numBuckets)))
            .apply(Redistribute.<Integer, T>byKey().withAllowDuplicates(this.allowDuplicates))
            .apply(Values.create());
      }
    }
  
    private static class RestoreMetadata<K, V>
        extends PTransform<PCollection<KV<K, ValueInSingleWindow<V>>>, PCollection<KV<K, V>>> {
      @Override
      public PCollection<KV<K, V>> expand(PCollection<KV<K, ValueInSingleWindow<V>>> input) {
        return input.apply(
            ParDo.of(
                new DoFn<KV<K, ValueInSingleWindow<V>>, KV<K, V>>() {
                  @Override
                  public Duration getAllowedTimestampSkew() {
                    return Duration.millis(Long.MAX_VALUE);
                  }
  
                  @ProcessElement
                  public void processElement(
                      @Element KV<K, ValueInSingleWindow<V>> kv,
                      OutputReceiver<KV<K, V>> outputReceiver) {
                    outputReceiver
                        .builder(KV.of(kv.getKey(), kv.getValue().getValue()))
                        .setTimestamp(kv.getValue().getTimestamp())
                        .setWindow(kv.getValue().getWindow())
                        .setPaneInfo(kv.getValue().getPaneInfo())
                        .output();
                  }
                }));
      }
    }
  
    static class AssignShardFn<T> extends DoFn<T, KV<Integer, T>> {
      private int shard;
      private @Nullable Integer numBuckets;
  
      public AssignShardFn(@Nullable Integer numBuckets) {
        this.numBuckets = numBuckets;
      }
  
      @Setup
      public void setup() {
        shard = ThreadLocalRandom.current().nextInt();
      }
  
      @ProcessElement
      public void processElement(@Element T element, OutputReceiver<KV<Integer, T>> r) {
        ++shard;
        // Smear the shard into something more random-looking, to avoid issues
        // with runners that don't properly hash the key being shuffled, but rely
        // on it being random-looking. E.g. Spark takes the Java hashCode() of keys,
        // which for Integer is a no-op and it is an issue:
        // http://hydronitrogen.com/poor-hash-partitioning-of-timestamps-integers-and-longs-in-
        // spark.html
        // This hashing strategy is copied from
        // org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Hashing.smear().
        int hashOfShard = 0x1b873593 * Integer.rotateLeft(shard * 0xcc9e2d51, 15);
        if (numBuckets != null) {
          UnsignedInteger unsignedNumBuckets = UnsignedInteger.fromIntBits(numBuckets);
          hashOfShard = UnsignedInteger.fromIntBits(hashOfShard).mod(unsignedNumBuckets).intValue();
        }
        r.output(KV.of(hashOfShard, element));
      }
    }
  
    static class RedistributeByKeyTranslator
        implements PTransformTranslation.TransformPayloadTranslator<RedistributeByKey<?, ?>> {
      @Override
      public String getUrn() {
        return PTransformTranslation.REDISTRIBUTE_BY_KEY_URN;
      }
  
      @Override
      @SuppressWarnings("nullness") // Cannot figure out how to make this typecheck
      public RunnerApi.FunctionSpec translate(
          AppliedPTransform<?, ?, RedistributeByKey<?, ?>> transform, SdkComponents components) {
        return RunnerApi.FunctionSpec.newBuilder()
            .setUrn(getUrn(transform.getTransform()))
            .setPayload(
                RunnerApi.RedistributePayload.newBuilder()
                    .setAllowDuplicates(transform.getTransform().getAllowDuplicates())
                    .build()
                    .toByteString())
            .build();
      }
    }
  
    static class RedistributeArbitrarilyTranslator
        implements PTransformTranslation.TransformPayloadTranslator<RedistributeArbitrarily<?>> {
      @Override
      public String getUrn() {
        return PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN;
      }
  
      @Override
      @SuppressWarnings("nullness") // Cannot figure out how to make this typecheck
      public RunnerApi.FunctionSpec translate(
          AppliedPTransform<?, ?, RedistributeArbitrarily<?>> transform, SdkComponents components) {
        return RunnerApi.FunctionSpec.newBuilder()
            .setUrn(getUrn(transform.getTransform()))
            .setPayload(
                RunnerApi.RedistributePayload.newBuilder()
                    .setAllowDuplicates(transform.getTransform().getAllowDuplicates())
                    .build()
                    .toByteString())
            .build();
      }
    }
  
    /** Registers translators for the Redistribute family of transforms. */
    @Internal
    @AutoService(TransformPayloadTranslatorRegistrar.class)
    public static class Registrar implements TransformPayloadTranslatorRegistrar {
  
      @Override
      @SuppressWarnings("rawtypes")
      public Map<
              ? extends Class<? extends PTransform>,
              ? extends PTransformTranslation.TransformPayloadTranslator>
          getTransformPayloadTranslators() {
        return ImmutableMap
            .<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder()
            .put(RedistributeByKey.class, new RedistributeByKeyTranslator())
            .put(RedistributeArbitrarily.class, new RedistributeArbitrarilyTranslator())
            .build();
      }
    }
  }
  