Spark Shuffle机制

/** Get the value for a given key */
def apply(key: K): V = {
  assert(!destroyed, destructionMessage)
  val k = key.asInstanceOf[AnyRef]
  if (k.eq(null)) {
    return nullValue
  }
  var pos = rehash(k.hashCode) & mask
  var i = 1
  while (true) {
    val curKey = data(2 * pos)
    if (k.eq(curKey) || k.equals(curKey)) {
      return data(2 * pos + 1).asInstanceOf[V]
    } else if (curKey.eq(null)) {
      return null.asInstanceOf[V]
    } else {
      val delta = i
      pos = (pos + delta) & mask
      i += 1
    }
  }
  null.asInstanceOf[V]
}

/** Set the value for a key */
def update(key: K, value: V): Unit = {
  assert(!destroyed, destructionMessage)
  val k = key.asInstanceOf[AnyRef]
  if (k.eq(null)) {
    if (!haveNullValue) {
      incrementSize()
    }
    nullValue = value
    haveNullValue = true
    return
  }
  var pos = rehash(key.hashCode) & mask
  var i = 1
  while (true) {
    val curKey = data(2 * pos)
    if (curKey.eq(null)) {
      data(2 * pos) = k
      data(2 * pos + 1) = value.asInstanceOf[AnyRef]
      incrementSize()  // Since we added a new key
      return
    } else if (k.eq(curKey) || k.equals(curKey)) {
      data(2 * pos + 1) = value.asInstanceOf[AnyRef]
      return
    } else {
      val delta = i
      pos = (pos + delta) & mask
      i += 1
    }
  }
}

/**
 * Return an iterator of the map in sorted order. This provides a way to sort the map without
 * using additional memory, at the expense of destroying the validity of the map.
 */
def destructiveSortedIterator(keyComparator: Comparator[K]): Iterator[(K, V)] = {
  destroyed = true
  // Pack KV pairs into the front of the underlying array
  var keyIndex, newIndex = 0
  while (keyIndex < capacity) {
    if (data(2 * keyIndex) != null) {
      data(2 * newIndex) = data(2 * keyIndex)
      data(2 * newIndex + 1) = data(2 * keyIndex + 1)
      newIndex += 1
    }
    keyIndex += 1
  }
  assert(curSize == newIndex + (if (haveNullValue) 1 else 0))

  new Sorter(new KVArraySortDataFormat[K, AnyRef]).sort(data, 0, newIndex, keyComparator)

  new Iterator[(K, V)] {
    var i = 0
    var nullValueReady = haveNullValue
    def hasNext: Boolean = (i < newIndex || nullValueReady)
    def next(): (K, V) = {
      if (nullValueReady) {
        nullValueReady = false
        (null.asInstanceOf[K], nullValue)
      } else {
        val item = (data(2 * i).asInstanceOf[K], data(2 * i + 1).asInstanceOf[V])
        i += 1
        item
      }
    }
  }
}

/**
 * A general interface for collections to keep track of their estimated sizes in bytes.
 * We sample with a slow exponential back-off using the SizeEstimator to amortize the time,
 * as each call to SizeEstimator is somewhat expensive (order of a few milliseconds).
 */
private[spark] trait SizeTracker {

  import SizeTracker._

  /**
   * Controls the base of the exponential which governs the rate of sampling.
   * E.g., a value of 2 would mean we sample at 1, 2, 4, 8, ... elements.
   */
  private val SAMPLE_GROWTH_RATE = 1.1

  /** Samples taken since last resetSamples(). Only the last two are kept for extrapolation. */
  private val samples = new mutable.Queue[Sample]

  /** The average number of bytes per update between our last two samples. */
  private var bytesPerUpdate: Double = _

  /** Total number of insertions and updates into the map since the last resetSamples(). */
  private var numUpdates: Long = _

  /** The value of 'numUpdates' at which we will take our next sample. */
  private var nextSampleNum: Long = _

  resetSamples()

  /**
   * Reset samples collected so far.
   * This should be called after the collection undergoes a dramatic change in size.
   */
  protected def resetSamples(): Unit = {
    numUpdates = 1
    nextSampleNum = 1
    samples.clear()
    takeSample()
  }

  /**
   * Callback to be invoked after every update.
   */
  protected def afterUpdate(): Unit = {
    numUpdates += 1
    if (nextSampleNum == numUpdates) {
      takeSample()
    }
  }

  /**
   * Take a new sample of the current collection's size.
   */
  private def takeSample(): Unit = {
    samples.enqueue(Sample(SizeEstimator.estimate(this), numUpdates))
    // Only use the last two samples to extrapolate
    if (samples.size > 2) {
      samples.dequeue()
    }
    val bytesDelta = samples.toList.reverse match {
      case latest :: previous :: tail =>
        (latest.size - previous.size).toDouble / (latest.numUpdates - previous.numUpdates)
      // If fewer than 2 samples, assume no change
      case _ => 0
    }
    bytesPerUpdate = math.max(0, bytesDelta)
    nextSampleNum = math.ceil(numUpdates * SAMPLE_GROWTH_RATE).toLong
  }

  /**
   * Estimate the current size of the collection in bytes. O(1) time.
   */
  def estimateSize(): Long = {
    assert(samples.nonEmpty)
    val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
    (samples.last.size + extrapolatedDelta).toLong
  }
}

private object SizeTracker {
  case class Sample(size: Long, numUpdates: Long)
}

private[spark] class SizeTrackingAppendOnlyMap[K, V]
  extends AppendOnlyMap[K, V] with SizeTracker
{
  override def update(key: K, value: V): Unit = {
    super.update(key, value)
    super.afterUpdate()
  }

  override def changeValue(key: K, updateFunc: (Boolean, V) => V): V = {
    val newValue = super.changeValue(key, updateFunc)
    super.afterUpdate()
    newValue
  }

  override protected def growTable(): Unit = {
    super.growTable()
    resetSamples()
  }
}

class ExternalAppendOnlyMap[K, V, C](
    createCombiner: V => C,
    mergeValue: (C, V) => C,
    mergeCombiners: (C, C) => C,
    serializer: Serializer = SparkEnv.get.serializer,
    blockManager: BlockManager = SparkEnv.get.blockManager,
    context: TaskContext = TaskContext.get(),
    serializerManager: SerializerManager = SparkEnv.get.serializerManager)
  extends Spillable[SizeTracker](context.taskMemoryManager())
  with Serializable
  with Logging
  with Iterable[(K, C)] {

  @volatile private[collection] var currentMap = new SizeTrackingAppendOnlyMap[K, C]
  private val spilledMaps = new ArrayBuffer[DiskMapIterator]
  private val sparkConf = SparkEnv.get.conf
  private val diskBlockManager = blockManager.diskBlockManager

  /**
   * Size of object batches when reading/writing from serializers.
   *
   * Objects are written in batches, with each batch using its own serialization stream. This
   * cuts down on the size of reference-tracking maps constructed when deserializing a stream.
   *
   * NOTE: Setting this too low can cause excessive copying when serializing, since some serializers
   * grow internal data structures by growing + copying every time the number of objects doubles.
   */
  private val serializerBatchSize = sparkConf.get(config.SHUFFLE_SPILL_BATCH_SIZE)

  // Number of bytes spilled in total
  private var _diskBytesSpilled = 0L
  def diskBytesSpilled: Long = _diskBytesSpilled

  // Use getSizeAsKb (not bytes) to maintain backwards compatibility if no units are provided
  private val fileBufferSize = sparkConf.get(config.SHUFFLE_FILE_BUFFER_SIZE).toInt * 1024

  // Write metrics
  private val writeMetrics: ShuffleWriteMetrics = new ShuffleWriteMetrics()

  // Peak size of the in-memory map observed so far, in bytes
  private var _peakMemoryUsedBytes: Long = 0L
  def peakMemoryUsedBytes: Long = _peakMemoryUsedBytes

  private val keyComparator = new HashComparator[K]
  private val ser = serializer.newInstance()

  @volatile private var readingIterator: SpillableIterator = null

/**
 * Insert the given iterator of keys and values into the map.
 *
 * When the underlying map needs to grow, check if the global pool of shuffle memory has
 * enough room for this to happen. If so, allocate the memory required to grow the map;
 * otherwise, spill the in-memory map to disk.
 *
 * The shuffle memory usage of the first trackMemoryThreshold entries is not tracked.
 */
def insertAll(entries: Iterator[Product2[K, V]]): Unit = {
  if (currentMap == null) {
    throw new IllegalStateException(
      "Cannot insert new elements into a map after calling iterator")
  }
  // An update function for the map that we reuse across entries to avoid allocating
  // a new closure each time
  var curEntry: Product2[K, V] = null
  val update: (Boolean, C) => C = (hadVal, oldVal) => {
    if (hadVal) mergeValue(oldVal, curEntry._2) else createCombiner(curEntry._2)
  }
	// 遍历每个K, V对
  while (entries.hasNext) {
    curEntry = entries.next()
    // 插入前，先估算当前currentMap的大小
    val estimatedSize = currentMap.estimateSize()
    // 更新巅峰占用内存
    if (estimatedSize > _peakMemoryUsedBytes) {
      _peakMemoryUsedBytes = estimatedSize
    }
    // 判断是否spill，如果发生spill，创建新的currentMap
    if (maybeSpill(currentMap, estimatedSize)) {
      currentMap = new SizeTrackingAppendOnlyMap[K, C]
    }
    // 插入新的记录
    currentMap.changeValue(curEntry._1, update)
    // 更新_elementsRead计数
    addElementsRead()
  }
}

// Initial threshold for the size of a collection before we start tracking its memory usage
// For testing only
private[this] val initialMemoryThreshold: Long =
  SparkEnv.get.conf.get(SHUFFLE_SPILL_INITIAL_MEM_THRESHOLD)

// Threshold for this collection's size in bytes before we start tracking its memory usage
// To avoid a large number of small spills, initialize this to a value orders of magnitude > 0
@volatile private[this] var myMemoryThreshold = initialMemoryThreshold
/**
 * Spills the current in-memory collection to disk if needed. Attempts to acquire more
 * memory before spilling.
 *
 * @param collection collection to spill to disk
 * @param currentMemory estimated size of the collection in bytes
 * @return true if `collection` was spilled to disk; false otherwise
 */
protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
  var shouldSpill = false
  // myMemoryThresold的作用类似于内存配额，或者说已经申请到的内存大小
  // 如果当前内存占用大于申请到的内存大小
  if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
    // Claim up to double our current memory from the shuffle memory pool
    // 尝试申请最多两倍的currentMemory
    val amountToRequest = 2 * currentMemory - myMemoryThreshold
    val granted = acquireMemory(amountToRequest)
    // 更新申请到的内存大小
    myMemoryThreshold += granted
    // If we were granted too little memory to grow further (either tryToAcquire returned 0,
    // or we already had more memory than myMemoryThreshold), spill the current collection
    // 如果内存申请不到，并且集合的内存占用大于申请到的内存
    shouldSpill = currentMemory >= myMemoryThreshold
  }
  shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
  // Actually spill
  if (shouldSpill) {
    _spillCount += 1
    logSpillage(currentMemory)
    spill(collection)
    _elementsRead = 0
    _memoryBytesSpilled += currentMemory
    releaseMemory()
  }
  shouldSpill
}

/**
 * Sort the existing contents of the in-memory map and spill them to a temporary file on disk.
 */
override protected[this] def spill(collection: SizeTracker): Unit = {
  val inMemoryIterator = currentMap.destructiveSortedIterator(keyComparator)
  val diskMapIterator = spillMemoryIteratorToDisk(inMemoryIterator)
  spilledMaps += diskMapIterator
}
/**
 * Spill the in-memory Iterator to a temporary file on disk.
 */
private[this] def spillMemoryIteratorToDisk(inMemoryIterator: Iterator[(K, C)])
    : DiskMapIterator = {
  // 利用uuid创建一个临时的block
  val (blockId, file) = diskBlockManager.createTempLocalBlock()
  // 获取文件的writer
  val writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSize, writeMetrics)
  var objectsWritten = 0

  // List of batch sizes (bytes) in the order they are written to disk
  val batchSizes = new ArrayBuffer[Long]

  // Flush the disk writer's contents to disk, and update relevant variables
  def flush(): Unit = {
    val segment = writer.commitAndGet()
    batchSizes += segment.length
    _diskBytesSpilled += segment.length
    objectsWritten = 0
  }

  var success = false
  try {
    // 遍历迭代器，分批将record flush到磁盘
    while (inMemoryIterator.hasNext) {
      val kv = inMemoryIterator.next()
      writer.write(kv._1, kv._2)
      objectsWritten += 1

      if (objectsWritten == serializerBatchSize) {
        flush()
      }
    }
    if (objectsWritten > 0) {
      flush()
      writer.close()
    } else {
      writer.revertPartialWritesAndClose()
    }
    success = true
  } finally {
    if (!success) {
      // This code path only happens if an exception was thrown above before we set success;
      // close our stuff and let the exception be thrown further
      writer.closeAndDelete()
    }
  }

  new DiskMapIterator(file, blockId, batchSizes)
}

/**
 * Returns a destructive iterator for iterating over the entries of this map.
 * If this iterator is forced spill to disk to release memory when there is not enough memory,
 * it returns pairs from an on-disk map.
 */
def destructiveIterator(inMemoryIterator: Iterator[(K, C)]): Iterator[(K, C)] = {
  readingIterator = new SpillableIterator(inMemoryIterator)
  readingIterator.toCompletionIterator
}
/**
 * Return a destructive iterator that merges the in-memory map with the spilled maps.
 * If no spill has occurred, simply return the in-memory map's iterator.
 */
override def iterator: Iterator[(K, C)] = {
  if (currentMap == null) {
    throw new IllegalStateException(
      "ExternalAppendOnlyMap.iterator is destructive and should only be called once.")
  }
  if (spilledMaps.isEmpty) {
    destructiveIterator(currentMap.iterator)
  } else {
    new ExternalIterator()
  }
}

/**
 * Wrapper around an iterator which calls a completion method after it successfully iterates
 * through all the elements.
 */
private[spark]
abstract class CompletionIterator[ +A, +I <: Iterator[A]](sub: I) extends Iterator[A] {

  private[this] var completed = false
  private[this] var iter = sub
  def next(): A = iter.next()
  def hasNext: Boolean = {
    val r = iter.hasNext
    if (!r && !completed) {
      completed = true
      // reassign to release resources of highly resource consuming iterators early
      iter = Iterator.empty.asInstanceOf[I]
      completion()
    }
    r
  }

  def completion(): Unit
}

private[spark] object CompletionIterator {
  def apply[A, I <: Iterator[A]](sub: I, completionFunction: => Unit) : CompletionIterator[A, I] = {
    new CompletionIterator[A, I](sub) {
      def completion(): Unit = completionFunction
    }
  }
}

private class SpillableIterator(var upstream: Iterator[(K, C)])
  extends Iterator[(K, C)] {

  private val SPILL_LOCK = new Object()

  private var cur: (K, C) = readNext()

  private var hasSpilled: Boolean = false

  def spill(): Boolean = SPILL_LOCK.synchronized {
    if (hasSpilled) {
      false
    } else {
      logInfo(log"Task ${MDC(TASK_ATTEMPT_ID, context.taskAttemptId())} force spilling" +
        log" in-memory map to disk and it will release " +
        log"${MDC(NUM_BYTES, org.apache.spark.util.Utils.bytesToString(getUsed()))} memory")
      val nextUpstream = spillMemoryIteratorToDisk(upstream)
      assert(!upstream.hasNext)
      hasSpilled = true
      upstream = nextUpstream
      true
    }
  }

  private def destroy(): Unit = {
    freeCurrentMap()
    upstream = Iterator.empty
  }

  def toCompletionIterator: CompletionIterator[(K, C), SpillableIterator] = {
    CompletionIterator[(K, C), SpillableIterator](this, this.destroy())
  }

  def readNext(): (K, C) = SPILL_LOCK.synchronized {
    if (upstream.hasNext) {
      upstream.next()
    } else {
      null
    }
  }

  override def hasNext: Boolean = cur != null

  override def next(): (K, C) = {
    val r = cur
    cur = readNext()
    r
  }
}

/**
 * An iterator that sort-merges (K, C) pairs from the in-memory map and the spilled maps
 */
private class ExternalIterator extends Iterator[(K, C)] {

  // A queue that maintains a buffer for each stream we are currently merging
  // This queue maintains the invariant that it only contains non-empty buffers
  // 堆，用作堆排序
  private val mergeHeap = new mutable.PriorityQueue[StreamBuffer]

  // Input streams are derived both from the in-memory map and spilled maps on disk
  // The in-memory map is sorted in place, while the spilled maps are already in sorted order
  // input stearam集合，包括内存中map和磁盘上map
  private val sortedMap = destructiveIterator(
    currentMap.destructiveSortedIterator(keyComparator))
  private val inputStreams = (Seq(sortedMap) ++ spilledMaps).map(it => it.buffered)
	// 将每个iterator中最小hash值对应的多条记录保存到kcPairs，并放入堆中进行排序
  // 每个iterator中，由于hash碰撞，可能出现多个Key对应同一个hash值
  // 但由于iterator中保存的都是聚合的结果，所以不可以有同一个key的多条记录
  inputStreams.foreach { it =>
    val kcPairs = new ArrayBuffer[(K, C)]
    readNextHashCode(it, kcPairs)
    if (kcPairs.length > 0) {
      mergeHeap.enqueue(new StreamBuffer(it, kcPairs))
    }
  }

  private def readNextHashCode(it: BufferedIterator[(K, C)], buf: ArrayBuffer[(K, C)]): Unit = {
    if (it.hasNext) {
      var kc = it.next()
      buf += kc
      val minHash = hashKey(kc)
      while (it.hasNext && it.head._1.hashCode() == minHash) {
        kc = it.next()
        buf += kc
      }
    }
  }

  /**
   * If the given buffer contains a value for the given key, merge that value into
   * baseCombiner and remove the corresponding (K, C) pair from the buffer.
   */
  private def mergeIfKeyExists(key: K, baseCombiner: C, buffer: StreamBuffer): C = {
    var i = 0
    // 遍历buffer中的记录，找到相同的key，并且merge，然后删除这条记录
    while (i < buffer.pairs.length) {
      val pair = buffer.pairs(i)
      if (pair._1 == key) {
        // Note that there's at most one pair in the buffer with a given key, since we always
        // merge stuff in a map before spilling, so it's safe to return after the first we find
        removeFromBuffer(buffer.pairs, i)
        return mergeCombiners(baseCombiner, pair._2)
      }
      i += 1
    }
    baseCombiner
  }

	// 这里通过交换元素并移除最后一个元素的方式实现高效的移除任意元素
  private def removeFromBuffer[T](buffer: ArrayBuffer[T], index: Int): T = {
    val elem = buffer(index)
    buffer(index) = buffer(buffer.size - 1)  // This also works if index == buffer.size - 1
    buffer.dropRightInPlace(1)
    elem
  }

  /**
   * Return true if there exists an input stream that still has unvisited pairs.
   */
  override def hasNext: Boolean = mergeHeap.nonEmpty

  /**
   * Select a key with the minimum hash, then combine all values with the same key from all
   * input streams.
   */
  override def next(): (K, C) = {
    if (mergeHeap.isEmpty) {
      throw new NoSuchElementException
    }
    // Select a key from the StreamBuffer that holds the lowest key hash
    // 取出拥有最小hash值的buffer
    val minBuffer = mergeHeap.dequeue()
    val minPairs = minBuffer.pairs
    val minHash = minBuffer.minKeyHash
    val minPair = removeFromBuffer(minPairs, 0)
    val minKey = minPair._1
    var minCombiner = minPair._2
    assert(hashKey(minPair) == minHash)

    // For all other streams that may have this key (i.e. have the same minimum key hash),
    // merge in the corresponding value (if any) from that stream
    // 查询堆顶元素，如果hash值相同，则出队，并尝试合并相同Key的记录
    val mergedBuffers = ArrayBuffer[StreamBuffer](minBuffer)
    while (mergeHeap.nonEmpty && mergeHeap.head.minKeyHash == minHash) {
      val newBuffer = mergeHeap.dequeue()
      minCombiner = mergeIfKeyExists(minKey, minCombiner, newBuffer)
      mergedBuffers += newBuffer
    }

    // Repopulate each visited stream buffer and add it back to the queue if it is non-empty
    // 检查出队的iterator是否为空，如果为空，则丢弃，否则如果buffer不为空，直接入队
    // 如果buffer为空，读取下一个最小的hash值的记录，并入队
    mergedBuffers.foreach { buffer =>
      if (buffer.isEmpty) {
        readNextHashCode(buffer.iterator, buffer.pairs)
      }
      if (!buffer.isEmpty) {
        mergeHeap.enqueue(buffer)
      }
    }

    (minKey, minCombiner)
  }

  private class StreamBuffer(
      val iterator: BufferedIterator[(K, C)],
      val pairs: ArrayBuffer[(K, C)])
    extends Comparable[StreamBuffer] {

    def isEmpty: Boolean = pairs.length == 0

    // Invalid if there are no more pairs in this stream
    def minKeyHash: Int = {
      assert(pairs.length > 0)
      hashKey(pairs.head)
    }

    override def compareTo(other: StreamBuffer): Int = {
      // descending order because mutable.PriorityQueue dequeues the max, not the min
      if (other.minKeyHash < minKeyHash) -1 else if (other.minKeyHash == minKeyHash) 0 else 1
    }
  }
}

/**
 * An iterator that returns (K, C) pairs in sorted order from an on-disk map
 */
private class DiskMapIterator(file: File, blockId: BlockId, batchSizes: ArrayBuffer[Long])
  extends Iterator[(K, C)]
{
  private val batchOffsets = batchSizes.scanLeft(0L)(_ + _)  // Size will be batchSize.length + 1
  assert(file.length() == batchOffsets.last,
    "File length is not equal to the last batch offset:\n" +
    s"    file length = ${file.length}\n" +
    s"    last batch offset = ${batchOffsets.last}\n" +
    s"    all batch offsets = ${batchOffsets.mkString(",")}"
  )

  private var batchIndex = 0  // Which batch we're in
  private var fileStream: FileInputStream = null

  // An intermediate stream that reads from exactly one batch
  // This guards against pre-fetching and other arbitrary behavior of higher level streams
  private var deserializeStream: DeserializationStream = null
  private var batchIterator: Iterator[(K, C)] = null
  private var objectsRead = 0

  /**
   * Construct a stream that reads only from the next batch.
   */
  private def nextBatchIterator(): Iterator[(K, C)] = {
    // Note that batchOffsets.length = numBatches + 1 since we did a scan above; check whether
    // we're still in a valid batch.
    if (batchIndex < batchOffsets.length - 1) {
      if (deserializeStream != null) {
        deserializeStream.close()
        fileStream.close()
        deserializeStream = null
        fileStream = null
      }

      val start = batchOffsets(batchIndex)
      fileStream = new FileInputStream(file)
      fileStream.getChannel.position(start)
      batchIndex += 1

      val end = batchOffsets(batchIndex)

      assert(end >= start, "start = " + start + ", end = " + end +
        ", batchOffsets = " + batchOffsets.mkString("[", ", ", "]"))

      val bufferedStream = new BufferedInputStream(ByteStreams.limit(fileStream, end - start))
      val wrappedStream = serializerManager.wrapStream(blockId, bufferedStream)
      deserializeStream = ser.deserializeStream(wrappedStream)
      deserializeStream.asKeyValueIterator.asInstanceOf[Iterator[(K, C)]]
    } else {
      // No more batches left
      cleanup()
      null
    }
  }

  /**
   * Return the next (K, C) pair from the deserialization stream.
   *
   * If the current batch is drained, construct a stream for the next batch and read from it.
   * If no more pairs are left, return null.
   */
  private def readNextItem(): (K, C) = {
    val item = batchIterator.next()
    objectsRead += 1
    if (objectsRead == serializerBatchSize) {
      objectsRead = 0
      batchIterator = nextBatchIterator()
    }
    item
  }

  override def hasNext: Boolean = {
    if (batchIterator == null) {
      // In case of batchIterator has not been initialized
      batchIterator = nextBatchIterator()
      if (batchIterator == null) {
        return false
      }
    }
    batchIterator.hasNext
  }

  override def next(): (K, C) = {
    if (!hasNext) {
      throw new NoSuchElementException
    }
    readNextItem()
  }

  private def cleanup(): Unit = {
    batchIndex = batchOffsets.length  // Prevent reading any other batch
    if (deserializeStream != null) {
      deserializeStream.close()
      deserializeStream = null
    }
    if (fileStream != null) {
      fileStream.close()
      fileStream = null
    }
    if (file.exists()) {
      if (!file.delete()) {
        logWarning(log"Error deleting ${MDC(FILE_NAME, file)}")
      }
    }
  }

  context.addTaskCompletionListener[Unit](context => cleanup())
}

/**
 * Implementation of WritablePartitionedPairCollection that wraps a map in which the keys are tuples
 * of (partition ID, K)
 */
private[spark] class PartitionedAppendOnlyMap[K, V]
  extends SizeTrackingAppendOnlyMap[(Int, K), V] with WritablePartitionedPairCollection[K, V] {

  def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
    : Iterator[((Int, K), V)] = {
    val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
    destructiveSortedIterator(comparator)
  }

  def insert(partition: Int, key: K, value: V): Unit = {
    update((partition, key), value)
  }
}

private[spark] class PartitionedPairBuffer[K, V](initialCapacity: Int = 64)
  extends WritablePartitionedPairCollection[K, V] with SizeTracker
{
  import PartitionedPairBuffer._

  require(initialCapacity <= MAXIMUM_CAPACITY,
    s"Can't make capacity bigger than ${MAXIMUM_CAPACITY} elements")
  require(initialCapacity >= 1, "Invalid initial capacity")

  // Basic growable array data structure. We use a single array of AnyRef to hold both the keys
  // and the values, so that we can sort them efficiently with KVArraySortDataFormat.
  private var capacity = initialCapacity
  private var curSize = 0
  private var data = new Array[AnyRef](2 * initialCapacity)

  /** Add an element into the buffer */
  def insert(partition: Int, key: K, value: V): Unit = {
    if (curSize == capacity) {
      growArray()
    }
    data(2 * curSize) = (partition, key.asInstanceOf[AnyRef])
    data(2 * curSize + 1) = value.asInstanceOf[AnyRef]
    curSize += 1
    afterUpdate()
  }

  /** Double the size of the array because we've reached capacity */
  private def growArray(): Unit = {
    if (capacity >= MAXIMUM_CAPACITY) {
      throw new IllegalStateException(s"Can't insert more than ${MAXIMUM_CAPACITY} elements")
    }
    val newCapacity =
      if (capacity * 2 > MAXIMUM_CAPACITY) { // Overflow
        MAXIMUM_CAPACITY
      } else {
        capacity * 2
      }
    val newArray = new Array[AnyRef](2 * newCapacity)
    System.arraycopy(data, 0, newArray, 0, 2 * capacity)
    data = newArray
    capacity = newCapacity
    resetSamples()
  }

  /** Iterate through the data in a given order. For this class this is not really destructive. */
  override def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
    : Iterator[((Int, K), V)] = {
    val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
    new Sorter(new KVArraySortDataFormat[(Int, K), AnyRef]).sort(data, 0, curSize, comparator)
    iterator()
  }

  private def iterator(): Iterator[((Int, K), V)] = new Iterator[((Int, K), V)] {
    var pos = 0

    override def hasNext: Boolean = pos < curSize

    override def next(): ((Int, K), V) = {
      if (!hasNext) {
        throw new NoSuchElementException
      }
      val pair = (data(2 * pos).asInstanceOf[(Int, K)], data(2 * pos + 1).asInstanceOf[V])
      pos += 1
      pair
    }
  }
}