/*
 * 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.spark

import java.io.{IOException, ObjectInputStream, ObjectOutputStream}

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
import scala.reflect.ClassTag
import scala.util.hashing.byteswap32

import org.apache.spark.rdd.{PartitionPruningRDD, RDD}
import org.apache.spark.serializer.JavaSerializer
import org.apache.spark.util.{CollectionsUtils, Utils}

/**
 * support data skew scenario
 * copy from spark: RangePartitioner
 *
 * RangePartitioner:
 * the rangeBounds are the distinct values, each rangeBound has a partition.
 * so for the data skew scenario, some partitions will include more data.
 *
 * DataSkewRangePartitioner:
 * the rangeBounds are also the distinct values, but it calculates the skew weight.
 * So some rangeBounds maybe have more than one partitions.
 *
 * for example, split following CSV file to 5 partitions by col2:
 * ---------------
 * col1,col2
 * 1,
 * 2,
 * 3,
 * 4,
 * 5,e
 * 6,f
 * 7,g
 * 8,h
 * 9,i
 * 10,j
 * ---------------
 * RangePartitioner will give the following rangeBounds.
 * -------------------------------------------------
 * | range bound| partition range| number of values|
 * -------------------------------------------------
 * | null       | <= null        | 4               |
 * | e          | (null, e]      | 1               |
 * | f          | (e, f]         | 1               |
 * | h          | (f, h]         | 2               |
 * |            | > h            | 2               |
 * -------------------------------------------------
 *
 * DataSkewRangePartitioner will give the following rangeBounds.
 * --------------------------------------------------------------
 * | range bound| skew weight| partition range| number of values|
 * --------------------------------------------------------------
 * | null       | 2          | <= null        | 2(4/2)          |
 * | f          |            | (null, f]      | 2               |
 * | h          |            | (f, h]         | 2               |
 * |            |            | > h            | 2               |
 * |            |            | <= null        | 2(4/2)          |
 * --------------------------------------------------------------
 * The skew weight of range bound "null" is 2.
 * So it will start two tasks for range bound "null" to create two partitions.
 * For a range bound, the number of final partitions is the same as the skew weight.
 */
class DataSkewRangePartitioner[K: Ordering : ClassTag, V](
    partitions: Int,
    rdd: RDD[_ <: Product2[K, V]],
    withoutSkew: Boolean)
  extends Partitioner {

  // We allow partitions = 0, which happens when sorting an empty RDD under the default settings.
  require(partitions >= 0, s"Number of partitions cannot be negative but found $partitions.")

  private var ordering = implicitly[Ordering[K]]

  // An array of upper bounds for the first (partitions - 1) partitions
  // dataSkewCount: how many bounds happened data skew
  // dataSkewIndex: the index of data skew bounds
  // dataSkewNum: how many partition of each data skew bound
  // Min and Max values of complete range
  var (rangeBounds: Array[K], skewCount: Int, skewIndexes: Array[Int],
  skewWeights: Array[Int]) = {
    if (partitions <= 1) {
      (Array.empty[K], 0, Array.empty[Int], Array.empty[Int])
    } else {
      // This is the sample size we need to have roughly balanced output partitions, capped at 1M.
      val sampleSize = math.min(20.0 * partitions, 1e6)
      // Assume the input partitions are roughly balanced and over-sample a little bit.
      val sampleSizePerPartition = math.ceil(3.0 * sampleSize / rdd.partitions.length).toInt
      val (numItems, sketched) = RangePartitioner.sketch(rdd.map(_._1), sampleSizePerPartition)
      if (numItems == 0L) {
        (Array.empty[K], 0, Array.empty[Int], Array.empty[Int], Array.empty[K])
      } else {
        // If a partition contains much more than the average number of items, we re-sample from it
        // to ensure that enough items are collected from that partition.
        val fraction = math.min(sampleSize / math.max(numItems, 1L), 1.0)
        val candidates = ArrayBuffer.empty[(K, Float)]
        val imbalancedPartitions = mutable.Set.empty[Int]
        sketched.foreach { case (idx, n, sample) =>
          if (fraction * n > sampleSizePerPartition) {
            imbalancedPartitions += idx
          } else {
            // The weight is 1 over the sampling probability.
            val weight = (n.toDouble / sample.length).toFloat
            for (key <- sample) {
              candidates += ((key, weight))
            }
          }
        }
        if (imbalancedPartitions.nonEmpty) {
          // Re-sample imbalanced partitions with the desired sampling probability.
          val imbalanced = new PartitionPruningRDD(rdd.map(_._1), imbalancedPartitions.contains)
          val seed = byteswap32(-rdd.id - 1)
          val reSampled = imbalanced.sample(withReplacement = false, fraction, seed).collect()
          val weight = (1.0 / fraction).toFloat
          candidates ++= reSampled.map(x => (x, weight))
        }
        // In case of compaction we do not need the skew handled ranges so we use RangePartitioner,
        // but we require the overall minmax for creating the separate ranges.
        // withoutSkew = true for Compaction only
        if (withoutSkew == false) {
          determineBounds(candidates, partitions, false)
        } else {
          var ranges = RangePartitioner.determineBounds(candidates, partitions)
          var otherRangeParams = determineBounds(candidates, partitions, true)
          (ranges, otherRangeParams._2, otherRangeParams._3,
            otherRangeParams._4)
        }
      }
    }
  }

  def determineBounds(
      candidates: ArrayBuffer[(K, Float)],
      partitions: Int,
      withoutSkew: Boolean): (Array[K], Int, Array[Int], Array[Int]) = {
    val ordered = candidates.sortBy(_._1)
    val numCandidates = ordered.size
    val sumWeights = ordered.map(_._2.toDouble).sum
    val step = sumWeights / partitions
    var cumWeight = 0.0
    var target = step
    val bounds = ArrayBuffer.empty[K]
    var i = 0
    var j = 0
    var previousBound = Option.empty[K]
    while ((i < numCandidates) && (j < partitions - 1)) {
      val (key, weight) = ordered(i)
      cumWeight += weight
      if (cumWeight >= target) {
        // Skip duplicate values.
        if (previousBound.isEmpty || ordering.gteq(key, previousBound.get)) {
          bounds += key
          target += step
          j += 1
          previousBound = Some(key)
        }
      }
      i += 1
    }

    if (bounds.size >= 2) {
      combineDataSkew(bounds)
    } else {
      (bounds.toArray, 0, Array.empty[Int], Array.empty[Int])
    }
  }

  def combineDataSkew(bounds: ArrayBuffer[K]): (Array[K], Int, Array[Int], Array[Int]) = {
    val finalBounds = ArrayBuffer.empty[K]
    var preBound = bounds(0)
    finalBounds += preBound
    val dataSkewIndexTmp = ArrayBuffer.empty[Int]
    val dataSkewNumTmp = ArrayBuffer.empty[Int]
    var dataSkewCountTmp = 1
    (1 until bounds.size).foreach { index =>
      val bound = bounds(index)
      if (ordering.equiv(bound, preBound)) {
        if (dataSkewCountTmp == 1) {
          dataSkewIndexTmp += (finalBounds.size - 1)
        }
        dataSkewCountTmp += 1
      } else {
        finalBounds += bound
        preBound = bound
        if (dataSkewCountTmp > 1) {
          dataSkewNumTmp += dataSkewCountTmp
          dataSkewCountTmp = 1
        }
      }
    }
    if (dataSkewCountTmp > 1) {
      dataSkewNumTmp += dataSkewCountTmp
    }
    if (dataSkewIndexTmp.size > 0) {
      (finalBounds.toArray, dataSkewIndexTmp.size, dataSkewIndexTmp.toArray, dataSkewNumTmp
        .toArray)
    } else {
      (finalBounds.toArray, 0, Array.empty[Int], Array.empty[Int])
    }
  }

  private var skewPartitions: Int = if (skewCount == 0) {
    0
  } else {
    skewWeights.map(_ - 1).sum
  }

  def numPartitions: Int = rangeBounds.length + 1 + skewPartitions

  private var binarySearch: ((Array[K], K) => Int) = CollectionsUtils.makeBinarySearch[K]
  private var skewIndexesLen = 0
  private var dsps: Array[DataSkewPartitioner] = null

  def initialize(): Unit = {
    if (skewCount > 0) {
      skewIndexesLen = skewIndexes.length
      dsps = new Array[DataSkewPartitioner](skewIndexesLen)
      var previousPart = rangeBounds.length
      for (i <- 0 until skewIndexesLen) {
        dsps(i) = new DataSkewPartitioner(skewIndexes(i), previousPart, skewWeights(i))
        previousPart = previousPart + skewWeights(i) - 1
      }
    }
  }

  private var needInit = true

  def getPartition(key: Any): Int = {
    if (needInit) {
      needInit = false
      initialize()
    }
    val k = key.asInstanceOf[K]
    var partition = 0
    if (rangeBounds.length <= 128) {
      // If we have less than 128 partitions naive search
      while (partition < rangeBounds.length && ordering.gt(k, rangeBounds(partition))) {
        partition += 1
      }
    } else {
      // Determine which binary search method to use only once.
      partition = binarySearch(rangeBounds, k)
      // binarySearch either returns the match location or -[insertion point]-1
      if (partition < 0) {
        partition = -partition - 1
      }
      if (partition > rangeBounds.length) {
        partition = rangeBounds.length
      }
    }
    if (skewCount == 0) {
      partition
    } else {
      getDataSkewPartition(partition)
    }
  }

  def getDataSkewPartition(partition: Int): Int = {
    var index = -1
    if (partition <= skewIndexes(skewIndexesLen - 1) && partition >= skewIndexes(0)) {
      for (i <- 0 until skewIndexesLen) {
        if (skewIndexes(i) == partition) {
          index = i
        }
      }
    }
    if (index == -1) {
      partition
    } else {
      nextPartition(index)
    }
  }

  def nextPartition(index: Int): Int = {
    dsps(index).nextPartition()
  }

  override def equals(other: Any): Boolean = {
    other match {
      case r: DataSkewRangePartitioner[_, _] =>
        r.rangeBounds.sameElements(rangeBounds)
      case _ =>
        false
    }
  }

  override def hashCode(): Int = {
    val prime = 31
    var result = 1
    var i = 0
    while (i < rangeBounds.length) {
      result = prime * result + rangeBounds(i).hashCode
      i += 1
    }
    result = prime * result
    result
  }

  @throws(classOf[IOException])
  private def writeObject(out: ObjectOutputStream): Unit = {
    Utils.tryOrIOException {
      val sfactory = SparkEnv.get.serializer
      sfactory match {
        case js: JavaSerializer => out.defaultWriteObject()
        case _ =>
          out.writeInt(skewCount)
          out.writeInt(skewPartitions)
          if (skewCount > 0) {
            out.writeObject(skewIndexes)
            out.writeObject(skewWeights)
          }
          out.writeObject(ordering)
          out.writeObject(binarySearch)

          val ser = sfactory.newInstance()
          Utils.serializeViaNestedStream(out, ser) { stream =>
            stream.writeObject(scala.reflect.classTag[Array[K]])
            stream.writeObject(rangeBounds)
          }
      }
    }
  }

  @throws(classOf[IOException])
  private def readObject(in: ObjectInputStream): Unit = {
    Utils.tryOrIOException {
      needInit = true
      val sfactory = SparkEnv.get.serializer
      sfactory match {
        case js: JavaSerializer => in.defaultReadObject()
        case _ =>
          skewCount = in.readInt()
          skewPartitions = in.readInt()
          if (skewCount > 0) {
            skewIndexes = in.readObject().asInstanceOf[Array[Int]]
            skewWeights = in.readObject().asInstanceOf[Array[Int]]
          }
          ordering = in.readObject().asInstanceOf[Ordering[K]]
          binarySearch = in.readObject().asInstanceOf[(Array[K], K) => Int]

          val ser = sfactory.newInstance()
          Utils.deserializeViaNestedStream(in, ser) { ds =>
            implicit val classTag = ds.readObject[ClassTag[Array[K]]]()
            rangeBounds = ds.readObject[Array[K]]()
          }
      }
    }
  }
}

class DataSkewPartitioner(originPart: Int, previousPart: Int, skewWeight: Int) {
  var index = 0

  def nextPartition(): Int = {
    if (index == 0) {
      index = index + 1
      originPart
    } else {
      val newPartition = previousPart + index
      index = index + 1
      if (index == skewWeight) {
        index = 0
      }
      newPartition
    }
  }
}