Simplifying StorageLevel checks

kellrott · kellrott · commit 95c7f67b1314 · 2014-02-24T13:35:05.000-08:00
diff --git a/core/src/main/scala/org/apache/spark/CacheManager.scala b/core/src/main/scala/org/apache/spark/CacheManager.scala
@@ -71,7 +71,7 @@ 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 }
-          if (storageLevel == StorageLevel.DISK_ONLY || storageLevel == StorageLevel.DISK_ONLY_2) {
+          if (storageLevel.useDisk || !storageLevel.useMemory) {
             blockManager.put(key, computedValues, storageLevel, tellMaster = true)
             return blockManager.get(key)  match {
               case Some(values) =>
