Changing CacheManager and BlockManager to pass iterators directly to the serializer when a 'DISK_ONLY' persist is called.

This is in response to SPARK-942.

Author: kellrott

core/src/main/scala/org/apache/spark/CacheManager.scala

@@ -71,10 +71,21 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging {
           val computedValues = rdd.computeOrReadCheckpoint(split, context)
           // Persist the result, so long as the task is not running locally
           if (context.runningLocally) { return computedValues }
-          val elements = new ArrayBuffer[Any]
-          elements ++= computedValues
-          blockManager.put(key, elements, storageLevel, tellMaster = true)
-          return elements.iterator.asInstanceOf[Iterator[T]]
+          if (storageLevel == StorageLevel.DISK_ONLY || storageLevel == StorageLevel.DISK_ONLY_2) {
+            blockManager.put(key, computedValues, storageLevel, tellMaster = true)
+            return blockManager.get(key)  match {
+              case Some(values) =>
+                return new InterruptibleIterator(context, values.asInstanceOf[Iterator[T]])
+              case None =>
+                logInfo("Failure to store %s".format(key));
+                return null;
+            }
+          } else {
+            val elements = new ArrayBuffer[Any]
+            elements ++= computedValues
+            blockManager.put(key, elements, storageLevel, tellMaster = true)
+            return elements.iterator.asInstanceOf[Iterator[T]]
+          }
         } finally {
           loading.synchronized {
             loading.remove(key)

core/src/main/scala/org/apache/spark/serializer/JavaSerializer.scala

@@ -24,7 +24,8 @@ import org.apache.spark.util.ByteBufferInputStream
 
 private[spark] class JavaSerializationStream(out: OutputStream) extends SerializationStream {
   val objOut = new ObjectOutputStream(out)
-  def writeObject[T](t: T): SerializationStream = { objOut.writeObject(t); this }
+  //Calling reset to avoid memory leak: http://stackoverflow.com/questions/1281549/memory-leak-traps-in-the-java-standard-api
+  def writeObject[T](t: T): SerializationStream = { objOut.writeObject(t); objOut.reset(); this }
   def flush() { objOut.flush() }
   def close() { objOut.close() }
 }

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -356,7 +356,7 @@ private[spark] class BlockManager(
                 // TODO: Consider creating a putValues that also takes in a iterator?
                 val valuesBuffer = new ArrayBuffer[Any]
                 valuesBuffer ++= values
-                memoryStore.putValues(blockId, valuesBuffer, level, true).data match {
+                memoryStore.putValues(blockId, valuesBuffer.toIterator, level, true).data match {
                   case Left(values2) =>
                     return Some(values2)
                   case _ =>
@@ -451,9 +451,7 @@ private[spark] class BlockManager(
 
   def put(blockId: BlockId, values: Iterator[Any], level: StorageLevel, tellMaster: Boolean)
     : Long = {
-    val elements = new ArrayBuffer[Any]
-    elements ++= values
-    put(blockId, elements, level, tellMaster)
+    doPut(blockId, Left(values), level, tellMaster)
   }
 
   /**
@@ -474,7 +472,7 @@ private[spark] class BlockManager(
   def put(blockId: BlockId, values: ArrayBuffer[Any], level: StorageLevel,
           tellMaster: Boolean = true) : Long = {
     require(values != null, "Values is null")
-    doPut(blockId, Left(values), level, tellMaster)
+    doPut(blockId, Left(values.toIterator), level, tellMaster)
   }
 
   /**
@@ -486,7 +484,7 @@ private[spark] class BlockManager(
     doPut(blockId, Right(bytes), level, tellMaster)
   }
 
-  private def doPut(blockId: BlockId, data: Either[ArrayBuffer[Any], ByteBuffer],
+  private def doPut(blockId: BlockId, data: Either[Iterator[Any], ByteBuffer],
                     level: StorageLevel, tellMaster: Boolean = true): Long = {
     require(blockId != null, "BlockId is null")
     require(level != null && level.isValid, "StorageLevel is null or invalid")
@@ -691,7 +689,7 @@ private[spark] class BlockManager(
           logInfo("Writing block " + blockId + " to disk")
           data match {
             case Left(elements) =>
-              diskStore.putValues(blockId, elements, level, false)
+              diskStore.putValues(blockId, elements.toIterator, level, false)
             case Right(bytes) =>
               diskStore.putBytes(blockId, bytes, level)
           }

core/src/main/scala/org/apache/spark/storage/BlockStore.scala

@@ -36,7 +36,7 @@ abstract class BlockStore(val blockManager: BlockManager) extends Logging {
    * @return a PutResult that contains the size of the data, as well as the values put if
    *         returnValues is true (if not, the result's data field can be null)
    */
-  def putValues(blockId: BlockId, values: ArrayBuffer[Any], level: StorageLevel,
+  def putValues(blockId: BlockId, values: Iterator[Any], level: StorageLevel,
     returnValues: Boolean) : PutResult
 
   /**

core/src/main/scala/org/apache/spark/storage/DiskStore.scala

@@ -57,7 +57,7 @@ private class DiskStore(blockManager: BlockManager, diskManager: DiskBlockManage
 
   override def putValues(
       blockId: BlockId,
-      values: ArrayBuffer[Any],
+      values: Iterator[Any],
       level: StorageLevel,
       returnValues: Boolean)
     : PutResult = {
@@ -66,7 +66,7 @@ private class DiskStore(blockManager: BlockManager, diskManager: DiskBlockManage
     val startTime = System.currentTimeMillis
     val file = diskManager.getFile(blockId)
     val outputStream = new FileOutputStream(file)
-    blockManager.dataSerializeStream(blockId, outputStream, values.iterator)
+    blockManager.dataSerializeStream(blockId, outputStream, values)
     val length = file.length
 
     val timeTaken = System.currentTimeMillis - startTime

core/src/main/scala/org/apache/spark/storage/MemoryStore.scala

@@ -65,17 +65,17 @@ private class MemoryStore(blockManager: BlockManager, maxMemory: Long)
 
   override def putValues(
       blockId: BlockId,
-      values: ArrayBuffer[Any],
+      values: Iterator[Any],
       level: StorageLevel,
       returnValues: Boolean)
     : PutResult = {
 
     if (level.deserialized) {
       val sizeEstimate = SizeEstimator.estimate(values.asInstanceOf[AnyRef])
       tryToPut(blockId, values, sizeEstimate, true)
-      PutResult(sizeEstimate, Left(values.iterator))
+      PutResult(sizeEstimate, Left(values))
     } else {
-      val bytes = blockManager.dataSerialize(blockId, values.iterator)
+      val bytes = blockManager.dataSerialize(blockId, values)
       tryToPut(blockId, bytes, bytes.limit, false)
       PutResult(bytes.limit(), Right(bytes.duplicate()))
     }