Avoid many small spills

Author: Andrew Or

core/src/main/scala/org/apache/spark/util/collection/Spillable.scala

@@ -24,9 +24,8 @@ import org.apache.spark.SparkEnv
  * Spills contents of an in-memory collection to disk when the memory threshold
  * has been exceeded.
  */
-private[spark] trait Spillable[C] {
-
-  this: Logging =>
+private[spark] trait Spillable[C] extends Logging {
+  import Spillable._
 
   /**
    * Spills the current in-memory collection to disk, and releases the memory.
@@ -41,11 +40,12 @@ private[spark] trait Spillable[C] {
   // Memory manager that can be used to acquire/release memory
   private[this] val shuffleMemoryManager = SparkEnv.get.shuffleMemoryManager
 
-  // What threshold of elementsRead we start estimating collection size at
+  // Threshold for `elementsRead` before we start tracking this collection's memory usage
   private[this] val trackMemoryThreshold = 1000
 
-  // How much of the shared memory pool this collection has claimed
-  private[this] var myMemoryThreshold = 0L
+  // 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
+  private[this] var myMemoryThreshold = INITIAL_MEMORY_TRACKING_THRESHOLD
 
   // Number of bytes spilled in total
   private[this] var _memoryBytesSpilled = 0L
@@ -94,8 +94,9 @@ private[spark] trait Spillable[C] {
    * Release our memory back to the shuffle pool so that other threads can grab it.
    */
   private def releaseMemoryForThisThread(): Unit = {
-    shuffleMemoryManager.release(myMemoryThreshold)
-    myMemoryThreshold = 0L
+    // The amount we requested does not include the initial memory tracking threshold
+    shuffleMemoryManager.release(myMemoryThreshold - INITIAL_MEMORY_TRACKING_THRESHOLD)
+    myMemoryThreshold = INITIAL_MEMORY_TRACKING_THRESHOLD
   }
 
   /**
@@ -106,7 +107,12 @@ private[spark] trait Spillable[C] {
   @inline private def logSpillage(size: Long) {
     val threadId = Thread.currentThread().getId
     logInfo("Thread %d spilling in-memory map of %s to disk (%d time%s so far)"
-        .format(threadId, org.apache.spark.util.Utils.bytesToString(size),
-            _spillCount, if (_spillCount > 1) "s" else ""))
+      .format(threadId, org.apache.spark.util.Utils.bytesToString(size),
+        _spillCount, if (_spillCount > 1) "s" else ""))
   }
 }
+
+private object Spillable {
+  // Initial threshold for the size of a collection before we start tracking its memory usage
+  val INITIAL_MEMORY_TRACKING_THRESHOLD: Long = 5 * 1024 * 1024
+}