/*
 * 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.io.hadoop.format;

import static java.util.Objects.requireNonNull;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkNotNull;

import com.google.auto.value.AutoValue;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.lang.reflect.InvocationTargetException;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.hadoop.SerializableConfiguration;
import org.apache.beam.sdk.io.hadoop.WritableCoder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.DefaultTrigger;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.AtomicDouble;
import org.apache.hadoop.conf.Configurable;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.ObjectWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapred.FileAlreadyExistsException;
import org.apache.hadoop.mapreduce.InputFormat;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.JobID;
import org.apache.hadoop.mapreduce.MRJobConfig;
import org.apache.hadoop.mapreduce.OutputCommitter;
import org.apache.hadoop.mapreduce.OutputFormat;
import org.apache.hadoop.mapreduce.Partitioner;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.TaskAttemptID;
import org.apache.hadoop.mapreduce.TaskID;
import org.apache.hadoop.mapreduce.lib.db.DBConfiguration;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.mapreduce.task.JobContextImpl;
import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl;
import org.joda.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A {@link HadoopFormatIO} is a Transform for reading data from any source or writing data to any
 * sink which implements Hadoop {@link InputFormat} or {@link OutputFormat}. For example: Cassandra,
 * Elasticsearch, HBase, Redis, Postgres etc. {@link HadoopFormatIO} has to make several performance
 * trade-offs in connecting to {@link InputFormat} or {@link OutputFormat}, so if there is another
 * Beam IO Transform specifically for connecting to your data source of choice, we would recommend
 * using that one, but this IO Transform allows you to connect to many data sources/sinks that do
 * not yet have a Beam IO Transform.
 *
 * <h3>Reading using Hadoop {@link HadoopFormatIO}</h3>
 *
 * <p>You will need to pass a Hadoop {@link Configuration} with parameters specifying how the read
 * will occur. Many properties of the Configuration are optional, and some are required for certain
 * {@link InputFormat} classes, but the following properties must be set for all InputFormats:
 *
 * <ul>
 *   <li>{@code mapreduce.job.inputformat.class}: The {@link InputFormat} class used to connect to
 *       your data source of choice.
 *   <li>{@code key.class}: The key class returned by the {@link InputFormat} in {@code
 *       mapreduce.job.inputformat.class}.
 *   <li>{@code value.class}: The value class returned by the {@link InputFormat} in {@code
 *       mapreduce.job.inputformat.class}.
 * </ul>
 *
 * For example:
 *
 * <pre>
 * {
 *   Configuration myHadoopConfiguration = new Configuration(false);
 *   // Set Hadoop InputFormat, key and value class in configuration
 *   myHadoopConfiguration.setClass(&quot;mapreduce.job.inputformat.class&quot;,
 *      MyDbInputFormatClass, InputFormat.class);
 *   myHadoopConfiguration.setClass(&quot;key.class&quot;, MyDbInputFormatKeyClass, Object.class);
 *   myHadoopConfiguration.setClass(&quot;value.class&quot;,
 *      MyDbInputFormatValueClass, Object.class);
 * }
 * </pre>
 *
 * <p>You will need to check to see if the key and value classes output by the {@link InputFormat}
 * have a Beam {@link Coder} available. If not, you can use withKeyTranslation/withValueTranslation
 * to specify a method transforming instances of those classes into another class that is supported
 * by a Beam {@link Coder}. These settings are optional and you don't need to specify translation
 * for both key and value. If you specify a translation, you will need to make sure the K or V of
 * the read transform match the output type of the translation.
 *
 * <p>You will need to set appropriate InputFormat key and value class (i.e. "key.class" and
 * "value.class") in Hadoop {@link Configuration}. If you set different InputFormat key or value
 * class than InputFormat's actual key or value class then, it may result in an error like
 * "unexpected extra bytes after decoding" while the decoding process of key/value object happens.
 * Hence, it is important to set appropriate InputFormat key and value class.
 *
 * <pre>{@code
 * Pipeline p = ...; // Create pipeline.
 * // Read data only with Hadoop configuration.
 * p.apply("read",
 *     HadoopFormatIO.<InputFormatKeyClass, InputFormatKeyClass>read()
 *              .withConfiguration(myHadoopConfiguration);
 * }
 * // Read data with configuration and key translation (Example scenario: Beam Coder is not
 * available for key class hence key translation is required.).
 * SimpleFunction&lt;InputFormatKeyClass, MyKeyClass&gt; myOutputKeyType =
 *       new SimpleFunction&lt;InputFormatKeyClass, MyKeyClass&gt;() {
 *         public MyKeyClass apply(InputFormatKeyClass input) {
 *           // ...logic to transform InputFormatKeyClass to MyKeyClass
 *         }
 * };
 * </pre>
 *
 * <pre>{@code
 * p.apply("read",
 *     HadoopFormatIO.<MyKeyClass, InputFormatKeyClass>read()
 *              .withConfiguration(myHadoopConfiguration)
 *              .withKeyTranslation(myOutputKeyType);
 * }</pre>
 *
 * <p>// Read data with configuration and value translation (Example scenario: Beam Coder is not
 * available for value class hence value translation is required.).
 *
 * <pre>{@code
 * SimpleFunction&lt;InputFormatValueClass, MyValueClass&gt; myOutputValueType =
 *      new SimpleFunction&lt;InputFormatValueClass, MyValueClass&gt;() {
 *          public MyValueClass apply(InputFormatValueClass input) {
 *            // ...logic to transform InputFormatValueClass to MyValueClass
 *          }
 *  };
 * }</pre>
 *
 * <pre>{@code
 * p.apply("read",
 *     HadoopFormatIO.<InputFormatKeyClass, MyValueClass>read()
 *              .withConfiguration(myHadoopConfiguration)
 *              .withValueTranslation(myOutputValueType);
 * }</pre>
 *
 * <p>IMPORTANT! In case of using {@code DBInputFormat} to read data from RDBMS, Beam parallelizes
 * the process by using LIMIT and OFFSET clauses of SQL query to fetch different ranges of records
 * (as a split) by different workers. To guarantee the same order and proper split of results you
 * need to order them by one or more keys (either PRIMARY or UNIQUE). It can be done during
 * configuration step, for example:
 *
 * <pre>{@code
 * Configuration conf = new Configuration();
 * conf.set(DBConfiguration.INPUT_TABLE_NAME_PROPERTY, tableName);
 * conf.setStrings(DBConfiguration.INPUT_FIELD_NAMES_PROPERTY, "id", "name");
 * conf.set(DBConfiguration.INPUT_ORDER_BY_PROPERTY, "id ASC");
 * }</pre>
 *
 * <h3>Writing using Hadoop {@link HadoopFormatIO}</h3>
 *
 * <p>You will need to pass a Hadoop {@link Configuration} with parameters specifying how the write
 * will occur. Many properties of the Configuration are optional, and some are required for certain
 * {@link OutputFormat} classes, but the following properties must be set for all OutputFormats:
 *
 * <ul>
 *   <li>{@code mapreduce.job.id}: The identifier of the write job. E.g.: end timestamp of window.
 *   <li>{@code mapreduce.job.outputformat.class}: The {@link OutputFormat} class used to connect to
 *       your data sink of choice.
 *   <li>{@code mapreduce.job.output.key.class}: The key class passed to the {@link OutputFormat} in
 *       {@code mapreduce.job.outputformat.class}.
 *   <li>{@code mapreduce.job.output.value.class}: The value class passed to the {@link
 *       OutputFormat} in {@code mapreduce.job.outputformat.class}.
 *   <li>{@code mapreduce.job.reduces}: Number of reduce tasks. Value is equal to number of write
 *       tasks which will be generated. This property is not required for {@link
 *       Write.PartitionedWriterBuilder#withoutPartitioning()} write.
 *   <li>{@code mapreduce.job.partitioner.class}: Hadoop partitioner class which will be used for
 *       distributing of records among partitions. This property is not required for {@link
 *       Write.PartitionedWriterBuilder#withoutPartitioning()} write.
 * </ul>
 *
 * <b>Note:</b> All mentioned values have appropriate constants. E.g.: {@link
 * #OUTPUT_FORMAT_CLASS_ATTR}.
 *
 * <p>For example:
 *
 * <pre>{@code
 * Configuration myHadoopConfiguration = new Configuration(false);
 * // Set Hadoop OutputFormat, key and value class in configuration
 * myHadoopConfiguration.setClass(&quot;mapreduce.job.outputformat.class&quot;,
 *    MyDbOutputFormatClass, OutputFormat.class);
 * myHadoopConfiguration.setClass(&quot;mapreduce.job.output.key.class&quot;,
 *    MyDbOutputFormatKeyClass, Object.class);
 * myHadoopConfiguration.setClass(&quot;mapreduce.job.output.value.class&quot;,
 *    MyDbOutputFormatValueClass, Object.class);
 * myHadoopConfiguration.setClass(&quot;mapreduce.job.partitioner.class&quot;,
 *    MyPartitionerClass, Object.class);
 * myHadoopConfiguration.setInt(&quot;mapreduce.job.reduces&quot;, 2);
 * }</pre>
 *
 * <p>You will need to set OutputFormat key and value class (i.e. "mapreduce.job.output.key.class"
 * and "mapreduce.job.output.value.class") in Hadoop {@link Configuration} which are equal to {@code
 * KeyT} and {@code ValueT}. If you set different OutputFormat key or value class than
 * OutputFormat's actual key or value class then, it will throw {@link IllegalArgumentException}
 *
 * <h4>Batch writing</h4>
 *
 * <pre>{@code
 * //Data which will we want to write
 * PCollection<KV<Text, LongWritable>> boundedWordsCount = ...
 *
 * //Hadoop configuration for write
 * //We have partitioned write, so Partitioner and reducers count have to be set - see withPartitioning() javadoc
 * Configuration myHadoopConfiguration = ...
 * //path to directory with locks
 * String locksDirPath = ...;
 *
 * boundedWordsCount.apply(
 *     "writeBatch",
 *     HadoopFormatIO.<Text, LongWritable>write()
 *         .withConfiguration(myHadoopConfiguration)
 *         .withPartitioning()
 *         .withExternalSynchronization(new HDFSSynchronization(locksDirPath)));
 * }</pre>
 *
 * <h4>Stream writing</h4>
 *
 * <pre>{@code
 *    // Data which will we want to write
 *   PCollection<KV<Text, LongWritable>> unboundedWordsCount = ...;
 *   // Transformation which transforms data of one window into one hadoop configuration
 *   PTransform<PCollection<? extends KV<Text, LongWritable>>, PCollectionView<Configuration>>
 *       configTransform = ...;
 *
 *   unboundedWordsCount.apply(
 *       "writeStream",
 *       HadoopFormatIO.<Text, LongWritable>write()
 *           .withConfigurationTransform(configTransform)
 *           .withExternalSynchronization(new HDFSSynchronization(locksDirPath)));
 * }
 * }</pre>
 */
@Experimental(Experimental.Kind.SOURCE_SINK)
public class HadoopFormatIO {
  private static final Logger LOG = LoggerFactory.getLogger(HadoopFormatIO.class);

  /** {@link MRJobConfig#OUTPUT_FORMAT_CLASS_ATTR}. */
  public static final String OUTPUT_FORMAT_CLASS_ATTR = MRJobConfig.OUTPUT_FORMAT_CLASS_ATTR;

  /** {@link MRJobConfig#OUTPUT_KEY_CLASS}. */
  public static final String OUTPUT_KEY_CLASS = MRJobConfig.OUTPUT_KEY_CLASS;

  /** {@link MRJobConfig#OUTPUT_VALUE_CLASS}. */
  public static final String OUTPUT_VALUE_CLASS = MRJobConfig.OUTPUT_VALUE_CLASS;

  /** {@link MRJobConfig#NUM_REDUCES}. */
  public static final String NUM_REDUCES = MRJobConfig.NUM_REDUCES;

  /** {@link MRJobConfig#PARTITIONER_CLASS_ATTR}. */
  public static final String PARTITIONER_CLASS_ATTR = MRJobConfig.PARTITIONER_CLASS_ATTR;

  /** {@link MRJobConfig#ID}. */
  public static final String JOB_ID = MRJobConfig.ID;

  /** {@link MRJobConfig#MAPREDUCE_JOB_DIR}. */
  public static final String OUTPUT_DIR = FileOutputFormat.OUTDIR;

  /**
   * Creates an uninitialized {@link HadoopFormatIO.Read}. Before use, the {@code Read} must be
   * initialized with a HadoopFormatIO.Read#withConfiguration(HadoopConfiguration) that specifies
   * the source. A key/value translation may also optionally be specified using {@link
   * HadoopFormatIO.Read#withKeyTranslation}/ {@link HadoopFormatIO.Read#withValueTranslation}.
   */
  public static <K, V> Read<K, V> read() {
    return new AutoValue_HadoopFormatIO_Read.Builder<K, V>().build();
  }

  /**
   * Creates an {@link Write.Builder} for creation of Write Transformation. Before creation of the
   * transformation, chain of builders must be set.
   *
   * @param <KeyT> Type of keys to be written.
   * @param <ValueT> Type of values to be written.
   * @return Write builder
   */
  public static <KeyT, ValueT> Write.WriteBuilder<KeyT, ValueT> write() {
    return new Write.Builder<>();
  }

  /**
   * A {@link PTransform} that reads from any data source which implements Hadoop InputFormat. For
   * e.g. Cassandra, Elasticsearch, HBase, Redis, Postgres, etc. See the class-level Javadoc on
   * {@link HadoopFormatIO} for more information.
   *
   * @param <K> Type of keys to be read.
   * @param <V> Type of values to be read.
   * @see HadoopFormatIO
   */
  @AutoValue
  public abstract static class Read<K, V> extends PTransform<PBegin, PCollection<KV<K, V>>> {

    // Returns the Hadoop Configuration which contains specification of source.
    @Nullable
    public abstract SerializableConfiguration getConfiguration();

    @Nullable
    public abstract SimpleFunction<?, K> getKeyTranslationFunction();

    @Nullable
    public abstract SimpleFunction<?, V> getValueTranslationFunction();

    @Nullable
    public abstract TypeDescriptor<K> getKeyTypeDescriptor();

    @Nullable
    public abstract TypeDescriptor<V> getValueTypeDescriptor();

    @Nullable
    public abstract TypeDescriptor<?> getinputFormatClass();

    @Nullable
    public abstract TypeDescriptor<?> getinputFormatKeyClass();

    @Nullable
    public abstract TypeDescriptor<?> getinputFormatValueClass();

    public abstract Builder<K, V> toBuilder();

    @AutoValue.Builder
    abstract static class Builder<K, V> {
      abstract Builder<K, V> setConfiguration(SerializableConfiguration configuration);

      abstract Builder<K, V> setKeyTranslationFunction(SimpleFunction<?, K> function);

      abstract Builder<K, V> setValueTranslationFunction(SimpleFunction<?, V> function);

      abstract Builder<K, V> setKeyTypeDescriptor(TypeDescriptor<K> keyTypeDescriptor);

      abstract Builder<K, V> setValueTypeDescriptor(TypeDescriptor<V> valueTypeDescriptor);

      abstract Builder<K, V> setInputFormatClass(TypeDescriptor<?> inputFormatClass);

      abstract Builder<K, V> setInputFormatKeyClass(TypeDescriptor<?> inputFormatKeyClass);

      abstract Builder<K, V> setInputFormatValueClass(TypeDescriptor<?> inputFormatValueClass);

      abstract Read<K, V> build();
    }

    /** Reads from the source using the options provided by the given configuration. */
    @SuppressWarnings("unchecked")
    public Read<K, V> withConfiguration(Configuration configuration) {
      validateConfiguration(configuration);
      TypeDescriptor<?> inputFormatClass =
          TypeDescriptor.of(configuration.getClass("mapreduce.job.inputformat.class", null));
      TypeDescriptor<?> inputFormatKeyClass =
          TypeDescriptor.of(configuration.getClass("key.class", null));
      TypeDescriptor<?> inputFormatValueClass =
          TypeDescriptor.of(configuration.getClass("value.class", null));
      Builder<K, V> builder =
          toBuilder().setConfiguration(new SerializableConfiguration(configuration));
      builder.setInputFormatClass(inputFormatClass);
      builder.setInputFormatKeyClass(inputFormatKeyClass);
      builder.setInputFormatValueClass(inputFormatValueClass);
      /*
       * Sets the output key class to InputFormat key class if withKeyTranslation() is not called
       * yet.
       */
      if (getKeyTranslationFunction() == null) {
        builder.setKeyTypeDescriptor((TypeDescriptor<K>) inputFormatKeyClass);
      }
      /*
       * Sets the output value class to InputFormat value class if withValueTranslation() is not
       * called yet.
       */
      if (getValueTranslationFunction() == null) {
        builder.setValueTypeDescriptor((TypeDescriptor<V>) inputFormatValueClass);
      }
      return builder.build();
    }

    /** Transforms the keys read from the source using the given key translation function. */
    public Read<K, V> withKeyTranslation(SimpleFunction<?, K> function) {
      checkArgument(function != null, "function can not be null");
      // Sets key class to key translation function's output class type.
      return toBuilder()
          .setKeyTranslationFunction(function)
          .setKeyTypeDescriptor(function.getOutputTypeDescriptor())
          .build();
    }

    /** Transforms the values read from the source using the given value translation function. */
    public Read<K, V> withValueTranslation(SimpleFunction<?, V> function) {
      checkArgument(function != null, "function can not be null");
      // Sets value class to value translation function's output class type.
      return toBuilder()
          .setValueTranslationFunction(function)
          .setValueTypeDescriptor(function.getOutputTypeDescriptor())
          .build();
    }

    @Override
    public PCollection<KV<K, V>> expand(PBegin input) {
      validateTransform();
      // Get the key and value coders based on the key and value classes.
      CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
      Coder<K> keyCoder = getDefaultCoder(getKeyTypeDescriptor(), coderRegistry);
      Coder<V> valueCoder = getDefaultCoder(getValueTypeDescriptor(), coderRegistry);
      HadoopInputFormatBoundedSource<K, V> source =
          new HadoopInputFormatBoundedSource<>(
              getConfiguration(),
              keyCoder,
              valueCoder,
              getKeyTranslationFunction(),
              getValueTranslationFunction());
      return input.getPipeline().apply(org.apache.beam.sdk.io.Read.from(source));
    }

    /**
     * Validates that the mandatory configuration properties such as InputFormat class, InputFormat
     * key and value classes are provided in the Hadoop configuration. In case of using {@code
     * DBInputFormat} you need to order results by one or more keys. It can be done by setting
     * configuration option "mapreduce.jdbc.input.orderby".
     */
    private void validateConfiguration(Configuration configuration) {
      checkArgument(configuration != null, "configuration can not be null");
      checkArgument(
          configuration.get("mapreduce.job.inputformat.class") != null,
          "Configuration must contain \"mapreduce.job.inputformat.class\"");
      checkArgument(
          configuration.get("key.class") != null, "configuration must contain \"key.class\"");
      checkArgument(
          configuration.get("value.class") != null, "configuration must contain \"value.class\"");
      if (configuration.get("mapreduce.job.inputformat.class").endsWith("DBInputFormat")) {
        checkArgument(
            configuration.get(DBConfiguration.INPUT_ORDER_BY_PROPERTY) != null,
            "Configuration must contain \""
                + DBConfiguration.INPUT_ORDER_BY_PROPERTY
                + "\" when using DBInputFormat");
      }
    }

    /** Validates construction of this transform. */
    //    @VisibleForTesting
    public void validateTransform() {
      checkArgument(getConfiguration() != null, "withConfiguration() is required");
      // Validate that the key translation input type must be same as key class of InputFormat.
      validateTranslationFunction(
          getinputFormatKeyClass(),
          getKeyTranslationFunction(),
          "Key translation's input type is not same as hadoop InputFormat : %s key class : %s");
      // Validate that the value translation input type must be same as value class of InputFormat.
      validateTranslationFunction(
          getinputFormatValueClass(),
          getValueTranslationFunction(),
          "Value translation's input type is not same as hadoop InputFormat :  "
              + "%s value class : %s");
    }

    /** Validates translation function given for key/value translation. */
    private void validateTranslationFunction(
        TypeDescriptor<?> inputType, SimpleFunction<?, ?> simpleFunction, String errorMsg) {
      if (simpleFunction != null && !simpleFunction.getInputTypeDescriptor().equals(inputType)) {
        throw new IllegalArgumentException(
            String.format(errorMsg, getinputFormatClass().getRawType(), inputType.getRawType()));
      }
    }

    /**
     * Returns the default coder for a given type descriptor. Coder Registry is queried for correct
     * coder, if not found in Coder Registry, then check if the type descriptor provided is of type
     * Writable, then WritableCoder is returned, else exception is thrown "Cannot find coder".
     */
    @SuppressWarnings({"unchecked", "WeakerAccess"})
    public <T> Coder<T> getDefaultCoder(TypeDescriptor<?> typeDesc, CoderRegistry coderRegistry) {
      Class classType = typeDesc.getRawType();
      try {
        return (Coder<T>) coderRegistry.getCoder(typeDesc);
      } catch (CannotProvideCoderException e) {
        if (Writable.class.isAssignableFrom(classType)) {
          return (Coder<T>) WritableCoder.of(classType);
        }
        throw new IllegalStateException(
            String.format("Cannot find coder for %s  : ", typeDesc) + e.getMessage(), e);
      }
    }
  }

  /**
   * Bounded source implementation for {@link HadoopFormatIO}.
   *
   * @param <K> Type of keys to be read.
   * @param <V> Type of values to be read.
   */
  public static class HadoopInputFormatBoundedSource<K, V> extends BoundedSource<KV<K, V>>
      implements Serializable {
    private final SerializableConfiguration conf;
    private final Coder<K> keyCoder;
    private final Coder<V> valueCoder;
    @Nullable private final SimpleFunction<?, K> keyTranslationFunction;
    @Nullable private final SimpleFunction<?, V> valueTranslationFunction;
    private final SerializableSplit inputSplit;
    private transient List<SerializableSplit> inputSplits;
    private long boundedSourceEstimatedSize = 0;
    private transient InputFormat<?, ?> inputFormatObj;
    private transient TaskAttemptContext taskAttemptContext;
    private static final Set<Class<?>> immutableTypes =
        new HashSet<>(
            Arrays.asList(
                String.class,
                Byte.class,
                Short.class,
                Integer.class,
                Long.class,
                Float.class,
                Double.class,
                Boolean.class,
                BigInteger.class,
                BigDecimal.class));

    HadoopInputFormatBoundedSource(
        SerializableConfiguration conf,
        Coder<K> keyCoder,
        Coder<V> valueCoder,
        @Nullable SimpleFunction<?, K> keyTranslationFunction,
        @Nullable SimpleFunction<?, V> valueTranslationFunction) {
      this(conf, keyCoder, valueCoder, keyTranslationFunction, valueTranslationFunction, null);
    }

    @SuppressWarnings("WeakerAccess")
    protected HadoopInputFormatBoundedSource(
        SerializableConfiguration conf,
        Coder<K> keyCoder,
        Coder<V> valueCoder,
        @Nullable SimpleFunction<?, K> keyTranslationFunction,
        @Nullable SimpleFunction<?, V> valueTranslationFunction,
        SerializableSplit inputSplit) {
      this.conf = conf;
      this.inputSplit = inputSplit;
      this.keyCoder = keyCoder;
      this.valueCoder = valueCoder;
      this.keyTranslationFunction = keyTranslationFunction;
      this.valueTranslationFunction = valueTranslationFunction;
    }

    @SuppressWarnings("WeakerAccess")
    public SerializableConfiguration getConfiguration() {
      return conf;
    }

    @Override
    public void validate() {
      checkArgument(conf != null, "conf can not be null");
      checkArgument(keyCoder != null, "keyCoder can not be null");
      checkArgument(valueCoder != null, "valueCoder can not be null");
    }

    @Override
    public void populateDisplayData(DisplayData.Builder builder) {
      super.populateDisplayData(builder);
      Configuration hadoopConfig = getConfiguration().get();
      if (hadoopConfig != null) {
        builder.addIfNotNull(
            DisplayData.item(
                    "mapreduce.job.inputformat.class",
                    hadoopConfig.get("mapreduce.job.inputformat.class"))
                .withLabel("InputFormat Class"));
        builder.addIfNotNull(
            DisplayData.item("key.class", hadoopConfig.get("key.class")).withLabel("Key Class"));
        builder.addIfNotNull(
            DisplayData.item("value.class", hadoopConfig.get("value.class"))
                .withLabel("Value Class"));
      }
    }

    @Override
    public List<BoundedSource<KV<K, V>>> split(long desiredBundleSizeBytes, PipelineOptions options)
        throws Exception {
      // desiredBundleSizeBytes is not being considered as splitting based on this
      // value is not supported by inputFormat getSplits() method.
      if (inputSplit != null) {
        LOG.info("Not splitting source {} because source is already split.", this);
        return ImmutableList.of(this);
      }
      computeSplitsIfNecessary();
      LOG.info(
          "Generated {} splits. Size of first split is {} ",
          inputSplits.size(),
          inputSplits.get(0).getSplit().getLength());
      return inputSplits.stream()
          .map(
              serializableInputSplit ->
                  new HadoopInputFormatBoundedSource<>(
                      conf,
                      keyCoder,
                      valueCoder,
                      keyTranslationFunction,
                      valueTranslationFunction,
                      serializableInputSplit))
          .collect(Collectors.toList());
    }

    @Override
    public long getEstimatedSizeBytes(PipelineOptions po) throws Exception {
      if (inputSplit == null) {
        // If there are no splits computed yet, then retrieve the splits.
        computeSplitsIfNecessary();
        return boundedSourceEstimatedSize;
      }
      return inputSplit.getSplit().getLength();
    }

    /**
     * This is a helper function to compute splits. This method will also calculate size of the data
     * being read. Note: This method is executed exactly once and the splits are retrieved and
     * cached in this. These splits are further used by split() and getEstimatedSizeBytes().
     */
    @VisibleForTesting
    void computeSplitsIfNecessary() throws IOException, InterruptedException {
      if (inputSplits != null) {
        return;
      }
      createInputFormatInstance();
      List<InputSplit> splits = inputFormatObj.getSplits(Job.getInstance(conf.get()));
      if (splits == null) {
        throw new IOException("Error in computing splits, getSplits() returns null.");
      }
      if (splits.isEmpty()) {
        throw new IOException("Error in computing splits, getSplits() returns a empty list");
      }
      boundedSourceEstimatedSize = 0;
      inputSplits = new ArrayList<>();
      for (InputSplit inputSplit : splits) {
        if (inputSplit == null) {
          throw new IOException(
              "Error in computing splits, split is null in InputSplits list "
                  + "populated by getSplits() : ");
        }
        boundedSourceEstimatedSize += inputSplit.getLength();
        inputSplits.add(new SerializableSplit(inputSplit));
      }
    }

    /**
     * Creates instance of InputFormat class. The InputFormat class name is specified in the Hadoop
     * configuration.
     */
    @SuppressWarnings("WeakerAccess")
    protected void createInputFormatInstance() throws IOException {
      if (inputFormatObj == null) {
        try {
          taskAttemptContext = new TaskAttemptContextImpl(conf.get(), new TaskAttemptID());
          inputFormatObj =
              (InputFormat<?, ?>)
                  conf.get()
                      .getClassByName(conf.get().get("mapreduce.job.inputformat.class"))
                      .getConstructor()
                      .newInstance();
          /*
           * If InputFormat explicitly implements interface {@link Configurable}, then setConf()
           * method of {@link Configurable} needs to be explicitly called to set all the
           * configuration parameters. For example: InputFormat classes which implement Configurable
           * are {@link org.apache.hadoop.mapreduce.lib.db.DBInputFormat DBInputFormat}, {@link
           * org.apache.hadoop.hbase.mapreduce.TableInputFormat TableInputFormat}, etc.
           */
          if (Configurable.class.isAssignableFrom(inputFormatObj.getClass())) {
            ((Configurable) inputFormatObj).setConf(conf.get());
          }
        } catch (InstantiationException
            | IllegalAccessException
            | ClassNotFoundException
            | NoSuchMethodException
            | InvocationTargetException e) {
          throw new IOException("Unable to create InputFormat object: ", e);
        }
      }
    }

    @VisibleForTesting
    InputFormat<?, ?> getInputFormat() {
      return inputFormatObj;
    }

    @VisibleForTesting
    void setInputFormatObj(InputFormat<?, ?> inputFormatObj) {
      this.inputFormatObj = inputFormatObj;
    }

    @Override
    public Coder<KV<K, V>> getOutputCoder() {
      return KvCoder.of(keyCoder, valueCoder);
    }

    @Override
    public BoundedReader<KV<K, V>> createReader(PipelineOptions options) throws IOException {
      this.validate();
      if (inputSplit == null) {
        throw new IOException("Cannot create reader as source is not split yet.");
      } else {
        createInputFormatInstance();
        return new HadoopInputFormatReader<>(
            this,
            keyTranslationFunction,
            valueTranslationFunction,
            inputSplit,
            inputFormatObj,
            taskAttemptContext);
      }
    }

    /**
     * BoundedReader for Hadoop InputFormat source.
     *
     * @param <T1> Type of keys RecordReader emits.
     * @param <T2> Type of values RecordReader emits.
     */
    class HadoopInputFormatReader<T1, T2> extends BoundedSource.BoundedReader<KV<K, V>> {

      private final HadoopInputFormatBoundedSource<K, V> source;
      @Nullable private final SimpleFunction<T1, K> keyTranslationFunction;
      @Nullable private final SimpleFunction<T2, V> valueTranslationFunction;
      private final SerializableSplit split;
      private RecordReader<T1, T2> recordReader;
      private volatile boolean doneReading = false;
      private final AtomicLong recordsReturned = new AtomicLong();
      // Tracks the progress of the RecordReader.
      private final AtomicDouble progressValue = new AtomicDouble();
      private final transient InputFormat<T1, T2> inputFormatObj;
      private final transient TaskAttemptContext taskAttemptContext;

      @SuppressWarnings("unchecked")
      private HadoopInputFormatReader(
          HadoopInputFormatBoundedSource<K, V> source,
          @Nullable SimpleFunction keyTranslationFunction,
          @Nullable SimpleFunction valueTranslationFunction,
          SerializableSplit split,
          InputFormat inputFormatObj,
          TaskAttemptContext taskAttemptContext) {
        this.source = source;
        this.keyTranslationFunction = keyTranslationFunction;
        this.valueTranslationFunction = valueTranslationFunction;
        this.split = split;
        this.inputFormatObj = inputFormatObj;
        this.taskAttemptContext = taskAttemptContext;
      }

      @Override
      public HadoopInputFormatBoundedSource<K, V> getCurrentSource() {
        return source;
      }

      @Override
      public boolean start() throws IOException {
        try {
          recordsReturned.set(0L);
          recordReader = inputFormatObj.createRecordReader(split.getSplit(), taskAttemptContext);
          if (recordReader != null) {
            recordReader.initialize(split.getSplit(), taskAttemptContext);
            progressValue.set(getProgress());
            if (recordReader.nextKeyValue()) {
              recordsReturned.incrementAndGet();
              doneReading = false;
              return true;
            }
          } else {
            throw new IOException(
                String.format(
                    "Null RecordReader object returned by %s", inputFormatObj.getClass()));
          }
          recordReader = null;
        } catch (InterruptedException e) {
          throw new IOException(
              "Could not read because the thread got interrupted while "
                  + "reading the records with an exception: ",
              e);
        }
        doneReading = true;
        return false;
      }

      @Override
      public boolean advance() throws IOException {
        try {
          progressValue.set(getProgress());
          if (recordReader.nextKeyValue()) {
            recordsReturned.incrementAndGet();
            return true;
          }
          doneReading = true;
        } catch (InterruptedException e) {
          throw new IOException("Unable to read data: ", e);
        }
        return false;
      }

      @Override
      public KV<K, V> getCurrent() {
        K key;
        V value;
        try {
          // Transform key if translation function is provided.
          key = transformKeyOrValue(recordReader.getCurrentKey(), keyTranslationFunction, keyCoder);
          // Transform value if translation function is provided.
          value =
              transformKeyOrValue(
                  recordReader.getCurrentValue(), valueTranslationFunction, valueCoder);
        } catch (IOException | InterruptedException e) {
          LOG.error("Unable to read data: ", e);
          throw new IllegalStateException("Unable to read data: " + "{}", e);
        }
        return KV.of(key, value);
      }

      /** Returns the serialized output of transformed key or value object. */
      @SuppressWarnings("unchecked")
      private <T, T3> T3 transformKeyOrValue(
          T input, @Nullable SimpleFunction<T, T3> simpleFunction, Coder<T3> coder)
          throws CoderException, ClassCastException {
        T3 output;
        if (null != simpleFunction) {
          output = simpleFunction.apply(input);
        } else {
          output = (T3) input;
        }
        return cloneIfPossiblyMutable(output, coder);
      }

      /**
       * Beam expects immutable objects, but the Hadoop InputFormats tend to re-use the same object
       * when returning them. Hence, mutable objects returned by Hadoop InputFormats are cloned.
       */
      private <T> T cloneIfPossiblyMutable(T input, Coder<T> coder)
          throws CoderException, ClassCastException {
        // If the input object is not of known immutable type, clone the object.
        if (!isKnownImmutable(input)) {
          input = CoderUtils.clone(coder, input);
        }
        return input;
      }

      /** Utility method to check if the passed object is of a known immutable type. */
      private boolean isKnownImmutable(Object o) {
        return immutableTypes.contains(o.getClass());
      }

      @Override
      public void close() throws IOException {
        LOG.info("Closing reader after reading {} records.", recordsReturned);
        if (recordReader != null) {
          recordReader.close();
          recordReader = null;
        }
      }

      @Override
      public Double getFractionConsumed() {
        if (doneReading) {
          return 1.0;
        } else if (recordReader == null || recordsReturned.get() == 0L) {
          return 0.0;
        }
        if (progressValue.get() == 0.0) {
          return null;
        }
        return progressValue.doubleValue();
      }

      /** Returns RecordReader's progress. */
      private Double getProgress() throws IOException, InterruptedException {
        try {
          float progress = recordReader.getProgress();
          return (double) progress < 0 || progress > 1 ? 0.0 : progress;
        } catch (IOException e) {
          LOG.error(
              "Error in computing the fractions consumed as RecordReader.getProgress() throws an "
                  + "exception : ",
              e);
          throw new IOException(
              "Error in computing the fractions consumed as RecordReader.getProgress() throws an "
                  + "exception : "
                  + e.getMessage(),
              e);
        }
      }

      @Override
      public final long getSplitPointsRemaining() {
        if (doneReading) {
          return 0;
        }
        /*
         This source does not currently support dynamic work rebalancing, so remaining parallelism
         is always 1.
        */
        return 1;
      }
    }
  }

  /**
   * A wrapper to allow Hadoop {@link org.apache.hadoop.mapreduce.InputSplit} to be serialized using
   * Java's standard serialization mechanisms.
   */
  public static class SerializableSplit implements Serializable {

    InputSplit inputSplit;

    public SerializableSplit() {}

    public SerializableSplit(InputSplit split) {
      checkArgument(
          split instanceof Writable, String.format("Split is not of type Writable: %s", split));
      this.inputSplit = split;
    }

    @SuppressWarnings("WeakerAccess")
    public InputSplit getSplit() {
      return inputSplit;
    }

    private void readObject(ObjectInputStream in) throws IOException {
      ObjectWritable ow = new ObjectWritable();
      ow.setConf(new Configuration(false));
      ow.readFields(in);
      this.inputSplit = (InputSplit) ow.get();
    }

    private void writeObject(ObjectOutputStream out) throws IOException {
      new ObjectWritable(inputSplit).write(out);
    }
  }

  /**
   * Generates tasks for output pairs and groups them by this key.
   *
   * <p>This transformation is used when is configured write with partitioning.
   *
   * @param <KeyT> type of key
   * @param <ValueT> type of value
   */
  private static class GroupDataByPartition<KeyT, ValueT>
      extends PTransform<
          PCollection<KV<KeyT, ValueT>>, PCollection<KV<Integer, KV<KeyT, ValueT>>>> {

    private final PCollectionView<Configuration> configView;

    private GroupDataByPartition(PCollectionView<Configuration> configView) {
      this.configView = configView;
    }

    @Override
    public PCollection<KV<Integer, KV<KeyT, ValueT>>> expand(PCollection<KV<KeyT, ValueT>> input) {
      return input
          .apply(
              "AssignTask",
              ParDo.of(new AssignTaskFn<KeyT, ValueT>(configView)).withSideInputs(configView))
          .setTypeDescriptor(
              TypeDescriptors.kvs(TypeDescriptors.integers(), input.getTypeDescriptor()))
          .apply("GroupByTaskId", GroupByKey.create())
          .apply("FlattenGroupedTasks", ParDo.of(new FlattenGroupedTasks<>()));
    }
  }

  /**
   * Flattens grouped iterable {@link KV} pairs into triplets of TaskID/Key/Value.
   *
   * @param <KeyT> Type of keys to be written.
   * @param <ValueT> Type of values to be written.
   */
  private static class FlattenGroupedTasks<KeyT, ValueT>
      extends DoFn<KV<Integer, Iterable<KV<KeyT, ValueT>>>, KV<Integer, KV<KeyT, ValueT>>> {

    @ProcessElement
    public void processElement(
        @Element KV<Integer, Iterable<KV<KeyT, ValueT>>> input,
        OutputReceiver<KV<Integer, KV<KeyT, ValueT>>> outputReceiver) {
      final Integer key = input.getKey();
      for (KV<KeyT, ValueT> element :
          requireNonNull(input.getValue(), "Iterable can not be null.")) {
        outputReceiver.output(KV.of(key, element));
      }
    }
  }

  /**
   * A {@link PTransform} that writes to any data sink which implements Hadoop OutputFormat. For
   * e.g. Cassandra, Elasticsearch, HBase, Redis, Postgres, etc. See the class-level Javadoc on
   * {@link HadoopFormatIO} for more information.
   *
   * @param <KeyT> Type of keys to be written.
   * @param <ValueT> Type of values to be written.
   * @see HadoopFormatIO
   */
  public static class Write<KeyT, ValueT> extends PTransform<PCollection<KV<KeyT, ValueT>>, PDone> {

    @Nullable private final transient Configuration configuration;

    @Nullable
    private final PTransform<
            PCollection<? extends KV<KeyT, ValueT>>, PCollectionView<Configuration>>
        configTransform;

    private final ExternalSynchronization externalSynchronization;

    private final boolean withPartitioning;

    Write(
        @Nullable Configuration configuration,
        @Nullable
            PTransform<PCollection<? extends KV<KeyT, ValueT>>, PCollectionView<Configuration>>
                configTransform,
        ExternalSynchronization externalSynchronization,
        boolean withPartitioning) {
      this.configuration = configuration;
      this.configTransform = configTransform;
      this.externalSynchronization = externalSynchronization;
      this.withPartitioning = withPartitioning;
    }

    /**
     * Builder for partitioning determining.
     *
     * @param <KeyT> Key type to write
     * @param <ValueT> Value type to write
     */
    public interface PartitionedWriterBuilder<KeyT, ValueT> {

      /**
       * Writes to the sink with partitioning by Task Id.
       *
       * <p>Following Hadoop configuration properties are required with this option:
       *
       * <ul>
       *   <li>{@code mapreduce.job.reduces}: Number of reduce tasks. Value is equal to number of
       *       write tasks which will be generated.
       *   <li>{@code mapreduce.job.partitioner.class}: Hadoop partitioner class which will be used
       *       for distributing of records among partitions.
       * </ul>
       *
       * @return WriteBuilder for write transformation
       */
      ExternalSynchronizationBuilder<KeyT, ValueT> withPartitioning();

      /**
       * Writes to the sink without need to partition output into specified number of partitions.
       *
       * <p>This write operation doesn't do shuffle by the partition so it saves transfer time
       * before write operation itself. As a consequence it generates random number of partitions.
       *
       * <p><b>Note:</b> Works only for {@link
       * org.apache.beam.sdk.values.PCollection.IsBounded#BOUNDED} {@link PCollection} with global
       * {@link WindowingStrategy}.
       *
       * @return WriteBuilder for write transformation
       */
      ExternalSynchronizationBuilder<KeyT, ValueT> withoutPartitioning();
    }

    /**
     * Builder for External Synchronization defining.
     *
     * @param <KeyT> Key type to write
     * @param <ValueT> Value type to write
     */
    public interface ExternalSynchronizationBuilder<KeyT, ValueT> {

      /**
       * Specifies class which will provide external synchronization required for hadoop write
       * operation.
       *
       * @param externalSynchronization provider of external synchronization
       * @return Write transformation
       */
      Write<KeyT, ValueT> withExternalSynchronization(
          ExternalSynchronization externalSynchronization);
    }

    /**
     * Main builder of Write transformation.
     *
     * @param <KeyT> Key type to write
     * @param <ValueT> Value type to write
     */
    public interface WriteBuilder<KeyT, ValueT> {

      /**
       * Writes to the sink using the options provided by the given hadoop configuration.
       *
       * <p><b>Note:</b> Works only for {@link
       * org.apache.beam.sdk.values.PCollection.IsBounded#BOUNDED} {@link PCollection} with global
       * {@link WindowingStrategy}.
       *
       * @param config hadoop configuration.
       * @return WriteBuilder with set configuration
       * @throws NullPointerException when the configuration is null
       * @see HadoopFormatIO for required hadoop {@link Configuration} properties
       */
      PartitionedWriterBuilder<KeyT, ValueT> withConfiguration(Configuration config);

      /**
       * Writes to the sink using configuration created by provided {@code
       * configurationTransformation}.
       *
       * <p>This type is useful especially for processing unbounded windowed data but can be used
       * also for batch processing.
       *
       * <p>Supports only {@link PCollection} with {@link DefaultTrigger}ing and without allowed
       * lateness
       *
       * @param configTransform configuration transformation interface
       * @return WriteBuilder with set configuration transformation
       * @throws NullPointerException when {@code configurationTransformation} is {@code null}
       * @see HadoopFormatIO for required hadoop {@link Configuration} properties
       */
      ExternalSynchronizationBuilder<KeyT, ValueT> withConfigurationTransform(
          PTransform<PCollection<? extends KV<KeyT, ValueT>>, PCollectionView<Configuration>>
              configTransform);
    }

    /**
     * Implementation of all builders.
     *
     * @param <KeyT> Key type to write
     * @param <ValueT> Value type to write
     */
    static class Builder<KeyT, ValueT>
        implements WriteBuilder<KeyT, ValueT>,
            PartitionedWriterBuilder<KeyT, ValueT>,
            ExternalSynchronizationBuilder<KeyT, ValueT> {

      private Configuration configuration;
      private PTransform<PCollection<? extends KV<KeyT, ValueT>>, PCollectionView<Configuration>>
          configTransform;
      private boolean isWithPartitioning;

      @Override
      public PartitionedWriterBuilder<KeyT, ValueT> withConfiguration(Configuration config) {
        checkNotNull(config, "Hadoop configuration cannot be null");
        this.configuration = config;
        return this;
      }

      @Override
      public ExternalSynchronizationBuilder<KeyT, ValueT> withConfigurationTransform(
          PTransform<PCollection<? extends KV<KeyT, ValueT>>, PCollectionView<Configuration>>
              configTransform) {
        checkNotNull(configTransform, "Configuration transformation cannot be null");
        this.isWithPartitioning = true;
        this.configTransform = configTransform;
        return this;
      }

      @Override
      public ExternalSynchronizationBuilder<KeyT, ValueT> withPartitioning() {
        this.isWithPartitioning = true;
        return this;
      }

      @Override
      public ExternalSynchronizationBuilder<KeyT, ValueT> withoutPartitioning() {
        this.isWithPartitioning = false;
        return this;
      }

      @Override
      public Write<KeyT, ValueT> withExternalSynchronization(
          ExternalSynchronization externalSynchronization) {
        checkNotNull(externalSynchronization, "External synchronization cannot be null");
        return new Write<>(
            this.configuration,
            this.configTransform,
            externalSynchronization,
            this.isWithPartitioning);
      }
    }

    @Override
    public void validate(PipelineOptions pipelineOptions) {}

    @Override
    public void populateDisplayData(DisplayData.Builder builder) {
      super.populateDisplayData(builder);
      Configuration hadoopConfig = configuration;
      if (hadoopConfig != null) {
        builder.addIfNotNull(
            DisplayData.item(OUTPUT_FORMAT_CLASS_ATTR, hadoopConfig.get(OUTPUT_FORMAT_CLASS_ATTR))
                .withLabel("OutputFormat Class"));
        builder.addIfNotNull(
            DisplayData.item(OUTPUT_KEY_CLASS, hadoopConfig.get(OUTPUT_KEY_CLASS))
                .withLabel("OutputFormat Key Class"));
        builder.addIfNotNull(
            DisplayData.item(OUTPUT_VALUE_CLASS, hadoopConfig.get(OUTPUT_VALUE_CLASS))
                .withLabel("OutputFormat Value Class"));
        builder.addIfNotNull(
            DisplayData.item(
                    PARTITIONER_CLASS_ATTR,
                    hadoopConfig.get(
                        PARTITIONER_CLASS_ATTR,
                        HadoopFormats.DEFAULT_PARTITIONER_CLASS_ATTR.getName()))
                .withLabel("Partitioner Class"));
      }
    }

    @Override
    public PDone expand(PCollection<KV<KeyT, ValueT>> input) {

      // streamed pipeline must have defined configuration transformation
      if (input.isBounded().equals(PCollection.IsBounded.UNBOUNDED)
          || !input.getWindowingStrategy().equals(WindowingStrategy.globalDefault())) {
        checkArgument(
            configTransform != null,
            "Writing of unbounded data can be processed only with configuration transformation provider. See %s.withConfigurationTransform()",
            Write.class);
      }

      verifyInputWindowing(input);

      TypeDescriptor<Configuration> configType = new TypeDescriptor<Configuration>() {};
      input
          .getPipeline()
          .getCoderRegistry()
          .registerCoderForType(configType, new ConfigurationCoder());

      PCollectionView<Configuration> configView = createConfigurationView(input);

      return processJob(input, configView);
    }

    /**
     * Processes write job. Write job is composed from following partial steps:
     *
     * <ul>
     *   <li>When partitioning is enabled:
     *       <ul>
     *         <li>Assigning of the {@link TaskID} (represented as {@link Integer}) to the {@link
     *             KV}s in {@link AssignTaskFn}
     *         <li>Grouping {@link KV}s by the {@link TaskID}
     *       </ul>
     *   <li>Otherwise creation of TaskId via {@link PrepareNonPartitionedTasksFn} where locks are
     *       created for each task id
     *   <li>Writing of {@link KV} records via {@link WriteFn}
     *   <li>Committing of whole job via {@link CommitJobFn}
     * </ul>
     *
     * @param input Collection with output data to write
     * @param configView configuration view
     * @return Successfully processed write
     */
    private PDone processJob(
        PCollection<KV<KeyT, ValueT>> input, PCollectionView<Configuration> configView) {

      TypeDescriptor<Iterable<Integer>> iterableIntType =
          TypeDescriptors.iterables(TypeDescriptors.integers());

      PCollection<KV<KeyT, ValueT>> validatedInput =
          input.apply(
              ParDo.of(
                      new SetupJobFn<>(
                          externalSynchronization, configView, input.getTypeDescriptor()))
                  .withSideInputs(configView));

      PCollection<KV<Integer, KV<KeyT, ValueT>>> writeData =
          withPartitioning
              ? validatedInput.apply("GroupDataByPartition", new GroupDataByPartition<>(configView))
              : validatedInput.apply(
                  "PrepareNonPartitionedTasks",
                  ParDo.of(
                          new PrepareNonPartitionedTasksFn<KeyT, ValueT>(
                              configView, externalSynchronization))
                      .withSideInputs(configView));

      PCollection<Iterable<Integer>> collectedFinishedWrites =
          writeData
              .apply(
                  "Write",
                  ParDo.of(new WriteFn<KeyT, ValueT>(configView, externalSynchronization))
                      .withSideInputs(configView))
              .setTypeDescriptor(TypeDescriptors.integers())
              .apply(
                  "CollectWriteTasks",
                  Combine.globally(new IterableCombinerFn<>(TypeDescriptors.integers()))
                      .withoutDefaults())
              .setTypeDescriptor(iterableIntType);

      return PDone.in(
          collectedFinishedWrites
              .apply(
                  "CommitWriteJob",
                  ParDo.of(new CommitJobFn<Integer>(configView, externalSynchronization))
                      .withSideInputs(configView))
              .getPipeline());
    }

    private void verifyInputWindowing(PCollection<KV<KeyT, ValueT>> input) {
      if (input.isBounded().equals(PCollection.IsBounded.UNBOUNDED)) {
        checkArgument(
            !input.getWindowingStrategy().equals(WindowingStrategy.globalDefault()),
            "Cannot work with %s and GLOBAL %s",
            PCollection.IsBounded.UNBOUNDED,
            WindowingStrategy.class.getSimpleName());
        checkArgument(
            input.getWindowingStrategy().getTrigger().getClass().equals(DefaultTrigger.class),
            "Cannot work with %s trigger. Write works correctly only with %s",
            input.getWindowingStrategy().getTrigger().getClass().getSimpleName(),
            DefaultTrigger.class.getSimpleName());
        checkArgument(
            input.getWindowingStrategy().getAllowedLateness().equals(Duration.ZERO),
            "Write does not allow late data.");
      }
    }

    /**
     * Creates {@link PCollectionView} with one {@link Configuration} based on the set source of the
     * configuration.
     *
     * @param input input data
     * @return PCollectionView with single {@link Configuration}
     * @see Builder#withConfiguration(Configuration)
     * @see Builder#withConfigurationTransform(PTransform)
     */
    private PCollectionView<Configuration> createConfigurationView(
        PCollection<KV<KeyT, ValueT>> input) {

      PCollectionView<Configuration> config;
      if (configuration != null) {
        config =
            input
                .getPipeline()
                .apply("CreateOutputConfig", Create.<Configuration>of(configuration))
                .apply(View.<Configuration>asSingleton().withDefaultValue(configuration));
      } else {
        config = input.apply("TransformDataIntoConfig", configTransform);
      }

      return config;
    }
  }

  /**
   * Represents context of one hadoop write task.
   *
   * @param <KeyT> Key type to write
   * @param <ValueT> Value type to write
   */
  private static class TaskContext<KeyT, ValueT> {

    private final RecordWriter<KeyT, ValueT> recordWriter;
    private final OutputCommitter outputCommitter;
    private final TaskAttemptContext taskAttemptContext;

    TaskContext(TaskAttemptID taskAttempt, Configuration conf) {
      taskAttemptContext = HadoopFormats.createTaskAttemptContext(conf, taskAttempt);
      OutputFormat<KeyT, ValueT> outputFormatObj = HadoopFormats.createOutputFormatFromConfig(conf);
      outputCommitter = initOutputCommitter(outputFormatObj, conf, taskAttemptContext);
      recordWriter = initRecordWriter(outputFormatObj, taskAttemptContext);
    }

    RecordWriter<KeyT, ValueT> getRecordWriter() {
      return recordWriter;
    }

    OutputCommitter getOutputCommitter() {
      return outputCommitter;
    }

    TaskAttemptContext getTaskAttemptContext() {
      return taskAttemptContext;
    }

    int getTaskId() {
      return taskAttemptContext.getTaskAttemptID().getTaskID().getId();
    }

    String getJobId() {
      return taskAttemptContext.getJobID().getJtIdentifier();
    }

    void abortTask() {
      try {
        outputCommitter.abortTask(taskAttemptContext);
      } catch (IOException e) {
        throw new IllegalStateException(
            String.format("Unable to abort task %s of job %s", getTaskId(), getJobId()));
      }
    }

    private RecordWriter<KeyT, ValueT> initRecordWriter(
        OutputFormat<KeyT, ValueT> outputFormatObj, TaskAttemptContext taskAttemptContext)
        throws IllegalStateException {
      try {
        LOG.info(
            "Creating new RecordWriter for task {} of Job with id {}.",
            taskAttemptContext.getTaskAttemptID().getTaskID().getId(),
            taskAttemptContext.getJobID().getJtIdentifier());
        return outputFormatObj.getRecordWriter(taskAttemptContext);
      } catch (InterruptedException | IOException e) {
        throw new IllegalStateException("Unable to create RecordWriter object: ", e);
      }
    }

    private static OutputCommitter initOutputCommitter(
        OutputFormat<?, ?> outputFormatObj,
        Configuration conf,
        TaskAttemptContext taskAttemptContext)
        throws IllegalStateException {
      OutputCommitter outputCommitter;
      try {
        outputCommitter = outputFormatObj.getOutputCommitter(taskAttemptContext);
        if (outputCommitter != null) {
          outputCommitter.setupJob(new JobContextImpl(conf, taskAttemptContext.getJobID()));
        }
      } catch (Exception e) {
        throw new IllegalStateException("Unable to create OutputCommitter object: ", e);
      }

      return outputCommitter;
    }

    @Override
    public String toString() {
      return "TaskContext{"
          + "jobId="
          + getJobId()
          + ", taskId="
          + getTaskId()
          + ", attemptId="
          + taskAttemptContext.getTaskAttemptID().getId()
          + '}';
    }
  }

  /** Coder of configuration instances. */
  private static class ConfigurationCoder extends AtomicCoder<Configuration> {

    @Override
    public void encode(Configuration value, OutputStream outStream) throws IOException {
      DataOutputStream dataOutputStream = new DataOutputStream(outStream);
      value.write(dataOutputStream);
      dataOutputStream.flush();
    }

    @Override
    public Configuration decode(InputStream inStream) throws IOException {
      DataInputStream dataInputStream = new DataInputStream(inStream);
      Configuration config = new Configuration(false);
      config.readFields(dataInputStream);

      return config;
    }
  }

  /**
   * DoFn with following responsibilities:
   *
   * <ul>
   *   <li>Validates configuration - checks whether all required properties are set.
   *   <li>Validates types of input PCollection elements.
   *   <li>Setups start of the {@link OutputFormat} job for given window.
   * </ul>
   *
   * <p>Logic of the setup job is called only for the very first element of the instance. All other
   * elements are directly sent to next processing.
   */
  private static class SetupJobFn<KeyT, ValueT> extends DoFn<KV<KeyT, ValueT>, KV<KeyT, ValueT>> {

    private final ExternalSynchronization externalSynchronization;
    private final PCollectionView<Configuration> configView;
    private final TypeDescriptor<KV<KeyT, ValueT>> inputTypeDescriptor;
    private boolean isSetupJobAttempted;

    SetupJobFn(
        ExternalSynchronization externalSynchronization,
        PCollectionView<Configuration> configView,
        TypeDescriptor<KV<KeyT, ValueT>> inputTypeDescriptor) {
      this.externalSynchronization = externalSynchronization;
      this.configView = configView;
      this.inputTypeDescriptor = inputTypeDescriptor;
    }

    @Setup
    public void setup() {
      isSetupJobAttempted = false;
    }

    @DoFn.ProcessElement
    public void processElement(
        @DoFn.Element KV<KeyT, ValueT> element,
        OutputReceiver<KV<KeyT, ValueT>> receiver,
        BoundedWindow window,
        ProcessContext c) {

      receiver.output(element);

      if (isSetupJobAttempted) {
        // setup of job was already attempted
        return;
      }

      Configuration conf = c.sideInput(configView);

      // validate configuration and input
      // must be done first, because in all later operations are required assumptions from
      // validation
      validateConfiguration(conf);
      validateInputData(conf);

      boolean isJobLockAcquired = externalSynchronization.tryAcquireJobLock(conf);
      isSetupJobAttempted = true;

      if (!isJobLockAcquired) {
        // some parallel execution acquired task
        return;
      }

      try {
        // setup job
        JobID jobId = HadoopFormats.getJobId(conf);

        trySetupJob(jobId, conf, window);

      } catch (Exception e) {
        throw new IllegalStateException(e);
      }
    }

    /**
     * Validates that the mandatory configuration properties such as OutputFormat class,
     * OutputFormat key and value classes are provided in the Hadoop configuration.
     */
    private void validateConfiguration(Configuration conf) {

      checkArgument(conf != null, "Configuration can not be null");
      checkArgument(
          conf.get(OUTPUT_FORMAT_CLASS_ATTR) != null,
          "Configuration must contain \"" + OUTPUT_FORMAT_CLASS_ATTR + "\"");
      checkArgument(
          conf.get(OUTPUT_KEY_CLASS) != null,
          "Configuration must contain \"" + OUTPUT_KEY_CLASS + "\"");
      checkArgument(
          conf.get(OUTPUT_VALUE_CLASS) != null,
          "Configuration must contain \"" + OUTPUT_VALUE_CLASS + "\"");
      checkArgument(conf.get(JOB_ID) != null, "Configuration must contain \"" + JOB_ID + "\"");
    }

    /**
     * Validates input data whether have correctly specified {@link TypeDescriptor}s of input data
     * and if the {@link TypeDescriptor}s match with output types set in the hadoop {@link
     * Configuration}.
     *
     * @param conf hadoop config
     */
    @SuppressWarnings("unchecked")
    private void validateInputData(Configuration conf) {
      TypeDescriptor<KeyT> outputFormatKeyClass =
          (TypeDescriptor<KeyT>) TypeDescriptor.of(conf.getClass(OUTPUT_KEY_CLASS, null));
      TypeDescriptor<ValueT> outputFormatValueClass =
          (TypeDescriptor<ValueT>) TypeDescriptor.of(conf.getClass(OUTPUT_VALUE_CLASS, null));

      checkArgument(
          inputTypeDescriptor != null,
          "Input %s must be set!",
          TypeDescriptor.class.getSimpleName());
      checkArgument(
          KV.class.equals(inputTypeDescriptor.getRawType()),
          "%s expects %s as input type.",
          Write.class.getSimpleName(),
          KV.class);
      checkArgument(
          inputTypeDescriptor.equals(
              TypeDescriptors.kvs(outputFormatKeyClass, outputFormatValueClass)),
          "%s expects following %ss: KV(Key: %s, Value: %s) but following %ss are set: KV(Key: %s, Value: %s)",
          Write.class.getSimpleName(),
          TypeDescriptor.class.getSimpleName(),
          outputFormatKeyClass.getRawType(),
          outputFormatValueClass.getRawType(),
          TypeDescriptor.class.getSimpleName(),
          inputTypeDescriptor.resolveType(KV.class.getTypeParameters()[0]),
          inputTypeDescriptor.resolveType(KV.class.getTypeParameters()[1]));
    }

    /**
     * Setups the hadoop write job as running. There is possibility that some parallel worker
     * already did setup. in this case {@link FileAlreadyExistsException} is catch.
     *
     * @param jobId jobId
     * @param conf hadoop configuration
     * @param window window
     */
    private void trySetupJob(JobID jobId, Configuration conf, BoundedWindow window) {
      try {
        TaskAttemptContext setupTaskContext = HadoopFormats.createSetupTaskContext(conf, jobId);
        OutputFormat<?, ?> jobOutputFormat = HadoopFormats.createOutputFormatFromConfig(conf);

        jobOutputFormat.checkOutputSpecs(setupTaskContext);
        jobOutputFormat.getOutputCommitter(setupTaskContext).setupJob(setupTaskContext);

        LOG.info(
            "Job with id {} successfully configured from window with max timestamp {}.",
            jobId.getJtIdentifier(),
            window.maxTimestamp());

      } catch (FileAlreadyExistsException e) {
        LOG.info("Job was already set by other worker. Skipping rest of the setup.");
      } catch (Exception e) {
        throw new RuntimeException("Unable to setup job.", e);
      }
    }
  }

  /**
   * Commits whole write job.
   *
   * @param <T> type of TaskId identifier
   */
  private static class CommitJobFn<T> extends DoFn<Iterable<T>, Void> {

    private final PCollectionView<Configuration> configView;
    private final ExternalSynchronization externalSynchronization;

    CommitJobFn(
        PCollectionView<Configuration> configView,
        ExternalSynchronization externalSynchronization) {
      this.configView = configView;
      this.externalSynchronization = externalSynchronization;
    }

    @ProcessElement
    public void processElement(ProcessContext c) {

      Configuration config = c.sideInput(configView);
      cleanupJob(config);
    }

    /**
     * Commits whole write job.
     *
     * @param config hadoop config
     */
    private void cleanupJob(Configuration config) {

      externalSynchronization.releaseJobIdLock(config);

      JobID jobID = HadoopFormats.getJobId(config);
      TaskAttemptContext cleanupTaskContext = HadoopFormats.createCleanupTaskContext(config, jobID);
      OutputFormat<?, ?> outputFormat = HadoopFormats.createOutputFormatFromConfig(config);
      try {
        OutputCommitter outputCommitter = outputFormat.getOutputCommitter(cleanupTaskContext);
        outputCommitter.commitJob(cleanupTaskContext);
      } catch (Exception e) {
        throw new RuntimeException("Unable to commit job.", e);
      }
    }
  }

  /**
   * Assigns {@link TaskID#getId()} to the given pair of key and value. {@link TaskID} is later used
   * for writing the pair to hadoop file.
   *
   * @param <KeyT> Type of key
   * @param <ValueT> Type of value
   */
  private static class AssignTaskFn<KeyT, ValueT>
      extends DoFn<KV<KeyT, ValueT>, KV<Integer, KV<KeyT, ValueT>>> {

    private final PCollectionView<Configuration> configView;

    // Transient properties because they are used only for one bundle
    /** Cache of created TaskIDs for given bundle. */
    private transient Map<Integer, TaskID> partitionToTaskContext;

    private transient Partitioner<KeyT, ValueT> partitioner;
    private transient Integer reducersCount;
    private transient JobID jobId;

    /**
     * Needs configuration view of given window.
     *
     * @param configView configuration view
     */
    AssignTaskFn(PCollectionView<Configuration> configView) {
      this.configView = configView;
    }

    /** Deletes cached fields used in previous bundle. */
    @StartBundle
    public void startBundle() {
      partitionToTaskContext = new HashMap<>();
      partitioner = null;
      jobId = null;
      reducersCount = null;
    }

    @ProcessElement
    public void processElement(
        @Element KV<KeyT, ValueT> element,
        OutputReceiver<KV<Integer, KV<KeyT, ValueT>>> receiver,
        ProcessContext c) {

      Configuration config = c.sideInput(configView);

      TaskID taskID = createTaskIDForKV(element, config);
      int taskId = taskID.getId();
      receiver.output(KV.of(taskId, element));
    }

    /**
     * Creates or reuses existing {@link TaskID} for given record.
     *
     * <p>The {@link TaskID} creation is based on the calculation hash function of {@code KeyT} of
     * the pair via {@link Partitioner} (stored in configuration)
     *
     * @param kv keyvalue pair which should be written
     * @param config hadoop configuration
     * @return TaskID assigned to given record
     */
    private TaskID createTaskIDForKV(KV<KeyT, ValueT> kv, Configuration config) {
      int taskContextKey =
          getPartitioner(config).getPartition(kv.getKey(), kv.getValue(), getReducersCount(config));

      return partitionToTaskContext.computeIfAbsent(
          taskContextKey, (key) -> HadoopFormats.createTaskID(getJobId(config), key));
    }

    private JobID getJobId(Configuration config) {
      if (jobId == null) {
        jobId = HadoopFormats.getJobId(config);
      }
      return jobId;
    }

    private int getReducersCount(Configuration config) {
      if (reducersCount == null) {
        reducersCount = HadoopFormats.getReducersCount(config);
      }
      return reducersCount;
    }

    private Partitioner<KeyT, ValueT> getPartitioner(Configuration config) {
      if (partitioner == null) {
        partitioner = HadoopFormats.getPartitioner(config);
      }
      return partitioner;
    }
  }

  /**
   * Writes all {@link KV}s pair for given {@link TaskID} (Task Id determines partition of writing).
   *
   * <p>For every {@link TaskID} are executed following steps:
   *
   * <ul>
   *   <li>Creation of {@link TaskContext} on start of bundle
   *   <li>Writing every single {@link KV} pair via {@link RecordWriter}.
   *   <li>Committing of task on bundle finish
   * </ul>
   *
   * @param <KeyT> Type of key
   * @param <ValueT> Type of value
   */
  private static class WriteFn<KeyT, ValueT> extends DoFn<KV<Integer, KV<KeyT, ValueT>>, Integer> {

    private final PCollectionView<Configuration> configView;
    private final ExternalSynchronization externalSynchronization;

    // Transient property because they are used only for one bundle
    /** Key by combination of Window and TaskId because different windows can use same TaskId. */
    private transient Map<KV<BoundedWindow, Integer>, TaskContext<KeyT, ValueT>>
        bundleTaskContextMap;

    WriteFn(
        PCollectionView<Configuration> configView,
        ExternalSynchronization externalSynchronization) {
      this.configView = configView;
      this.externalSynchronization = externalSynchronization;
    }

    /** Deletes cached map from previous bundle. */
    @StartBundle
    public void startBundle() {
      bundleTaskContextMap = new HashMap<>();
    }

    @ProcessElement
    public void processElement(
        @Element KV<Integer, KV<KeyT, ValueT>> element, ProcessContext c, BoundedWindow b) {

      Integer taskID = element.getKey();
      KV<BoundedWindow, Integer> win2TaskId = KV.of(b, taskID);

      TaskContext<KeyT, ValueT> taskContext =
          bundleTaskContextMap.computeIfAbsent(win2TaskId, w2t -> setupTask(w2t.getValue(), c));

      write(element.getValue(), taskContext);
    }

    @FinishBundle
    public void finishBundle(FinishBundleContext c) {
      if (bundleTaskContextMap == null) {
        return;
      }

      for (Map.Entry<KV<BoundedWindow, Integer>, TaskContext<KeyT, ValueT>> entry :
          bundleTaskContextMap.entrySet()) {
        TaskContext<KeyT, ValueT> taskContext = entry.getValue();

        try {
          taskContext.getRecordWriter().close(taskContext.getTaskAttemptContext());
          taskContext.getOutputCommitter().commitTask(taskContext.getTaskAttemptContext());

          LOG.info("Write task for {} was successfully committed!", taskContext);
        } catch (Exception e) {
          processTaskException(taskContext, e);
        }

        BoundedWindow window = entry.getKey().getKey();
        c.output(taskContext.getTaskId(), Objects.requireNonNull(window).maxTimestamp(), window);
      }
    }

    private void processTaskException(TaskContext<KeyT, ValueT> taskContext, Exception e) {
      LOG.warn("Write task for {} failed. Will abort task.", taskContext);
      taskContext.abortTask();
      throw new IllegalArgumentException(e);
    }

    /**
     * Writes one {@link KV} pair for given {@link TaskID}.
     *
     * @param kv Iterable of pairs to write
     * @param taskContext taskContext
     */
    private void write(KV<KeyT, ValueT> kv, TaskContext<KeyT, ValueT> taskContext) {

      try {
        RecordWriter<KeyT, ValueT> recordWriter = taskContext.getRecordWriter();
        recordWriter.write(kv.getKey(), kv.getValue());
      } catch (Exception e) {
        processTaskException(taskContext, e);
      }
    }

    /**
     * Creates {@link TaskContext} and setups write for given {@code taskId}.
     *
     * @param taskId id of the write Task
     * @param c process context
     * @return created TaskContext
     * @throws IllegalStateException if the setup of the write task failed
     */
    private TaskContext<KeyT, ValueT> setupTask(Integer taskId, ProcessContext c)
        throws IllegalStateException {

      final Configuration conf = c.sideInput(configView);
      TaskAttemptID taskAttemptID = externalSynchronization.acquireTaskAttemptIdLock(conf, taskId);

      TaskContext<KeyT, ValueT> taskContext = new TaskContext<>(taskAttemptID, conf);

      try {
        taskContext.getOutputCommitter().setupTask(taskContext.getTaskAttemptContext());
      } catch (Exception e) {
        processTaskException(taskContext, e);
      }

      LOG.info(
          "Task with id {} of job {} was successfully setup!",
          taskId,
          HadoopFormats.getJobId(conf).getJtIdentifier());

      return taskContext;
    }
  }

  /**
   * Registers task ID for each bundle without need to group data by taskId. Every bundle reserves
   * its own taskId via particular implementation of {@link ExternalSynchronization} class.
   *
   * @param <KeyT> Type of keys to be written.
   * @param <ValueT> Type of values to be written.
   */
  private static class PrepareNonPartitionedTasksFn<KeyT, ValueT>
      extends DoFn<KV<KeyT, ValueT>, KV<Integer, KV<KeyT, ValueT>>> {

    private transient TaskID taskId;

    private final PCollectionView<Configuration> configView;
    private final ExternalSynchronization externalSynchronization;

    private PrepareNonPartitionedTasksFn(
        PCollectionView<Configuration> configView,
        ExternalSynchronization externalSynchronization) {
      this.configView = configView;
      this.externalSynchronization = externalSynchronization;
    }

    @DoFn.StartBundle
    public void startBundle() {
      taskId = null;
    }

    @DoFn.ProcessElement
    public void processElement(
        @DoFn.Element KV<KeyT, ValueT> element,
        OutputReceiver<KV<Integer, KV<KeyT, ValueT>>> output,
        ProcessContext c) {

      if (taskId == null) {
        Configuration conf = c.sideInput(configView);
        taskId = externalSynchronization.acquireTaskIdLock(conf);
      }

      output.output(KV.of(taskId.getId(), element));
    }
  }
}