diff options
author | Aaron Davidson <aaron@databricks.com> | 2013-09-05 21:43:26 -0700 |
---|---|---|
committer | Aaron Davidson <aaron@databricks.com> | 2013-09-05 21:43:26 -0700 |
commit | 3a04e76c89ac44a1345417597971e6e37944d371 (patch) | |
tree | 9d59caba4b2a94f71b40619254b772dcc0d844a3 | |
parent | 4f2236a1c5fdc6c3a5711cef73407f926bf1a779 (diff) | |
download | spark-3a04e76c89ac44a1345417597971e6e37944d371.tar.gz spark-3a04e76c89ac44a1345417597971e6e37944d371.tar.bz2 spark-3a04e76c89ac44a1345417597971e6e37944d371.zip |
Reynold's second round of comments
-rw-r--r-- | core/src/main/scala/org/apache/spark/CacheManager.scala | 2 | ||||
-rw-r--r-- | core/src/test/scala/org/apache/spark/CacheManagerSuite.scala | 34 |
2 files changed, 19 insertions, 17 deletions
diff --git a/core/src/main/scala/org/apache/spark/CacheManager.scala b/core/src/main/scala/org/apache/spark/CacheManager.scala index 85a2b2b331..68b99ca125 100644 --- a/core/src/main/scala/org/apache/spark/CacheManager.scala +++ b/core/src/main/scala/org/apache/spark/CacheManager.scala @@ -69,7 +69,7 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging { logInfo("Computing partition " + split) val computedValues = rdd.computeOrReadCheckpoint(split, context) // Persist the result, so long as the task is not running locally - if (context.runningLocally) return computedValues + if (context.runningLocally) { return computedValues } val elements = new ArrayBuffer[Any] elements ++= computedValues blockManager.put(key, elements, storageLevel, true) diff --git a/core/src/test/scala/org/apache/spark/CacheManagerSuite.scala b/core/src/test/scala/org/apache/spark/CacheManagerSuite.scala index a85d666ace..214f8244d5 100644 --- a/core/src/test/scala/org/apache/spark/CacheManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/CacheManagerSuite.scala @@ -17,7 +17,6 @@ package org.apache.spark -import scala.Array import scala.collection.mutable.ArrayBuffer import org.scalatest.{BeforeAndAfter, FunSuite} @@ -25,27 +24,30 @@ import org.scalatest.mock.EasyMockSugar import org.apache.spark.rdd.RDD import org.apache.spark.storage.{BlockManager, StorageLevel} -import org.junit.Assert // TODO: Test the CacheManager's thread-safety aspects class CacheManagerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar { - val sc = new SparkContext("local", "test") - - val split = new Partition { override def index: Int = 0 } - - // An RDD which returns the values [1, 2, 3, 4]. - val rdd = new RDD[Int](sc, Nil) { - override def getPartitions: Array[Partition] = Array(split) - override val getDependencies = List[Dependency[_]]() - override def compute(split: Partition, context: TaskContext) = Array(1, 2, 3, 4).iterator - } - + var sc : SparkContext = _ var blockManager: BlockManager = _ var cacheManager: CacheManager = _ + var split: Partition = _ + /** An RDD which returns the values [1, 2, 3, 4]. */ + var rdd: RDD[Int] = _ before { + sc = new SparkContext("local", "test") blockManager = mock[BlockManager] cacheManager = new CacheManager(blockManager) + split = new Partition { override def index: Int = 0 } + rdd = new RDD(sc, Nil) { + override def getPartitions: Array[Partition] = Array(split) + override val getDependencies = List[Dependency[_]]() + override def compute(split: Partition, context: TaskContext) = Array(1, 2, 3, 4).iterator + } + } + + after { + sc.stop() } test("get uncached rdd") { @@ -58,7 +60,7 @@ class CacheManagerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar whenExecuting(blockManager) { val context = new TaskContext(0, 0, 0, runningLocally = false, null) val value = cacheManager.getOrCompute(rdd, split, context, StorageLevel.MEMORY_ONLY) - Assert.assertEquals(value.toList, List(1, 2, 3, 4)) + assert(value.toList === List(1, 2, 3, 4)) } } @@ -70,7 +72,7 @@ class CacheManagerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar whenExecuting(blockManager) { val context = new TaskContext(0, 0, 0, runningLocally = false, null) val value = cacheManager.getOrCompute(rdd, split, context, StorageLevel.MEMORY_ONLY) - Assert.assertEquals(value.toList, List(5, 6, 7)) + assert(value.toList === List(5, 6, 7)) } } @@ -83,7 +85,7 @@ class CacheManagerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar whenExecuting(blockManager) { val context = new TaskContext(0, 0, 0, runningLocally = true, null) val value = cacheManager.getOrCompute(rdd, split, context, StorageLevel.MEMORY_ONLY) - Assert.assertEquals(value.toList, List(1, 2, 3, 4)) + assert(value.toList === List(1, 2, 3, 4)) } } } |