/*
 * 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 java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeMap;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.spark.util.SideInputBroadcast;
import org.apache.beam.runners.spark.util.SparkSideInputReader;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.IterableCoder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.CombineWithContext;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.apache.spark.api.java.function.Function;
import org.joda.time.Instant;

/**
 * A {@link org.apache.beam.sdk.transforms.CombineFnBase.GlobalCombineFn} with a {@link
 * org.apache.beam.sdk.transforms.CombineWithContext.Context} for the SparkRunner.
 */
public class SparkCombineFn<InputT, ValueT, AccumT, OutputT> implements Serializable {

  /** Accumulator of WindowedValues holding values for different windows. */
  public interface WindowedAccumulator<
      InputT, ValueT, AccumT, ImplT extends WindowedAccumulator<InputT, ValueT, AccumT, ImplT>> {

    /**
     * Type of the accumulator. The type depends (mostly) on {@link WindowingStrategy}, and more
     * specialized versions enable more optimizations.
     */
    enum Type {
      MERGING,
      NON_MERGING,
      EXPLODE_WINDOWS,
      SINGLE_WINDOW;

      boolean isMapBased() {
        return this == NON_MERGING || this == MERGING;
      }
    }

    /**
     * Check if this accumulator is empty.
     *
     * @return {@code true} if this accumulator is empty
     */
    boolean isEmpty();

    /** Add value with unexploded windows into the accumulator. */
    void add(WindowedValue<InputT> value, SparkCombineFn<InputT, ValueT, AccumT, ?> context)
        throws Exception;

    /**
     * Merge other acccumulator into this one.
     *
     * @param other the other accumulator to merge
     */
    void merge(ImplT other, SparkCombineFn<?, ?, AccumT, ?> context) throws Exception;

    /** Extract output. */
    Collection<WindowedValue<AccumT>> extractOutput();

    /** Create concrete accumulator for given type. */
    static <InputT, ValueT, AccumT> WindowedAccumulator<InputT, ValueT, AccumT, ?> create(
        SparkCombineFn<InputT, ValueT, AccumT, ?> context,
        Function<InputT, ValueT> toValue,
        WindowingStrategy<?, ?> windowingStrategy,
        Comparator<BoundedWindow> windowComparator) {
      return create(toValue, context.getType(windowingStrategy), windowComparator);
    }

    static <InputT, ValueT, AccumT> WindowedAccumulator<InputT, ValueT, AccumT, ?> create(
        Function<InputT, ValueT> toValue, Type type, Comparator<BoundedWindow> windowComparator) {
      switch (type) {
        case MERGING:
          return MergingWindowedAccumulator.create(toValue, windowComparator);
        case NON_MERGING:
          return NonMergingWindowedAccumulator.create(toValue);
        case SINGLE_WINDOW:
        case EXPLODE_WINDOWS:
          return SingleWindowWindowedAccumulator.create(toValue);
        default:
          throw new IllegalArgumentException("Unknown type: " + type);
      }
    }

    /** Create concrete accumulator for given type. */
    static <InputT, ValueT, AccumT> WindowedAccumulator<InputT, ValueT, AccumT, ?> create(
        Function<InputT, ValueT> toValue,
        Type type,
        Iterable<WindowedValue<AccumT>> values,
        Comparator<BoundedWindow> windowComparator) {
      switch (type) {
        case MERGING:
          return MergingWindowedAccumulator.from(toValue, values, windowComparator);
        case NON_MERGING:
          return NonMergingWindowedAccumulator.from(toValue, values);
        case SINGLE_WINDOW:
        case EXPLODE_WINDOWS:
          Iterator<WindowedValue<AccumT>> iter = values.iterator();
          if (iter.hasNext()) {
            return SingleWindowWindowedAccumulator.create(toValue, iter.next());
          }
          return SingleWindowWindowedAccumulator.create(toValue);
        default:
          throw new IllegalArgumentException("Unknown type: " + type);
      }
    }
  }

  /**
   * Accumulator for {@link WindowFn WindowFns} which create only single window per key, and
   * therefore don't need any map to hold different windows per key.
   *
   * @param <AccumT> type of accumulator
   */
  static class SingleWindowWindowedAccumulator<InputT, ValueT, AccumT>
      implements WindowedAccumulator<
          InputT, ValueT, AccumT, SingleWindowWindowedAccumulator<InputT, ValueT, AccumT>> {

    static <InputT, ValueT, AccumT> SingleWindowWindowedAccumulator<InputT, ValueT, AccumT> create(
        Function<InputT, ValueT> toValue) {
      return new SingleWindowWindowedAccumulator<>(toValue);
    }

    static <InputT, ValueT, AccumT> WindowedAccumulator<InputT, ValueT, AccumT, ?> create(
        Function<InputT, ValueT> toValue, WindowedValue<AccumT> accumulator) {
      return new SingleWindowWindowedAccumulator<>(toValue, accumulator);
    }

    final Function<InputT, ValueT> toValue;
    AccumT windowAccumulator = null;
    Instant accTimestamp = null;
    BoundedWindow accWindow = null;

    SingleWindowWindowedAccumulator(Function<InputT, ValueT> toValue) {
      this.toValue = toValue;
    }

    SingleWindowWindowedAccumulator(
        Function<InputT, ValueT> toValue, WindowedValue<AccumT> accumulator) {
      this.toValue = toValue;
      this.windowAccumulator = accumulator.getValue();
      this.accTimestamp =
          accumulator.getTimestamp().equals(BoundedWindow.TIMESTAMP_MIN_VALUE)
              ? null
              : accumulator.getTimestamp();
      this.accWindow = getWindow(accumulator);
    }

    @Override
    public void add(WindowedValue<InputT> value, SparkCombineFn<InputT, ValueT, AccumT, ?> context)
        throws Exception {
      BoundedWindow window = getWindow(value);
      SparkCombineContext ctx = context.ctxtForValue(value);
      TimestampCombiner combiner = context.windowingStrategy.getTimestampCombiner();
      Instant windowTimestamp =
          combiner.assign(
              window,
              context.windowingStrategy.getWindowFn().getOutputTime(value.getTimestamp(), window));
      final AccumT acc;
      final Instant timestamp;
      if (windowAccumulator == null) {
        acc = context.combineFn.createAccumulator(ctx);
        timestamp = windowTimestamp;
      } else {
        acc = windowAccumulator;
        timestamp = accTimestamp;
      }
      AccumT result = context.combineFn.addInput(acc, toValue(value), ctx);
      Instant timestampCombined = combiner.combine(windowTimestamp, timestamp);
      windowAccumulator = result;
      accTimestamp = timestampCombined;
      accWindow = window;
    }

    @Override
    public void merge(
        SingleWindowWindowedAccumulator<InputT, ValueT, AccumT> other,
        SparkCombineFn<?, ?, AccumT, ?> context) {
      if (windowAccumulator != null && other.windowAccumulator != null) {
        List<AccumT> accumulators = Arrays.asList(windowAccumulator, other.windowAccumulator);
        AccumT merged =
            context.combineFn.mergeAccumulators(
                accumulators, context.ctxtForWindows(Arrays.asList(accWindow)));
        Instant combined =
            context
                .windowingStrategy
                .getTimestampCombiner()
                .combine(accTimestamp, other.accTimestamp);
        windowAccumulator = merged;
        accTimestamp = combined;
      } else if (windowAccumulator == null) {
        windowAccumulator = other.windowAccumulator;
        accTimestamp = other.accTimestamp;
        accWindow = other.accWindow;
      }
    }

    @Override
    public Collection<WindowedValue<AccumT>> extractOutput() {
      if (windowAccumulator != null) {
        return Arrays.asList(
            WindowedValue.of(
                windowAccumulator, accTimestamp, accWindow, PaneInfo.ON_TIME_AND_ONLY_FIRING));
      }
      return Collections.emptyList();
    }

    private ValueT toValue(WindowedValue<InputT> input) {
      try {
        return toValue.call(input.getValue());
      } catch (Exception ex) {
        throw UserCodeException.wrap(ex);
      }
    }

    @Override
    public boolean isEmpty() {
      return windowAccumulator == null;
    }
  }

  /**
   * Abstract base class for {@link WindowedAccumulator WindowedAccumulators} which hold
   * accumulators in map.
   *
   * @param <AccumT> type of accumulator
   * @param <ImplT> type of final subclass
   */
  abstract static class MapBasedWindowedAccumulator<
          InputT,
          ValueT,
          AccumT,
          ImplT extends MapBasedWindowedAccumulator<InputT, ValueT, AccumT, ImplT>>
      implements WindowedAccumulator<InputT, ValueT, AccumT, ImplT> {

    final Function<InputT, ValueT> toValue;
    final Map<BoundedWindow, WindowedValue<AccumT>> map;

    MapBasedWindowedAccumulator(
        Function<InputT, ValueT> toValue, Map<BoundedWindow, WindowedValue<AccumT>> map) {
      this.toValue = toValue;
      this.map = map;
    }

    @Override
    public void add(WindowedValue<InputT> value, SparkCombineFn<InputT, ValueT, AccumT, ?> context)
        throws Exception {
      for (WindowedValue<InputT> v : value.explodeWindows()) {
        SparkCombineContext ctx = context.ctxtForValue(v);
        BoundedWindow window = getWindow(v);
        TimestampCombiner combiner = context.windowingStrategy.getTimestampCombiner();
        Instant windowTimestamp =
            combiner.assign(
                window,
                context.windowingStrategy.getWindowFn().getOutputTime(v.getTimestamp(), window));
        map.compute(
            window,
            (w, windowAccumulator) -> {
              final AccumT acc;
              final Instant timestamp;
              if (windowAccumulator == null) {
                acc = context.combineFn.createAccumulator(ctx);
                timestamp = windowTimestamp;
              } else {
                acc = windowAccumulator.getValue();
                timestamp = windowAccumulator.getTimestamp();
              }
              AccumT result = context.combineFn.addInput(acc, toValue(v), ctx);
              Instant timestampCombined = combiner.combine(windowTimestamp, timestamp);
              return WindowedValue.of(result, timestampCombined, window, PaneInfo.NO_FIRING);
            });
      }
      mergeWindows(context);
    }

    @Override
    public void merge(ImplT other, SparkCombineFn<?, ?, AccumT, ?> context) throws Exception {
      other.map.forEach(
          (window, acc) -> {
            WindowedValue<AccumT> thisAcc = this.map.get(window);
            if (thisAcc == null) {
              // just copy
              this.map.put(window, acc);
            } else {
              // merge
              this.map.put(
                  window,
                  WindowedValue.of(
                      context.combineFn.mergeAccumulators(
                          Lists.newArrayList(thisAcc.getValue(), acc.getValue()),
                          context.ctxtForValue(acc)),
                      context
                          .windowingStrategy
                          .getTimestampCombiner()
                          .combine(acc.getTimestamp(), thisAcc.getTimestamp()),
                      window,
                      PaneInfo.NO_FIRING));
            }
          });
      mergeWindows(context);
    }

    @Override
    public Collection<WindowedValue<AccumT>> extractOutput() {
      return map.values();
    }

    @Override
    public boolean isEmpty() {
      return map.isEmpty();
    }

    void mergeWindows(SparkCombineFn<?, ?, AccumT, ?> fn) throws Exception {}

    private ValueT toValue(WindowedValue<InputT> value) {
      try {
        return toValue.call(value.getValue());
      } catch (Exception ex) {
        throw UserCodeException.wrap(ex);
      }
    }
  }

  /** Accumulator for non-merging windows. */
  static class NonMergingWindowedAccumulator<InputT, ValueT, AccumT>
      extends MapBasedWindowedAccumulator<
          InputT, ValueT, AccumT, NonMergingWindowedAccumulator<InputT, ValueT, AccumT>> {

    static <InputT, ValueT, AccumT> NonMergingWindowedAccumulator<InputT, ValueT, AccumT> create(
        Function<InputT, ValueT> toValue) {
      return new NonMergingWindowedAccumulator<>(toValue);
    }

    static <InputT, ValueT, AccumT> NonMergingWindowedAccumulator<InputT, ValueT, AccumT> from(
        Function<InputT, ValueT> toValue, Iterable<WindowedValue<AccumT>> values) {
      return new NonMergingWindowedAccumulator<>(toValue, values);
    }

    @SuppressWarnings("unchecked")
    private NonMergingWindowedAccumulator(Function<InputT, ValueT> toValue) {
      super(toValue, new HashMap<>());
    }

    private NonMergingWindowedAccumulator(
        Function<InputT, ValueT> toValue, Iterable<WindowedValue<AccumT>> values) {
      super(toValue, asMap(values, new HashMap<>()));
    }
  }

  /** Accumulator for merging windows. */
  static class MergingWindowedAccumulator<InputT, ValueT, AccumT>
      extends MapBasedWindowedAccumulator<
          InputT, ValueT, AccumT, MergingWindowedAccumulator<InputT, ValueT, AccumT>> {

    static <InputT, ValueT, AccumT> MergingWindowedAccumulator<InputT, ValueT, AccumT> create(
        Function<InputT, ValueT> toValue, Comparator<BoundedWindow> windowComparator) {
      return new MergingWindowedAccumulator<>(toValue, windowComparator);
    }

    static <InputT, ValueT, AccumT> MergingWindowedAccumulator<InputT, ValueT, AccumT> from(
        Function<InputT, ValueT> toValue,
        Iterable<WindowedValue<AccumT>> values,
        Comparator<BoundedWindow> windowComparator) {
      return new MergingWindowedAccumulator<>(toValue, values, windowComparator);
    }

    @SuppressWarnings("unchecked")
    private MergingWindowedAccumulator(
        Function<InputT, ValueT> toValue, Comparator<BoundedWindow> windowComparator) {
      super(toValue, new TreeMap<>(windowComparator));
    }

    private MergingWindowedAccumulator(
        Function<InputT, ValueT> toValue,
        Iterable<WindowedValue<AccumT>> values,
        Comparator<BoundedWindow> windowComparator) {
      super(toValue, asMap(values, new TreeMap<>(windowComparator)));
    }

    @Override
    void mergeWindows(SparkCombineFn<?, ?, AccumT, ?> fn) throws Exception {

      SparkCombineContext ctx = fn.ctxtForWindows(this.map.keySet());

      @SuppressWarnings("unchecked")
      WindowFn<Object, BoundedWindow> windowFn = (WindowFn) fn.windowingStrategy.getWindowFn();
      windowFn.mergeWindows(
          asMergeContext(
              windowFn,
              (a, b) -> fn.combineFn.mergeAccumulators(Lists.newArrayList(a, b), ctx),
              (toBeMerged, mergeResult) -> {
                Instant mergedInstant =
                    fn.windowingStrategy
                        .getTimestampCombiner()
                        .merge(
                            mergeResult.getKey(),
                            toBeMerged.stream()
                                .map(w -> map.get(w).getTimestamp())
                                .collect(Collectors.toList()));
                toBeMerged.forEach(this.map::remove);
                this.map.put(
                    mergeResult.getKey(),
                    WindowedValue.of(
                        mergeResult.getValue(),
                        mergedInstant,
                        mergeResult.getKey(),
                        PaneInfo.NO_FIRING));
              },
              map));
    }

    private WindowFn<Object, BoundedWindow>.MergeContext asMergeContext(
        WindowFn<Object, BoundedWindow> windowFn,
        BiFunction<AccumT, AccumT, AccumT> mergeFn,
        BiConsumer<Collection<BoundedWindow>, KV<BoundedWindow, AccumT>> afterMerge,
        Map<BoundedWindow, WindowedValue<AccumT>> map) {

      return windowFn.new MergeContext() {

        @Override
        public Collection<BoundedWindow> windows() {
          return map.keySet();
        }

        @Override
        public void merge(Collection<BoundedWindow> toBeMerged, BoundedWindow mergeResult) {
          AccumT accumulator = null;
          for (BoundedWindow w : toBeMerged) {
            WindowedValue<AccumT> windowAccumulator = Objects.requireNonNull(map.get(w));
            if (accumulator == null) {
              accumulator = windowAccumulator.getValue();
            } else {
              accumulator = mergeFn.apply(accumulator, windowAccumulator.getValue());
            }
          }
          afterMerge.accept(toBeMerged, KV.of(mergeResult, accumulator));
        }
      };
    }

    @Override
    public String toString() {
      return "MergingWindowedAccumulator(" + this.map + ")";
    }
  }

  static class WindowedAccumulatorCoder<InputT, ValueT, AccumT>
      extends Coder<WindowedAccumulator<InputT, ValueT, AccumT, ?>> {

    private final Function<InputT, ValueT> toValue;
    private final IterableCoder<WindowedValue<AccumT>> wrap;
    private final Coder<WindowedValue<AccumT>> accumCoder;
    private final Comparator<BoundedWindow> windowComparator;
    private final WindowedAccumulator.Type type;

    @SuppressWarnings("unchecked")
    WindowedAccumulatorCoder(
        Function<InputT, ValueT> toValue,
        Coder<BoundedWindow> windowCoder,
        Comparator<BoundedWindow> windowComparator,
        Coder<AccumT> accumCoder,
        WindowedAccumulator.Type type) {

      this.toValue = toValue;
      this.accumCoder = WindowedValue.FullWindowedValueCoder.of(accumCoder, windowCoder);
      this.windowComparator = windowComparator;
      this.wrap = IterableCoder.of(this.accumCoder);
      this.type = type;
    }

    @Override
    public void encode(WindowedAccumulator<InputT, ValueT, AccumT, ?> value, OutputStream outStream)
        throws CoderException, IOException {
      if (type.isMapBased()) {
        wrap.encode(((MapBasedWindowedAccumulator<?, ?, AccumT, ?>) value).map.values(), outStream);
      } else {
        SingleWindowWindowedAccumulator<?, ?, AccumT> swwa =
            (SingleWindowWindowedAccumulator<?, ?, AccumT>) value;
        if (swwa.isEmpty()) {
          outStream.write(0);
        } else {
          outStream.write(1);
          accumCoder.encode(
              WindowedValue.of(
                  swwa.windowAccumulator, swwa.accTimestamp, swwa.accWindow, PaneInfo.NO_FIRING),
              outStream);
        }
      }
    }

    @Override
    public WindowedAccumulator<InputT, ValueT, AccumT, ?> decode(InputStream inStream)
        throws CoderException, IOException {
      if (type.isMapBased()) {
        return WindowedAccumulator.create(toValue, type, wrap.decode(inStream), windowComparator);
      }
      boolean empty = inStream.read() == 0;
      if (empty) {
        return WindowedAccumulator.create(toValue, type, windowComparator);
      }
      return WindowedAccumulator.create(
          toValue, type, Arrays.asList(accumCoder.decode(inStream)), windowComparator);
    }

    @Override
    public List<? extends Coder<?>> getCoderArguments() {
      return wrap.getComponents();
    }

    @Override
    public void verifyDeterministic() throws NonDeterministicException {}
  }

  /** An implementation of {@link CombineWithContext.Context} for the SparkRunner. */
  static class SparkCombineContext extends CombineWithContext.Context {
    private final PipelineOptions pipelineOptions;
    private final SideInputReader sideInputReader;

    SparkCombineContext(PipelineOptions pipelineOptions, SideInputReader sideInputReader) {
      this.pipelineOptions = pipelineOptions;
      this.sideInputReader = sideInputReader;
    }

    Collection<? extends BoundedWindow> windows = null;

    SparkCombineContext forInput(final Collection<? extends BoundedWindow> windows) {
      this.windows = Objects.requireNonNull(windows);
      return this;
    }

    @Override
    public PipelineOptions getPipelineOptions() {
      return pipelineOptions;
    }

    @Override
    public <T> T sideInput(PCollectionView<T> view) {
      // validate element window.
      Preconditions.checkState(
          windows.size() == 1,
          "sideInput can only be called when the main " + "input element is in exactly one window");
      return sideInputReader.get(view, windows.iterator().next());
    }
  }

  @VisibleForTesting
  static <K, V, AccumT, OutputT> SparkCombineFn<KV<K, V>, V, AccumT, OutputT> keyed(
      CombineWithContext.CombineFnWithContext<V, AccumT, OutputT> combineFn,
      SerializablePipelineOptions options,
      Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs,
      WindowingStrategy<?, ?> windowingStrategy,
      WindowedAccumulator.Type nonMergingStrategy) {
    return new SparkCombineFn<>(
        false, KV::getValue, combineFn, options, sideInputs, windowingStrategy, nonMergingStrategy);
  }

  public static <K, V, AccumT, OutputT> SparkCombineFn<KV<K, V>, V, AccumT, OutputT> keyed(
      CombineWithContext.CombineFnWithContext<V, AccumT, OutputT> combineFn,
      SerializablePipelineOptions options,
      Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs,
      WindowingStrategy<?, ?> windowingStrategy) {
    return new SparkCombineFn<>(
        false, KV::getValue, combineFn, options, sideInputs, windowingStrategy);
  }

  public static <InputT, AccumT, OutputT> SparkCombineFn<InputT, InputT, AccumT, OutputT> globally(
      CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn,
      SerializablePipelineOptions options,
      Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs,
      WindowingStrategy<?, ?> windowingStrategy) {
    return new SparkCombineFn<>(true, e -> e, combineFn, options, sideInputs, windowingStrategy);
  }

  private static <T> Map<BoundedWindow, WindowedValue<T>> asMap(
      Iterable<WindowedValue<T>> values, Map<BoundedWindow, WindowedValue<T>> res) {
    for (WindowedValue<T> v : values) {
      res.put(getWindow(v), v);
    }
    return res;
  }

  /**
   * Create comparator for given type descriptor. Note that the returned {@link Comparator} has to
   * be {@link Serializable}.
   */
  @SuppressWarnings({"unchecked", "rawtypes"})
  private static Comparator<BoundedWindow> asWindowComparator(
      @Nullable TypeDescriptor<?> windowType) {
    final Comparator<BoundedWindow> comparator;
    if (windowType != null
        && StreamSupport.stream(windowType.getInterfaces().spliterator(), false)
            .anyMatch(t -> t.isSubtypeOf(TypeDescriptor.of(Comparable.class)))) {
      comparator = (Comparator) Comparator.naturalOrder();
    } else {
      java.util.function.Function<BoundedWindow, Instant> keyExtractor =
          (java.util.function.Function<BoundedWindow, Instant> & Serializable)
              BoundedWindow::maxTimestamp;
      comparator = Comparator.comparing(keyExtractor);
    }
    return comparator;
  }

  private final boolean globalCombine;
  private final SerializablePipelineOptions options;
  private final Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs;
  final WindowingStrategy<?, BoundedWindow> windowingStrategy;

  private final Function<InputT, ValueT> toValue;
  private final WindowedAccumulator.Type defaultNonMergingCombineStrategy;
  private final CombineWithContext.CombineFnWithContext<ValueT, AccumT, OutputT> combineFn;
  private final Comparator<BoundedWindow> windowComparator;

  // each Spark task should get it's own copy of this SparkCombineFn, and since Spark tasks
  // are single-threaded, it is safe to reuse the context.
  // the combine context is not Serializable so we'll use lazy initialization.
  // ** DO NOT attempt to turn this into a Singleton as Spark may run multiple tasks in parallel
  // in the same JVM (Executor). **
  // ** DO NOT use combineContext directly inside this class, use ctxtForValue instead. **
  private transient SparkCombineContext combineContext;

  SparkCombineFn(
      boolean global,
      Function<InputT, ValueT> toValue,
      CombineWithContext.CombineFnWithContext<ValueT, AccumT, OutputT> combineFn,
      SerializablePipelineOptions options,
      Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs,
      WindowingStrategy<?, ?> windowingStrategy) {
    this(
        global,
        toValue,
        combineFn,
        options,
        sideInputs,
        windowingStrategy,
        WindowedAccumulator.Type.EXPLODE_WINDOWS);
  }

  @VisibleForTesting
  SparkCombineFn(
      boolean global,
      Function<InputT, ValueT> toValue,
      CombineWithContext.CombineFnWithContext<ValueT, AccumT, OutputT> combineFn,
      SerializablePipelineOptions options,
      Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs,
      WindowingStrategy<?, ?> windowingStrategy,
      WindowedAccumulator.Type defaultNonMergingCombineStrategy) {

    this.globalCombine = global;
    this.options = options;
    this.sideInputs = sideInputs;
    @SuppressWarnings("unchecked")
    WindowingStrategy<?, BoundedWindow> castStrategy = (WindowingStrategy) windowingStrategy;
    this.windowingStrategy = castStrategy;
    this.toValue = toValue;
    this.defaultNonMergingCombineStrategy = defaultNonMergingCombineStrategy;
    this.combineFn = combineFn;
    @SuppressWarnings("unchecked")
    TypeDescriptor<BoundedWindow> untyped =
        (TypeDescriptor<BoundedWindow>) windowingStrategy.getWindowFn().getWindowTypeDescriptor();
    this.windowComparator = asWindowComparator(untyped);
  }

  /** Create empty combiner. Implements Spark's zeroValue for aggregateFn. */
  WindowedAccumulator<InputT, ValueT, AccumT, ?> createCombiner() {
    return WindowedAccumulator.create(this, toValue, windowingStrategy, windowComparator);
  }

  /**
   * Implements Spark's createCombiner function in:
   *
   * <p>{@link org.apache.spark.rdd.PairRDDFunctions#combineByKey}.
   */
  WindowedAccumulator<InputT, ValueT, AccumT, ?> createCombiner(WindowedValue<InputT> value) {
    try {
      WindowedAccumulator<InputT, ValueT, AccumT, ?> accumulator =
          WindowedAccumulator.create(this, toValue, windowingStrategy, windowComparator);
      accumulator.add(value, this);
      return accumulator;
    } catch (Exception ex) {
      throw new IllegalStateException(ex);
    }
  }

  /**
   * Implements Spark's mergeValue function in:
   *
   * <p>{@link org.apache.spark.rdd.PairRDDFunctions#combineByKey}.
   */
  WindowedAccumulator<InputT, ValueT, AccumT, ?> mergeValue(
      WindowedAccumulator<InputT, ValueT, AccumT, ?> accumulator, WindowedValue<InputT> value) {
    try {
      accumulator.add(value, this);
      return accumulator;
    } catch (Exception ex) {
      throw new IllegalStateException(ex);
    }
  }

  /**
   * Implements Spark's mergeCombiners function in:
   *
   * <p>{@link org.apache.spark.rdd.PairRDDFunctions#combineByKey}.
   */
  @SuppressWarnings({"unchecked", "rawtypes"})
  WindowedAccumulator<InputT, ValueT, AccumT, ?> mergeCombiners(
      WindowedAccumulator ac1, WindowedAccumulator ac2) {
    try {
      ac1.merge(ac2, this);
      return ac1;
    } catch (Exception ex) {
      throw new IllegalStateException(ex);
    }
  }

  Iterable<WindowedValue<OutputT>> extractOutput(WindowedAccumulator<?, ?, AccumT, ?> accumulator) {
    return extractOutputStream(accumulator).collect(Collectors.toList());
  }

  /** Extracts the stream of accumulated values. */
  @Internal
  public Stream<WindowedValue<OutputT>> extractOutputStream(
      WindowedAccumulator<?, ?, AccumT, ?> accumulator) {
    return accumulator.extractOutput().stream()
        .filter(Objects::nonNull)
        .map(
            windowAcc ->
                windowAcc.withValue(
                    combineFn.extractOutput(windowAcc.getValue(), ctxtForValue(windowAcc))));
  }

  WindowedAccumulatorCoder<InputT, ValueT, AccumT> accumulatorCoder(
      Coder<BoundedWindow> windowCoder,
      Coder<AccumT> accumulatorCoder,
      WindowingStrategy<?, ?> windowingStrategy) {
    return new WindowedAccumulatorCoder<>(
        toValue, windowCoder, windowComparator, accumulatorCoder, getType(windowingStrategy));
  }

  CombineWithContext.CombineFnWithContext<ValueT, AccumT, OutputT> getCombineFn() {
    return combineFn;
  }

  boolean mustBringWindowToKey() {
    return !getType(windowingStrategy).isMapBased();
  }

  private WindowedAccumulator.Type getType(WindowingStrategy<?, ?> windowingStrategy) {
    if (windowingStrategy.getWindowFn().isNonMerging()) {
      if (globalCombine) {
        /* global combine must use map-based accumulator to incorporate multiple windows */
        return WindowedAccumulator.Type.NON_MERGING;
      }
      if (windowingStrategy.getWindowFn().assignsToOneWindow()
          && GroupNonMergingWindowsFunctions.isEligibleForGroupByWindow(windowingStrategy)) {
        return WindowedAccumulator.Type.SINGLE_WINDOW;
      }
      return defaultNonMergingCombineStrategy;
    }
    return WindowedAccumulator.Type.MERGING;
  }

  private static BoundedWindow getWindow(WindowedValue<?> value) {
    if (value.isSingleWindowedValue()) {
      return ((WindowedValue.SingleWindowedValue) value).getWindow();
    }
    return Iterables.getOnlyElement(value.getWindows());
  }

  @SuppressWarnings("unchecked")
  SparkCombineContext ctxtForValue(WindowedValue<?> input) {
    return ctxtForWindows((Collection) input.getWindows());
  }

  SparkCombineContext ctxtForWindows(Collection<BoundedWindow> windows) {
    if (combineContext == null) {
      combineContext = new SparkCombineContext(options.get(), new SparkSideInputReader(sideInputs));
    }
    return combineContext.forInput(windows);
  }
}