[Improvement][Spark] Change VertexWriter constructor signature (#314)

On branch 313-replace-scala-option Changes to be committed: modified: spark/src/main/scala/com/alibaba/graphar/graph/GraphWriter.scala modified: spark/src/main/scala/com/alibaba/graphar/writer/VertexWriter.scala
apache · Jan 3, 2024 · 9ca5bc3 · 9ca5bc3
1 parent 4d626cd
commit 9ca5bc3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/spark/src/main/scala/com/alibaba/graphar/graph/GraphWriter.scala b/spark/src/main/scala/com/alibaba/graphar/graph/GraphWriter.scala
@@ -94,7 +94,7 @@ class GraphWriter() {
         val df_with_index = df_and_mapping._1
         indexMappings += label -> df_and_mapping._2
         val writer =
-          new VertexWriter(prefix, vertexInfo, df_with_index, Some(vertex_num))
+          new VertexWriter(prefix, vertexInfo, df_with_index, vertex_num)
         writer.writeVertexProperties()
       }
     }

diff --git a/spark/src/main/scala/com/alibaba/graphar/writer/VertexWriter.scala b/spark/src/main/scala/com/alibaba/graphar/writer/VertexWriter.scala
@@ -63,19 +63,20 @@ object VertexWriter {
  *   the vertex info that describes the vertex type.
  * @param vertexDf
  *   the input vertex DataFrame.
+ * @param numVertices
+ *   the number of vertices, if negative value is passed, then vertexDF.count
+ *   will be used
  */
 class VertexWriter(
     prefix: String,
     vertexInfo: VertexInfo,
     vertexDf: DataFrame,
-    numVertices: Option[Long] = None
+    numVertices: Long = -1
 ) {
   private val spark = vertexDf.sparkSession
   validate()
-  private val vertexNum: Long = numVertices match {
-    case None => vertexDf.count()
-    case _    => numVertices.get
-  }
+  private val vertexNum: Long =
+    if (numVertices < 0) vertexDf.count else numVertices
   writeVertexNum()
 
   private var chunks: DataFrame = VertexWriter.repartitionAndSort(