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.spark.translation;
  
  import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
  import java.io.Serializable;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.Map;
  import javax.annotation.Nonnull;
  import org.apache.beam.runners.core.InMemoryStateInternals;
  import org.apache.beam.runners.core.StateInternals;
  import org.apache.beam.runners.core.StateInternalsFactory;
  import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
  import org.apache.beam.runners.spark.SparkPipelineOptions;
  import org.apache.beam.runners.spark.SparkRunner;
  import org.apache.beam.runners.spark.coders.CoderHelpers;
  import org.apache.beam.runners.spark.util.ByteArray;
  import org.apache.beam.runners.spark.util.SideInputBroadcast;
  import org.apache.beam.sdk.coders.Coder;
  import org.apache.beam.sdk.state.TimeDomain;
  import org.apache.beam.sdk.state.TimerSpec;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
  import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
  import org.apache.beam.sdk.transforms.windowing.Window;
  import org.apache.beam.sdk.transforms.windowing.WindowFn;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.PCollectionView;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.sdk.values.WindowedValue;
  import org.apache.beam.sdk.values.WindowedValues;
  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.collect.Iterators;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
  import org.apache.spark.api.java.JavaSparkContext;
  import org.apache.spark.api.java.function.FlatMapFunction;
  import org.apache.spark.api.java.function.Function;
  import org.apache.spark.api.java.function.PairFlatMapFunction;
  import org.apache.spark.api.java.function.PairFunction;
  import org.apache.spark.api.java.function.VoidFunction;
  import org.apache.spark.storage.StorageLevel;
  import org.apache.spark.streaming.Duration;
  import org.apache.spark.streaming.api.java.JavaDStream;
  import org.apache.spark.streaming.api.java.JavaPairDStream;
  import org.apache.spark.streaming.dstream.DStream;
  import scala.Tuple2;
  
  /** A set of utilities to help translating Beam transformations into Spark transformations. */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public final class TranslationUtils {
  
    private TranslationUtils() {}
  
    /**
     * In-memory state internals factory.
     *
     * @param <K> State key type.
     */
    static class InMemoryStateInternalsFactory<K> implements StateInternalsFactory<K>, Serializable {
      @Override
      public StateInternals stateInternalsForKey(K key) {
        return InMemoryStateInternals.forKey(key);
      }
    }
  
    /**
     * A SparkCombineFn function applied to grouped KVs.
     *
     * @param <K> Grouped key type.
     * @param <InputT> Grouped values type.
     * @param <OutputT> Output type.
     */
    public static class CombineGroupedValues<K, InputT, OutputT>
        implements Function<WindowedValue<KV<K, Iterable<InputT>>>, WindowedValue<KV<K, OutputT>>> {
      private final SparkCombineFn<KV<K, InputT>, InputT, ?, OutputT> fn;
  
      public CombineGroupedValues(SparkCombineFn<KV<K, InputT>, InputT, ?, OutputT> fn) {
        this.fn = fn;
      }
  
      @SuppressWarnings("unchecked")
      @Override
      public WindowedValue<KV<K, OutputT>> call(WindowedValue<KV<K, Iterable<InputT>>> windowedKv)
          throws Exception {
        return WindowedValues.of(
            KV.of(
                windowedKv.getValue().getKey(),
                fn.getCombineFn()
                    .apply(windowedKv.getValue().getValue(), fn.ctxtForValue(windowedKv))),
            windowedKv.getTimestamp(),
            windowedKv.getWindows(),
            windowedKv.getPaneInfo());
      }
    }
  
    /**
     * Checks if the window transformation should be applied or skipped.
     *
     * <p>Avoid running assign windows if both source and destination are global window or if the user
     * has not specified the WindowFn (meaning they are just messing with triggering or allowed
     * lateness).
     *
     * @param transform The {@link Window.Assign} transformation.
     * @param context The {@link EvaluationContext}.
     * @param <T> PCollection type.
     * @param <W> {@link BoundedWindow} type.
     * @return if to apply the transformation.
     */
    public static <T, W extends BoundedWindow> boolean skipAssignWindows(
        Window.Assign<T> transform, EvaluationContext context) {
      @SuppressWarnings("unchecked")
      WindowFn<? super T, W> windowFn = (WindowFn<? super T, W>) transform.getWindowFn();
      return windowFn == null
          || (context.getInput(transform).getWindowingStrategy().getWindowFn()
                  instanceof GlobalWindows
              && windowFn instanceof GlobalWindows);
    }
  
    /** Transform a pair stream into a value stream. */
    public static <T1, T2> JavaDStream<T2> dStreamValues(JavaPairDStream<T1, T2> pairDStream) {
      return pairDStream.map(Tuple2::_2);
    }
  
    /** {@link KV} to pair function. */
    public static <K, V> PairFunction<KV<K, V>, K, V> toPairFunction() {
      return kv -> new Tuple2<>(kv.getKey(), kv.getValue());
    }
  
    /** {@link KV} to pair flatmap function. */
    public static <K, V> PairFlatMapFunction<Iterator<KV<K, V>>, K, V> toPairFlatMapFunction() {
      return itr -> Iterators.transform(itr, kv -> new Tuple2<>(kv.getKey(), kv.getValue()));
    }
  
    /** A pair to {@link KV} function . */
    static class FromPairFunction<K, V>
        implements Function<Tuple2<K, V>, KV<K, V>>,
            org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Function<
                Tuple2<K, V>, KV<K, V>> {
      @Override
      public KV<K, V> call(Tuple2<K, V> t2) {
        return KV.of(t2._1(), t2._2());
      }
  
      @SuppressFBWarnings(
          value = "NP_METHOD_PARAMETER_TIGHTENS_ANNOTATION",
          justification = "https://github.com/google/guava/issues/920")
      @Override
      public KV<K, V> apply(@Nonnull Tuple2<K, V> t2) {
        return call(t2);
      }
    }
  
    /** A pair to {@link KV} flatmap function . */
    static <K, V> FlatMapFunction<Iterator<Tuple2<K, V>>, KV<K, V>> fromPairFlatMapFunction() {
      return itr -> Iterators.transform(itr, t2 -> KV.of(t2._1(), t2._2()));
    }
  
    /** Extract key from a {@link WindowedValue} {@link KV} into a pair. */
    public static <K, V>
        PairFunction<WindowedValue<KV<K, V>>, ByteArray, WindowedValue<KV<K, V>>>
            toPairByKeyInWindowedValue(final Coder<K> keyCoder) {
      return windowedKv ->
          new Tuple2<>(
              new ByteArray(CoderHelpers.toByteArray(windowedKv.getValue().getKey(), keyCoder)),
              windowedKv);
    }
  
    /** Extract window from a {@link KV} with {@link WindowedValue} value. */
    static class ToKVByWindowInValueFunction<K, V>
        implements Function<KV<K, WindowedValue<V>>, WindowedValue<KV<K, V>>>,
            org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Function<
                KV<K, WindowedValue<V>>, WindowedValue<KV<K, V>>> {
  
      @Override
      public WindowedValue<KV<K, V>> call(KV<K, WindowedValue<V>> kv) {
        WindowedValue<V> wv = kv.getValue();
        return wv.withValue(KV.of(kv.getKey(), wv.getValue()));
      }
  
      @SuppressFBWarnings(
          value = "NP_METHOD_PARAMETER_TIGHTENS_ANNOTATION",
          justification = "https://github.com/google/guava/issues/920")
      @Override
      public WindowedValue<KV<K, V>> apply(@Nonnull KV<K, WindowedValue<V>> kv) {
        return call(kv);
      }
    }
  
    /**
     * A utility class to filter {@link TupleTag}s.
     *
     * @param <V> TupleTag type.
     */
    public static final class TupleTagFilter<V>
        implements Function<Tuple2<TupleTag<V>, WindowedValue<?>>, Boolean> {
  
      private final TupleTag<V> tag;
  
      public TupleTagFilter(TupleTag<V> tag) {
        this.tag = tag;
      }
  
      @Override
      public Boolean call(Tuple2<TupleTag<V>, WindowedValue<?>> input) {
        return tag.equals(input._1());
      }
    }
  
    /**
     * Create SideInputs as Broadcast variables.
     *
     * @param views The {@link PCollectionView}s.
     * @param context The {@link EvaluationContext}.
     * @return a map of tagged {@link SideInputBroadcast}s and their {@link WindowingStrategy}.
     */
    static Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> getSideInputs(
        Iterable<PCollectionView<?>> views, EvaluationContext context) {
      return getSideInputs(views, context.getSparkContext(), context.getPViews());
    }
  
    /**
     * Create SideInputs as Broadcast variables.
     *
     * @param views The {@link PCollectionView}s.
     * @param context The {@link JavaSparkContext}.
     * @param pviews The {@link SparkPCollectionView}.
     * @return a map of tagged {@link SideInputBroadcast}s and their {@link WindowingStrategy}.
     */
    public static Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> getSideInputs(
        Iterable<PCollectionView<?>> views, JavaSparkContext context, SparkPCollectionView pviews) {
      if (views == null) {
        return ImmutableMap.of();
      } else {
        Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs =
            Maps.newHashMap();
        for (PCollectionView<?> view : views) {
          SideInputBroadcast helper = pviews.getPCollectionView(view, context);
          WindowingStrategy<?, ?> windowingStrategy = view.getWindowingStrategyInternal();
          sideInputs.put(view.getTagInternal(), KV.of(windowingStrategy, helper));
        }
        return sideInputs;
      }
    }
  
    /**
     * Retrieves the batch duration in milliseconds from Spark pipeline options.
     *
     * @param options The serializable pipeline options containing Spark-specific settings
     * @return The checkpoint duration in milliseconds as specified in SparkPipelineOptions
     */
    public static Long getBatchDuration(final SerializablePipelineOptions options) {
      return options.get().as(SparkPipelineOptions.class).getCheckpointDurationMillis();
    }
  
    /**
     * Checks if the given DoFn uses any timers.
     *
     * @param doFn the DoFn to check for timer usage
     * @return true if the DoFn uses timers, false otherwise
     */
    public static boolean hasTimers(DoFn<?, ?> doFn) {
      final DoFnSignature signature = DoFnSignatures.signatureForDoFn(doFn);
      return signature.usesTimers();
    }
  
    /**
     * Checks if the given DoFn uses event time timers.
     *
     * @param doFn the DoFn to check for event time timer usage
     * @return true if the DoFn uses event time timers, false otherwise. Note: Returns false if the
     *     DoFn has no timers at all.
     */
    public static boolean hasEventTimers(DoFn<?, ?> doFn) {
      for (DoFnSignature.TimerDeclaration timerDeclaration :
          DoFnSignatures.signatureForDoFn(doFn).timerDeclarations().values()) {
        final TimerSpec timerSpec = DoFnSignatures.getTimerSpecOrThrow(timerDeclaration, doFn);
        return timerSpec.getTimeDomain().equals(TimeDomain.EVENT_TIME);
      }
      return false;
    }
  
    /**
     * Checkpoints the given DStream if checkpointing is enabled in the pipeline options.
     *
     * @param dStream The DStream to be checkpointed
     * @param options The SerializablePipelineOptions containing configuration settings including
     *     batch duration
     */
    public static void checkpointIfNeeded(
        final DStream<?> dStream, final SerializablePipelineOptions options) {
  
      final Long checkpointDurationMillis = getBatchDuration(options);
  
      if (checkpointDurationMillis > 0) {
        dStream.checkpoint(new Duration(checkpointDurationMillis));
      }
    }
  
    /**
     * Reject state and timers {@link DoFn}.
     *
     * @param doFn the {@link DoFn} to possibly reject.
     */
    public static void rejectStateAndTimers(DoFn<?, ?> doFn) {
      DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());
  
      if (signature.stateDeclarations().size() > 0) {
        throw new UnsupportedOperationException(
            String.format(
                "Found %s annotations on %s, but %s cannot yet be used with state in the %s.",
                DoFn.StateId.class.getSimpleName(),
                doFn.getClass().getName(),
                DoFn.class.getSimpleName(),
                SparkRunner.class.getSimpleName()));
      }
  
      if (signature.timerDeclarations().size() > 0
          || signature.timerFamilyDeclarations().size() > 0) {
        throw new UnsupportedOperationException(
            String.format(
                "Found %s annotations on %s, but %s cannot yet be used with timers in the %s.",
                DoFn.TimerId.class.getSimpleName(),
                doFn.getClass().getName(),
                DoFn.class.getSimpleName(),
                SparkRunner.class.getSimpleName()));
      }
    }
  
    public static <T> VoidFunction<T> emptyVoidFunction() {
      return t -> {
        // Empty implementation.
      };
    }
  
    /**
     * A utility method that adapts {@link PairFunction} to a {@link PairFlatMapFunction} with an
     * {@link Iterator} input. This is particularly useful because it allows to use functions written
     * for mapToPair functions in flatmapToPair functions.
     *
     * @param pairFunction the {@link PairFunction} to adapt.
     * @param <T> the input type.
     * @param <K> the output key type.
     * @param <V> the output value type.
     * @return a {@link PairFlatMapFunction} that accepts an {@link Iterator} as an input and applies
     *     the {@link PairFunction} on every element.
     */
    public static <T, K, V> PairFlatMapFunction<Iterator<T>, K, V> pairFunctionToPairFlatMapFunction(
        final PairFunction<T, K, V> pairFunction) {
      return itr ->
          Iterators.transform(
              itr,
              t -> {
                try {
                  return pairFunction.call(t);
                } catch (Exception e) {
                  throw new RuntimeException(e);
                }
              });
    }
  
    /**
     * A utility method that adapts {@link Function} to a {@link FlatMapFunction} with an {@link
     * Iterator} input. This is particularly useful because it allows to use functions written for map
     * functions in flatmap functions.
     *
     * @param func the {@link Function} to adapt.
     * @param <InputT> the input type.
     * @param <OutputT> the output type.
     * @return a {@link FlatMapFunction} that accepts an {@link Iterator} as an input and applies the
     *     {@link Function} on every element.
     */
    public static <InputT, OutputT>
        FlatMapFunction<Iterator<InputT>, OutputT> functionToFlatMapFunction(
            final Function<InputT, OutputT> func) {
      return itr ->
          Iterators.transform(
              itr,
              t -> {
                try {
                  return func.call(t);
                } catch (Exception e) {
                  throw new RuntimeException(e);
                }
              });
    }
  
    /**
     * Utility to get mapping between TupleTag and a coder.
     *
     * @param outputs - A map of tuple tags and pcollections
     * @return mapping between TupleTag and a coder
     */
    public static Map<TupleTag<?>, Coder<WindowedValue<?>>> getTupleTagCoders(
        Map<TupleTag<?>, PCollection<?>> outputs) {
      Map<TupleTag<?>, Coder<WindowedValue<?>>> coderMap = new HashMap<>(outputs.size());
  
      for (Map.Entry<TupleTag<?>, PCollection<?>> output : outputs.entrySet()) {
        // we get the first PValue as all of them are fro the same type.
        PCollection<?> pCollection = output.getValue();
        Coder<?> coder = pCollection.getCoder();
        Coder<? extends BoundedWindow> wCoder =
            pCollection.getWindowingStrategy().getWindowFn().windowCoder();
        @SuppressWarnings("unchecked")
        Coder<WindowedValue<?>> windowedValueCoder =
            (Coder<WindowedValue<?>>) (Coder<?>) WindowedValues.getFullCoder(coder, wCoder);
        coderMap.put(output.getKey(), windowedValueCoder);
      }
      return coderMap;
    }
  
    /**
     * Returns a pair function to convert value to bytes via coder.
     *
     * @param coderMap - mapping between TupleTag and a coder
     * @return a pair function to convert value to bytes via coder
     */
    public static PairFunction<
            Tuple2<TupleTag<?>, WindowedValue<?>>,
            TupleTag<?>,
            ValueAndCoderLazySerializable<WindowedValue<?>>>
        getTupleTagEncodeFunction(final Map<TupleTag<?>, Coder<WindowedValue<?>>> coderMap) {
      return tuple2 -> {
        TupleTag<?> tupleTag = tuple2._1;
        WindowedValue<?> windowedValue = tuple2._2;
        return new Tuple2<>(
            tupleTag, ValueAndCoderLazySerializable.of(windowedValue, coderMap.get(tupleTag)));
      };
    }
  
    /**
     * Returns a pair function to convert bytes to value via coder.
     *
     * @param coderMap - mapping between TupleTag and a coder
     * @return a pair function to convert bytes to value via coder
     */
    public static PairFunction<
            Tuple2<TupleTag<?>, ValueAndCoderLazySerializable<WindowedValue<?>>>,
            TupleTag<?>,
            WindowedValue<?>>
        getTupleTagDecodeFunction(final Map<TupleTag<?>, Coder<WindowedValue<?>>> coderMap) {
      return tuple2 -> {
        TupleTag<?> tupleTag = tuple2._1;
        ValueAndCoderLazySerializable<WindowedValue<?>> windowedByteValue = tuple2._2;
        return new Tuple2<>(tupleTag, windowedByteValue.getOrDecode(coderMap.get(tupleTag)));
      };
    }
  
    /**
     * Check if we can avoid Serialization. It is only relevant to RDDs. DStreams are memory ser in
     * spark.
     *
     * <p>See <a
     * href="https://spark.apache.org/docs/latest/rdd-programming-guide.html#rdd-persistence">Spark
     * RDD programming guide</a> for more details.
     *
     * @param level StorageLevel required
     * @return true if the level is memory only
     */
    static boolean canAvoidRddSerialization(StorageLevel level) {
      return level.equals(StorageLevel.MEMORY_ONLY());
    }
  }
  