SPARK-660: Add StorageLevel support in Python

It uses reflection... I am not proud of that fact, but it at least ensures
compatibility (sans refactoring of the StorageLevel stuff).

1 files changed, 11 insertions, 0 deletions

diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala
index ccd3833964..6ca56b3af6 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala
@@ -28,6 +28,7 @@ import org.apache.spark.broadcast.Broadcast
 import org.apache.spark._
 import org.apache.spark.rdd.RDD
 import org.apache.spark.rdd.PipedRDD
+import org.apache.spark.storage.StorageLevel
 import org.apache.spark.util.Utils
 
 
@@ -270,6 +271,16 @@ private[spark] object PythonRDD {
     JavaRDD.fromRDD(sc.sc.parallelize(objs, parallelism))
   }
 
+  /**
+   * Returns the StorageLevel with the given string name.
+   * Throws an exception if the name is not a valid StorageLevel.
+   */
+  def getStorageLevel(name: String) : StorageLevel = {
+    // In Scala, "val MEMORY_ONLY" produces a public getter by the same name.
+    val storageLevelGetter = StorageLevel.getClass().getDeclaredMethod(name)
+    return storageLevelGetter.invoke(StorageLevel).asInstanceOf[StorageLevel]
+  }
+
   def writeIteratorToPickleFile[T](items: java.util.Iterator[T], filename: String) {
     import scala.collection.JavaConverters._
     writeIteratorToPickleFile(items.asScala, filename)