make intermediate storage level configurable in CC (graphframes#215)

The default storage level MEMORY_AND_DISK might not be the best option for certain setup. We can make it configurable by users.

Author: mengxr

src/main/scala/org/graphframes/lib/ConnectedComponents.scala

@@ -132,14 +132,30 @@ class ConnectedComponents private[graphframes] (
    */
   def getCheckpointInterval: Int = checkpointInterval
 
+  private var intermediateStorageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK
+
+  /**
+   * Sets storage level for intermediate datasets that require multiple passes (default: ``MEMORY_AND_DISK``).
+   */
+  def setIntermediateStorageLevel(value: StorageLevel): this.type = {
+    intermediateStorageLevel = value
+    this
+  }
+
+  /**
+   * Gets storage level for intermediate datasets that require multiple passes.
+   */
+  def getIntermediateStorageLevel: StorageLevel = intermediateStorageLevel
+
   /**
    * Runs the algorithm.
    */
   def run(): DataFrame = {
     ConnectedComponents.run(graph,
       algorithm = algorithm,
       broadcastThreshold = broadcastThreshold,
-      checkpointInterval = checkpointInterval)
+      checkpointInterval = checkpointInterval,
+      intermediateStorageLevel = intermediateStorageLevel)
   }
 }
 
@@ -256,7 +272,8 @@ object ConnectedComponents extends Logging {
       graph: GraphFrame,
       algorithm: String,
       broadcastThreshold: Int,
-      checkpointInterval: Int): DataFrame = {
+      checkpointInterval: Int,
+      intermediateStorageLevel: StorageLevel): DataFrame = {
     require(supportedAlgorithms.contains(algorithm),
       s"Supported algorithms are {${supportedAlgorithms.mkString(", ")}}, but got $algorithm.")
 
@@ -290,7 +307,7 @@ object ConnectedComponents extends Logging {
     logger.info(s"$logPrefix Preparing the graph for connected component computation ...")
     val g = prepare(graph)
     val vv = g.vertices
-    var ee = g.edges.persist(StorageLevel.MEMORY_AND_DISK) // src < dst
+    var ee = g.edges.persist(intermediateStorageLevel) // src < dst
     val numEdges = ee.count()
     logger.info(s"$logPrefix Found $numEdges edges after preparation.")
 
@@ -301,17 +318,17 @@ object ConnectedComponents extends Logging {
       // large-star step
       // compute min neighbors (including self-min)
       val minNbrs1 = minNbrs(ee) // src >= min_nbr
-        .persist(StorageLevel.MEMORY_AND_DISK)
+        .persist(intermediateStorageLevel)
       // connect all strictly larger neighbors to the min neighbor (including self)
       ee = skewedJoin(ee, minNbrs1, broadcastThreshold, logPrefix)
         .select(col(DST).as(SRC), col(MIN_NBR).as(DST)) // src > dst
         .distinct()
-        .persist(StorageLevel.MEMORY_AND_DISK)
+        .persist(intermediateStorageLevel)
 
       // small-star step
       // compute min neighbors (excluding self-min)
       val minNbrs2 = ee.groupBy(col(SRC)).agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT)) // src > min_nbr
-        .persist(StorageLevel.MEMORY_AND_DISK)
+        .persist(intermediateStorageLevel)
       // connect all smaller neighbors to the min neighbor
       ee = skewedJoin(ee, minNbrs2, broadcastThreshold, logPrefix)
         .select(col(MIN_NBR).as(SRC), col(DST)) // src <= dst
@@ -337,7 +354,7 @@ object ConnectedComponents extends Logging {
         System.gc() // hint Spark to clean shuffle directories
       }
 
-      ee.persist(StorageLevel.MEMORY_AND_DISK)
+      ee.persist(intermediateStorageLevel)
 
       // test convergence
 

src/test/scala/org/graphframes/lib/ConnectedComponentsSuite.scala

@@ -25,6 +25,7 @@ import scala.reflect.runtime.universe.TypeTag
 import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.sql.functions.{col, lit}
 import org.apache.spark.sql.types.DataTypes
+import org.apache.spark.storage.StorageLevel
 
 import org.graphframes._
 import org.graphframes.GraphFrame._
@@ -213,6 +214,22 @@ class ConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkCon
       "The result shouldn't depend on checkpoint data if converged before first checkpoint.")
   }
 
+  test("intermediate storage level") {
+    val friends = Graphs.friends
+    val expected = Set(Set("a", "b", "c", "d", "e", "f"), Set("g"))
+
+    val cc = friends.connectedComponents
+    assert(cc.getIntermediateStorageLevel === StorageLevel.MEMORY_AND_DISK)
+
+    for (storageLevel <- Seq(StorageLevel.DISK_ONLY, StorageLevel.MEMORY_ONLY, StorageLevel.NONE)) {
+      // TODO: it is not trivial to confirm the actual storage level used
+      val components = cc
+        .setIntermediateStorageLevel(storageLevel)
+        .run()
+      assertComponents(components, expected)
+    }
+  }
+
   private def assertComponents[T: ClassTag:TypeTag](
       actual: DataFrame,
       expected: Set[Set[T]]): Unit = {