Fixing issues (formatting, variable names, etc.) from review comments

Author: kellrott

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

@@ -72,15 +72,24 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging {
           // Persist the result, so long as the task is not running locally
           if (context.runningLocally) { return computedValues }
           if (storageLevel.useDisk && !storageLevel.useMemory) {
+            // In the case that this RDD is to be persisted using DISK_ONLY
+            // the iterator will be passed directly to the blockManager (rather then
+            // caching it to an ArrayBuffer first), then the resulting block data iterator
+            // will be passed back to the user. If the iterator generates a lot of data,
+            // this means that it doesn't all have to be held in memory at one time.
+            // This could also apply to MEMORY_ONLY_SER storage, but we need to make sure
+            // blocks aren't dropped by the block store before enabling that.
             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
+                throw new Exception("Block manager failed to return persisted valued")
             }
           } else {
+            // In this case the RDD is cached to an array buffer. This will save the results
+            // if we're dealing with a 'one-time' iterator
             val elements = new ArrayBuffer[Any]
             elements ++= computedValues
             blockManager.put(key, elements, storageLevel, tellMaster = true)

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

@@ -540,7 +540,7 @@ private[spark] class BlockManager(
     // If we're storing bytes, then initiate the replication before storing them locally.
     // This is faster as data is already serialized and ready to send.
     val replicationFuture = if (data.isInstanceOf[ByteBufferValues] && level.replication > 1) {
-      //Duplicate doesn't copy the bytes, just creates a wrapper
+      // Duplicate doesn't copy the bytes, just creates a wrapper
       val bufferView = data.asInstanceOf[ByteBufferValues].buffer.duplicate()
       Future {
         replicate(blockId, bufferView, level)
@@ -559,13 +559,13 @@ private[spark] class BlockManager(
           // Save it just to memory first, even if it also has useDisk set to true; we will
           // drop it to disk later if the memory store can't hold it.
           val res = data match {
-            case IteratorValues(values_i) =>
-              memoryStore.putValues(blockId, values_i, level, true)
-            case ArrayBufferValues(values_a) =>
-              memoryStore.putValues(blockId, values_a, level, true)
-            case ByteBufferValues(value_bytes) => {
-              value_bytes.rewind();
-              memoryStore.putBytes(blockId, value_bytes, level)
+            case IteratorValues(iterator) =>
+              memoryStore.putValues(blockId, iterator, level, true)
+            case ArrayBufferValues(array) =>
+              memoryStore.putValues(blockId, array, level, true)
+            case ByteBufferValues(bytes) => {
+              bytes.rewind();
+              memoryStore.putBytes(blockId, bytes, level)
             }
           }
           size = res.size
@@ -579,13 +579,13 @@ private[spark] class BlockManager(
           val askForBytes = level.replication > 1
 
           val res = data match {
-            case IteratorValues(values_i) =>
-              diskStore.putValues(blockId, values_i, level, askForBytes)
-            case ArrayBufferValues(values_a) =>
-              diskStore.putValues(blockId, values_a, level, askForBytes)
-            case ByteBufferValues(value_bytes) => {
-              value_bytes.rewind();
-              diskStore.putBytes(blockId, value_bytes, level)
+            case IteratorValues(iterator) =>
+              diskStore.putValues(blockId, iterator, level, askForBytes)
+            case ArrayBufferValues(array) =>
+              diskStore.putValues(blockId, array, level, askForBytes)
+            case ByteBufferValues(bytes) => {
+              bytes.rewind();
+              diskStore.putBytes(blockId, bytes, level)
             }
           }
           size = res.size

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

@@ -41,7 +41,7 @@ abstract class BlockStore(val blockManager: BlockManager) extends Logging {
     returnValues: Boolean) : PutResult
 
   def putValues(blockId: BlockId, values: ArrayBuffer[Any], level: StorageLevel,
-                returnValues: Boolean) : PutResult
+    returnValues: Boolean) : PutResult
 
   /**
    * Return the size of a block in bytes.

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

@@ -56,10 +56,10 @@ private class DiskStore(blockManager: BlockManager, diskManager: DiskBlockManage
   }
 
   override def putValues(
-                          blockId: BlockId,
-                          values: ArrayBuffer[Any],
-                          level: StorageLevel,
-                          returnValues: Boolean)
+      blockId: BlockId,
+      values: ArrayBuffer[Any],
+      level: StorageLevel,
+      returnValues: Boolean)
   : PutResult = {
     return putValues(blockId, values.toIterator, level, returnValues)
   }

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

@@ -67,10 +67,10 @@ private class MemoryStore(blockManager: BlockManager, maxMemory: Long)
   }
 
   override def putValues(
-                          blockId: BlockId,
-                          values: ArrayBuffer[Any],
-                          level: StorageLevel,
-                          returnValues: Boolean)
+      blockId: BlockId,
+      values: ArrayBuffer[Any],
+      level: StorageLevel,
+      returnValues: Boolean)
   : PutResult = {
     if (level.deserialized) {
       val sizeEstimate = SizeEstimator.estimate(values.asInstanceOf[AnyRef])

core/src/test/scala/org/apache/spark/storage/FlatmapIteratorSuite.scala

@@ -36,48 +36,39 @@ class FlatmapIteratorSuite extends FunSuite with LocalSparkContext {
     val sconf = new SparkConf().setMaster("local-cluster[1,1,512]")
       .setAppName("iterator_to_disk_test")
     sc = new SparkContext(sconf)
-    try {
-      val expand_size = 100
-      val data = sc.parallelize( (1 to 5).toSeq ).
-        flatMap( x => Stream.range(0, expand_size) )
-      var persisted = data.persist(StorageLevel.DISK_ONLY)
-      println(persisted.count())
-      assert( persisted.count() == 500)
-      assert( persisted.filter( _==1 ).count() == 5 )
-    } catch {
-      case _ : OutOfMemoryError => assert(false)
-    }
+    val expand_size = 100
+    val data = sc.parallelize((1 to 5).toSeq).
+      flatMap( x => Stream.range(0, expand_size))
+    var persisted = data.persist(StorageLevel.DISK_ONLY)
+    println(persisted.count())
+    assert(persisted.count()===500)
+    assert(persisted.filter(_==1).count()===5)
   }
 
   test("Flatmap Iterator to Memory") {
     val sconf = new SparkConf().setMaster("local-cluster[1,1,512]")
       .setAppName("iterator_to_disk_test")
     sc = new SparkContext(sconf)
-    try {
-      val expand_size = 100
-      val data = sc.parallelize( (1 to 5).toSeq ).
-        flatMap( x => Stream.range(0, expand_size) )
-      var persisted = data.persist(StorageLevel.MEMORY_ONLY)
-      println(persisted.count())
-      assert( persisted.count() == 500)
-      assert( persisted.filter( _==1 ).count() == 5 )
-    } catch {
-      case _ : OutOfMemoryError => assert(false)
-    }
+    val expand_size = 100
+    val data = sc.parallelize((1 to 5).toSeq).
+      flatMap(x => Stream.range(0, expand_size))
+    var persisted = data.persist(StorageLevel.MEMORY_ONLY)
+    println(persisted.count())
+    assert(persisted.count()===500)
+    assert(persisted.filter(_==1).count()===5)
   }
 
   test("Serializer Reset") {
     val sconf = new SparkConf().setMaster("local-cluster[1,1,512]")
       .setAppName("serializer_reset_test")
       .set("spark.serializer.objectStreamReset", "10")
-
     sc = new SparkContext(sconf)
     val expand_size = 500
-    val data = sc.parallelize( Seq(1,2) ).
-      flatMap( x => Stream.range(1, expand_size).
-      map( y => "%d: string test %d".format(y,x) ) )
+    val data = sc.parallelize(Seq(1,2)).
+      flatMap(x => Stream.range(1, expand_size).
+      map(y => "%d: string test %d".format(y,x)))
     var persisted = data.persist(StorageLevel.MEMORY_ONLY_SER)
-    assert( persisted.filter( _.startsWith("1:") ).count() == 2 )
+    assert(persisted.filter(_.startsWith("1:")).count()===2)
   }
 
 }