Memoize StorageLevels read from JVM

core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala

@@ -275,7 +275,7 @@ private[spark] object PythonRDD {
    * Returns the StorageLevel with the given string name.
    * Throws an exception if the name is not a valid StorageLevel.
    */
-  def getStorageLevel(name: String) : StorageLevel = {
+  def getStorageLevelByName(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]

python/pyspark/context.py

@@ -281,16 +281,23 @@ class SparkContext(object):
 
 class StorageLevelReader:
     """
-    Mimics the Scala StorageLevel by directing all attribute requests
+    Mimics the Scala StorageLevel by delegating all attribute requests
     (e.g., StorageLevel.DISK_ONLY) to the JVM for reflection.
+    Memoizes results to reduce JVM call/memory overheads.
     """
 
     def __init__(self, sc):
         self.sc = sc
+        self.memoized = {}
 
     def __getattr__(self, name):
+        if name in self.memoized: 
+          return self.memoized[name]
+
         try:
-            return self.sc._jvm.PythonRDD.getStorageLevel(name)
+            storageLevel = self.sc._jvm.PythonRDD.getStorageLevelByName(name)
+            self.memoized[name] = storageLevel
+            return storageLevel
         except:
             print "Failed to find StorageLevel:", name