<commit_before>/*
 * 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.util.collection

import java.io._
import java.util.Comparator

import it.unimi.dsi.fastutil.io.FastBufferedInputStream

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

import org.apache.spark.{Logging, SparkEnv}
import org.apache.spark.serializer.{KryoDeserializationStream, KryoSerializationStream, Serializer}
import org.apache.spark.storage.{BlockId, BlockManager, DiskBlockManager, DiskBlockObjectWriter}

/**
 * An append-only map that spills sorted content to disk when there is insufficient space for it
 * to grow.
 *
 * This map takes two passes over the data:
 *
 *   (1) Values are merged into combiners, which are sorted and spilled to disk as necessary
 *   (2) Combiners are read from disk and merged together
 *
 * The setting of the spill threshold faces the following trade-off: If the spill threshold is
 * too high, the in-memory map may occupy more memory than is available, resulting in OOM.
 * However, if the spill threshold is too low, we spill frequently and incur unnecessary disk
 * writes. This may lead to a performance regression compared to the normal case of using the
 * non-spilling AppendOnlyMap.
 *
 * Two parameters control the memory threshold:
 *
 *   `spark.shuffle.memoryFraction` specifies the collective amount of memory used for storing
 *   these maps as a fraction of the executor's total memory. Since each concurrently running
 *   task maintains one map, the actual threshold for each map is this quantity divided by the
 *   number of running tasks.
 *
 *   `spark.shuffle.safetyFraction` specifies an additional margin of safety as a fraction of
 *   this threshold, in case map size estimation is not sufficiently accurate.
 */

private[spark] class ExternalAppendOnlyMap[K, V, C](
    createCombiner: V => C,
    mergeValue: (C, V) => C,
    mergeCombiners: (C, C) => C,
    serializer: Serializer = SparkEnv.get.serializerManager.default,
    blockManager: BlockManager = SparkEnv.get.blockManager)
  extends Iterable[(K, C)] with Serializable with Logging {

  import ExternalAppendOnlyMap._

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

  // Collective memory threshold shared across all running tasks
  private val maxMemoryThreshold = {
    val memoryFraction = sparkConf.getDouble("spark.shuffle.memoryFraction", 0.3)
    val safetyFraction = sparkConf.getDouble("spark.shuffle.safetyFraction", 0.8)
    (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
  }

  // Number of pairs in the in-memory map
  private var numPairsInMemory = 0L

  // Number of in-memory pairs inserted before tracking the map's shuffle memory usage
  private val trackMemoryThreshold = 1000

  // 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 excess copying when serializing, since some serializers
  // grow internal data structures by growing + copying every time the number of objects doubles.
  private val serializerBatchSize = sparkConf.getLong("spark.shuffle.spill.batchSize", 10000)

  // How many times we have spilled so far
  private var spillCount = 0

  // Number of bytes spilled in total
  private var _memoryBytesSpilled = 0L
  private var _diskBytesSpilled = 0L

  private val fileBufferSize = sparkConf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
  private val syncWrites = sparkConf.getBoolean("spark.shuffle.sync", false)
  private val comparator = new KCComparator[K, C]
  private val ser = serializer.newInstance()

  /**
   * Insert the given key and value into the map.
   *
   * If the underlying map is about 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 insert(key: K, value: V) {
    val update: (Boolean, C) => C = (hadVal, oldVal) => {
      if (hadVal) mergeValue(oldVal, value) else createCombiner(value)
    }
    if (numPairsInMemory > trackMemoryThreshold && currentMap.atGrowThreshold) {
      val mapSize = currentMap.estimateSize()
      var shouldSpill = false
      val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap

      // Atomically check whether there is sufficient memory in the global pool for
      // this map to grow and, if possible, allocate the required amount
      shuffleMemoryMap.synchronized {
        val threadId = Thread.currentThread().getId
        val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
        val availableMemory = maxMemoryThreshold -
          (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))

        // Assume map growth factor is 2x
        shouldSpill = availableMemory < mapSize * 2
        if (!shouldSpill) {
          shuffleMemoryMap(threadId) = mapSize * 2
        }
      }
      // Do not synchronize spills
      if (shouldSpill) {
        spill(mapSize)
      }
    }
    currentMap.changeValue(key, update)
    numPairsInMemory += 1
  }

  /**
   * Sort the existing contents of the in-memory map and spill them to a temporary file on disk
   */
  private def spill(mapSize: Long) {
    spillCount += 1
    logWarning("Spilling in-memory map of %d MB to disk (%d time%s so far)"
      .format(mapSize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
    val (blockId, file) = diskBlockManager.createTempBlock()

    val compressStream: OutputStream => OutputStream = blockManager.wrapForCompression(blockId, _)
    def getNewWriter = new DiskBlockObjectWriter(blockId, file, serializer, fileBufferSize,
      compressStream, syncWrites)

    var writer = getNewWriter
    var objectsWritten = 0
    try {
      val it = currentMap.destructiveSortedIterator(comparator)
      while (it.hasNext) {
        val kv = it.next()
        writer.write(kv)
        objectsWritten += 1

        if (objectsWritten == serializerBatchSize) {
          writer.commit()
          writer = getNewWriter
          objectsWritten = 0
        }
      }

      if (objectsWritten > 0) writer.commit()
    } finally {
      // Partial failures cannot be tolerated; do not revert partial writes
      _diskBytesSpilled += writer.bytesWritten
      writer.close()
    }
    currentMap = new SizeTrackingAppendOnlyMap[K, C]
    spilledMaps.append(new DiskMapIterator(file, blockId))

    // Reset the amount of shuffle memory used by this map in the global pool
    val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
    shuffleMemoryMap.synchronized {
      shuffleMemoryMap(Thread.currentThread().getId) = 0
    }
    numPairsInMemory = 0
    _memoryBytesSpilled += mapSize
  }

  def memoryBytesSpilled: Long = _memoryBytesSpilled
  def diskBytesSpilled: Long = _diskBytesSpilled

  /**
   * Return an 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 (spilledMaps.isEmpty) {
      currentMap.iterator
    } else {
      new ExternalIterator()
    }<commit_msg>Bug fix for reporting of spill output<commit_after>/*
 * 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.util.collection

import java.io._
import java.util.Comparator

import it.unimi.dsi.fastutil.io.FastBufferedInputStream

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

import org.apache.spark.{Logging, SparkEnv}
import org.apache.spark.serializer.{KryoDeserializationStream, KryoSerializationStream, Serializer}
import org.apache.spark.storage.{BlockId, BlockManager, DiskBlockManager, DiskBlockObjectWriter}

/**
 * An append-only map that spills sorted content to disk when there is insufficient space for it
 * to grow.
 *
 * This map takes two passes over the data:
 *
 *   (1) Values are merged into combiners, which are sorted and spilled to disk as necessary
 *   (2) Combiners are read from disk and merged together
 *
 * The setting of the spill threshold faces the following trade-off: If the spill threshold is
 * too high, the in-memory map may occupy more memory than is available, resulting in OOM.
 * However, if the spill threshold is too low, we spill frequently and incur unnecessary disk
 * writes. This may lead to a performance regression compared to the normal case of using the
 * non-spilling AppendOnlyMap.
 *
 * Two parameters control the memory threshold:
 *
 *   `spark.shuffle.memoryFraction` specifies the collective amount of memory used for storing
 *   these maps as a fraction of the executor's total memory. Since each concurrently running
 *   task maintains one map, the actual threshold for each map is this quantity divided by the
 *   number of running tasks.
 *
 *   `spark.shuffle.safetyFraction` specifies an additional margin of safety as a fraction of
 *   this threshold, in case map size estimation is not sufficiently accurate.
 */

private[spark] class ExternalAppendOnlyMap[K, V, C](
    createCombiner: V => C,
    mergeValue: (C, V) => C,
    mergeCombiners: (C, C) => C,
    serializer: Serializer = SparkEnv.get.serializerManager.default,
    blockManager: BlockManager = SparkEnv.get.blockManager)
  extends Iterable[(K, C)] with Serializable with Logging {

  import ExternalAppendOnlyMap._

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

  // Collective memory threshold shared across all running tasks
  private val maxMemoryThreshold = {
    val memoryFraction = sparkConf.getDouble("spark.shuffle.memoryFraction", 0.3)
    val safetyFraction = sparkConf.getDouble("spark.shuffle.safetyFraction", 0.8)
    (Runtime.getRuntime.maxMemory * memoryFraction * safetyFraction).toLong
  }

  // Number of pairs in the in-memory map
  private var numPairsInMemory = 0L

  // Number of in-memory pairs inserted before tracking the map's shuffle memory usage
  private val trackMemoryThreshold = 1000

  // 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 excess copying when serializing, since some serializers
  // grow internal data structures by growing + copying every time the number of objects doubles.
  private val serializerBatchSize = sparkConf.getLong("spark.shuffle.spill.batchSize", 10000)

  // How many times we have spilled so far
  private var spillCount = 0

  // Number of bytes spilled in total
  private var _memoryBytesSpilled = 0L
  private var _diskBytesSpilled = 0L

  private val fileBufferSize = sparkConf.getInt("spark.shuffle.file.buffer.kb", 100) * 1024
  private val syncWrites = sparkConf.getBoolean("spark.shuffle.sync", false)
  private val comparator = new KCComparator[K, C]
  private val ser = serializer.newInstance()

  /**
   * Insert the given key and value into the map.
   *
   * If the underlying map is about 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 insert(key: K, value: V) {
    val update: (Boolean, C) => C = (hadVal, oldVal) => {
      if (hadVal) mergeValue(oldVal, value) else createCombiner(value)
    }
    if (numPairsInMemory > trackMemoryThreshold && currentMap.atGrowThreshold) {
      val mapSize = currentMap.estimateSize()
      var shouldSpill = false
      val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap

      // Atomically check whether there is sufficient memory in the global pool for
      // this map to grow and, if possible, allocate the required amount
      shuffleMemoryMap.synchronized {
        val threadId = Thread.currentThread().getId
        val previouslyOccupiedMemory = shuffleMemoryMap.get(threadId)
        val availableMemory = maxMemoryThreshold -
          (shuffleMemoryMap.values.sum - previouslyOccupiedMemory.getOrElse(0L))

        // Assume map growth factor is 2x
        shouldSpill = availableMemory < mapSize * 2
        if (!shouldSpill) {
          shuffleMemoryMap(threadId) = mapSize * 2
        }
      }
      // Do not synchronize spills
      if (shouldSpill) {
        spill(mapSize)
      }
    }
    currentMap.changeValue(key, update)
    numPairsInMemory += 1
  }

  /**
   * Sort the existing contents of the in-memory map and spill them to a temporary file on disk
   */
  private def spill(mapSize: Long) {
    spillCount += 1
    logWarning("Spilling in-memory map of %d MB to disk (%d time%s so far)"
      .format(mapSize / (1024 * 1024), spillCount, if (spillCount > 1) "s" else ""))
    val (blockId, file) = diskBlockManager.createTempBlock()

    val compressStream: OutputStream => OutputStream = blockManager.wrapForCompression(blockId, _)
    def getNewWriter = new DiskBlockObjectWriter(blockId, file, serializer, fileBufferSize,
      compressStream, syncWrites)

    var writer = getNewWriter
    var objectsWritten = 0
    try {
      val it = currentMap.destructiveSortedIterator(comparator)
      while (it.hasNext) {
        val kv = it.next()
        writer.write(kv)
        objectsWritten += 1

        if (objectsWritten == serializerBatchSize) {
          writer.commit()
          writer.close()
          _diskBytesSpilled += writer.bytesWritten
          writer = getNewWriter
          objectsWritten = 0
        }
      }

      if (objectsWritten > 0) writer.commit()
    } finally {
      // Partial failures cannot be tolerated; do not revert partial writes
      writer.close()
      _diskBytesSpilled += writer.bytesWritten
    }
    currentMap = new SizeTrackingAppendOnlyMap[K, C]
    spilledMaps.append(new DiskMapIterator(file, blockId))

    // Reset the amount of shuffle memory used by this map in the global pool
    val shuffleMemoryMap = SparkEnv.get.shuffleMemoryMap
    shuffleMemoryMap.synchronized {
      shuffleMemoryMap(Thread.currentThread().getId) = 0
    }
    numPairsInMemory = 0
    _memoryBytesSpilled += mapSize
  }

  def memoryBytesSpilled: Long = _memoryBytesSpilled
  def diskBytesSpilled: Long = _diskBytesSpilled

  /**
   * Return an 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 (spilledMaps.isEmpty) {
      currentMap.iterator
    } else {
      new ExternalIterator()
    }
{'fraction_non_alphanumeric': 0.07322872087494056, 'fraction_numerical': 0.006062767475035664, 'mean_word_length': 3.461681251657385, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1894044', 'n_tokens_mistral': 4778, 'n_tokens_neox': 4509, 'n_words': 1747}