fix python tests

pdeyhim · Mar 22, 2014 · bf278fa · bf278fa
1 parent e82909c
commit bf278fa
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/python/pyspark/context.py b/python/pyspark/context.py
@@ -383,7 +383,7 @@ def _getJavaStorageLevel(self, storageLevel):
             raise Exception("storageLevel must be of type pyspark.StorageLevel")
 
         newStorageLevel = self._jvm.org.apache.spark.storage.StorageLevel
-        return newStorageLevel(storageLevel.useDisk, storageLevel.useMemory,
+        return newStorageLevel(storageLevel.useDisk, storageLevel.useMemory, storageLevel.userTachyon,
             storageLevel.deserialized, storageLevel.replication)
 
     def setJobGroup(self, groupId, description):

diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
@@ -1178,11 +1178,12 @@ def getStorageLevel(self):
         Get the RDD's current storage level.
         >>> rdd1 = sc.parallelize([1,2])
         >>> rdd1.getStorageLevel()
-        StorageLevel(False, False, False, 1)
+        StorageLevel(False, False, False, False, 1)
         """
         java_storage_level = self._jrdd.getStorageLevel()
         storage_level = StorageLevel(java_storage_level.useDisk(),
                                      java_storage_level.useMemory(),
+                                     java_storage_level.useTachyon(),
                                      java_storage_level.deserialized(),
                                      java_storage_level.replication())
         return storage_level

diff --git a/python/pyspark/storagelevel.py b/python/pyspark/storagelevel.py
@@ -33,8 +33,8 @@ def __init__(self, useDisk, useMemory, useTachyon, deserialized, replication = 1
         self.replication = replication
 
     def __repr__(self):
-        return "StorageLevel(%s, %s, %s, %s)" % (
-            self.useDisk, self.useMemory, self.deserialized, self.replication)
+        return "StorageLevel(%s, %s, %s, %s, %s)" % (
+            self.useDisk, self.useMemory, self.useTachyon, self.deserialized, self.replication)
 
 StorageLevel.DISK_ONLY = StorageLevel(True, False, False, False)
 StorageLevel.DISK_ONLY_2 = StorageLevel(True, False, False, False, 2)