def test_storage_levels(self):
sqlContext = SQLContext(self.sc)
df = sqlContext.createDataFrame(
[(0, 0, 4.0), (0, 1, 2.0), (1, 1, 3.0), (1, 2, 4.0), (2, 1, 1.0), (2, 2, 5.0)],
["user", "item", "rating"])
als = ALS().setMaxIter(1).setRank(1)
# test default params
als.fit(df)
self.assertEqual(als.getIntermediateRDDStorageLevel(), "MEMORY_AND_DISK")
self.assertEqual(als._java_obj.getIntermediateRDDStorageLevel(), "MEMORY_AND_DISK")
self.assertEqual(als.getFinalRDDStorageLevel(), "MEMORY_AND_DISK")
self.assertEqual(als._java_obj.getFinalRDDStorageLevel(), "MEMORY_AND_DISK")
# test non-default params
als.setIntermediateRDDStorageLevel("MEMORY_ONLY_2")
als.setFinalRDDStorageLevel("DISK_ONLY")
als.fit(df)
self.assertEqual(als.getIntermediateRDDStorageLevel(), "MEMORY_ONLY_2")
self.assertEqual(als._java_obj.getIntermediateRDDStorageLevel(), "MEMORY_ONLY_2")
self.assertEqual(als.getFinalRDDStorageLevel(), "DISK_ONLY")
self.assertEqual(als._java_obj.getFinalRDDStorageLevel(), "DISK_ONLY")
