[Feat][Spark] Memory tuning for GraphAr spark with persist and storage level

spark/src/main/java/com/alibaba/graphar/GeneralParams.java

@@ -16,6 +16,8 @@
 
 package com.alibaba.graphar;
 
+import org.apache.spark.storage.StorageLevel;
+
 /** General constant parameters for graphar. */
 public class GeneralParams {
     // column name
@@ -33,4 +35,5 @@ public class GeneralParams {
     public static final Long defaultEdgeChunkSize = 4194304L; // 2^22
     public static final String defaultFileType = "parquet";
     public static final String defaultVersion = "v1";
+    public static final StorageLevel defaultStorageLevel = StorageLevel.MEMORY_AND_DISK_SER();
 }

spark/src/main/scala/com/alibaba/graphar/graph/GraphWriter.scala

@@ -52,12 +52,11 @@ class GraphWriter() {
       )
     }
     vertices += label -> df
-    vertexNums += label -> df.count
     primaryKeys += label -> primaryKey
   }
 
   /**
-   * Put the egde datafrme into writer.
+   * Put the egde dataframe into writer.
    * @param relation
    *   3-Tuple (source label, edge label, target label) to indicate edge type.
    * @param df
@@ -87,15 +86,26 @@ class GraphWriter() {
       scala.collection.mutable.Map[String, DataFrame]()
     vertexInfos.foreach {
       case (label, vertexInfo) => {
-        val vertex_num = vertexNums(label)
         val primaryKey = primaryKeys(label)
+        vertices(label).persist(
+          GeneralParams.defaultStorageLevel
+        ) // cache the vertex DataFrame
         val df_and_mapping = IndexGenerator
           .generateVertexIndexColumnAndIndexMapping(vertices(label), primaryKey)
+        df_and_mapping._1.persist(
+          GeneralParams.defaultStorageLevel
+        ) // cache the vertex DataFrame with index
+        df_and_mapping._2.persist(
+          GeneralParams.defaultStorageLevel
+        ) // cache the index mapping DataFrame
+        vertices(label).unpersist() // unpersist the vertex DataFrame
         val df_with_index = df_and_mapping._1
         indexMappings += label -> df_and_mapping._2
         val writer =
-          new VertexWriter(prefix, vertexInfo, df_with_index, vertex_num)
+          new VertexWriter(prefix, vertexInfo, df_with_index)
+        vertexNums += label -> writer.getVertexNum()
         writer.writeVertexProperties()
+        df_with_index.unpersist()
       }
     }
 
@@ -117,6 +127,9 @@ class GraphWriter() {
             src_vertex_index_mapping,
             dst_vertex_index_mapping
           )
+        edge_df_with_index.persist(
+          GeneralParams.defaultStorageLevel
+        ) // cache the edge DataFrame with index
 
         val adj_lists = edgeInfo.getAdj_lists
         val adj_list_it = adj_lists.iterator
@@ -140,6 +153,7 @@ class GraphWriter() {
           )
           writer.writeEdges()
         }
+        edge_df_with_index.unpersist()
       }
     }
   }

spark/src/main/scala/com/alibaba/graphar/writer/EdgeWriter.scala

@@ -64,6 +64,7 @@ object EdgeWriter {
 
     // sort by primary key and generate continue edge id for edge records
     val sortedDfRDD = edgeDf.sort(colName).rdd
+    sortedDfRDD.persist(GeneralParams.defaultStorageLevel)
     // generate continue edge id for every edge
     val partitionCounts = sortedDfRDD
       .mapPartitionsWithIndex(
@@ -82,6 +83,7 @@ object EdgeWriter {
       val start = broadcastedPartitionCounts.value(i)
       for { (row, j) <- ps.zipWithIndex } yield (start + j, row)
     })
+    rddWithEid.persist(GeneralParams.defaultStorageLevel)
 
     // Construct partitioner for edge chunk
     // get edge num of every vertex chunk
@@ -101,6 +103,8 @@ object EdgeWriter {
         edgeNumMutableMap(i) = 0
       }
     }
+    sortedDfRDD.unpersist() // unpersist the sortedDfRDD
+
     var eidBeginOfVertexChunks =
       new Array[Long](vertexChunkNum + 1) // eid begin of vertex chunks
     var aggEdgeChunkNumOfVertexChunks =
@@ -130,7 +134,8 @@ object EdgeWriter {
     val partitionRDD =
       rddWithEid.repartitionAndSortWithinPartitions(partitioner).values
     val partitionEdgeDf = spark.createDataFrame(partitionRDD, edgeSchema)
-    partitionEdgeDf.cache()
+    rddWithEid.unpersist() // unpersist the rddWithEid
+    partitionEdgeDf.persist(GeneralParams.defaultStorageLevel)
 
     // generate offset DataFrames
     if (
@@ -141,6 +146,7 @@ object EdgeWriter {
           iterator.map(row => (row(colIndex).asInstanceOf[Long], 1))
         })
         .reduceByKey(_ + _)
+      edgeCountsByPrimaryKey.persist(GeneralParams.defaultStorageLevel)
       val offsetDfSchema = StructType(
         Seq(StructField(GeneralParams.offsetCol, IntegerType))
       )
@@ -168,10 +174,11 @@ object EdgeWriter {
             })
             .map { case (k, v) => Row(v) }
           val offsetChunk = spark.createDataFrame(offsetRDD, offsetDfSchema)
-          offsetChunk.cache()
+          offsetChunk.persist(GeneralParams.defaultStorageLevel)
           offsetChunk
         }
       }
+      edgeCountsByPrimaryKey.unpersist() // unpersist the edgeCountsByPrimaryKey
       return (
         partitionEdgeDf,
         offsetDfArray,
@@ -216,6 +223,8 @@ class EdgeWriter(
   validate()
   writeVertexNum()
 
+  edgeDf.persist(GeneralParams.defaultStorageLevel)
+
   // validate data and info
   private def validate(): Unit = {
     // chunk if edge info contains the adj list type
@@ -290,6 +299,7 @@ class EdgeWriter(
         Some(chunkIndex),
         None
       )
+      offsetChunk.unpersist()
       chunkIndex = chunkIndex + 1
     }
   }
@@ -366,4 +376,8 @@ class EdgeWriter(
     writeAdjList()
     writeEdgeProperties()
   }
+
+  override def finalize(): Unit = {
+    edgeDfAndOffsetDf._1.unpersist()
+  }
 }

spark/src/main/scala/com/alibaba/graphar/writer/VertexWriter.scala

@@ -74,6 +74,9 @@ class VertexWriter(
     numVertices: Long = -1
 ) {
   private val spark = vertexDf.sparkSession
+  vertexDf.persist(
+    GeneralParams.defaultStorageLevel
+  ) // cache the vertex DataFrame
   validate()
   private val vertexNum: Long =
     if (numVertices < 0) vertexDf.count else numVertices
@@ -84,6 +87,8 @@ class VertexWriter(
     vertexInfo.getChunk_size(),
     vertexNum
   )
+  vertexDf.unpersist() // unpersist the vertex DataFrame
+  chunks.persist(GeneralParams.defaultStorageLevel)
 
   private def validate(): Unit = {
     // check if vertex DataFrame contains the index_filed
@@ -104,6 +109,8 @@ class VertexWriter(
     )
   }
 
+  def getVertexNum(): Long = vertexNum
+
   /**
    * Generate chunks of the property group for vertex DataFrame.
    *
@@ -146,4 +153,9 @@ class VertexWriter(
       writeVertexProperties(property_group)
     }
   }
+
+  override def finalize(): Unit = {
+    chunks.unpersist()
+  }
+
 }