Repository: Angel-ML/angel (language: java)

Issue:
[Bug] v3.2.0 duplicate nodes in pagerank result
duplicate nodes in pagerank result， when version is vertex-cut or router type is hash




Discussion / hints:

Fix patch:
diff --git a/spark-on-angel/examples/src/main/scala/com/tencent/angel/spark/examples/cluster/PageRankExample.scala b/spark-on-angel/examples/src/main/scala/com/tencent/angel/spark/examples/cluster/PageRankExample.scala
index 0b5a3ed61..87a84bf94 100644
--- a/spark-on-angel/examples/src/main/scala/com/tencent/angel/spark/examples/cluster/PageRankExample.scala
+++ b/spark-on-angel/examples/src/main/scala/com/tencent/angel/spark/examples/cluster/PageRankExample.scala
@@ -48,6 +48,7 @@ object PageRankExample {
     val balancePartitionPercent = params.getOrElse("balancePartitionPercent", "0.7").toFloat
     val version = params.getOrElse("version", "edge-cut")
     val numBatch = params.getOrElse("numBatch", "1").toInt
+    val batchSize = params.getOrElse("batchSize", "1000").toInt
 
     val sep = params.getOrElse("sep", "space") match {
       case "space" => " "
@@ -63,10 +64,10 @@ object PageRankExample {
     val ranks = version match {
       case "edge-cut" => edgeCutPageRank(edges, partitionNum, psPartitionNum,
         storageLevel, tol, resetProp, isWeight,
-        useBalancePartition, balancePartitionPercent, numBatch)
+        useBalancePartition, balancePartitionPercent, numBatch, batchSize)
       case "vertex-cut" => vertexCutPageRank(edges, partitionNum, psPartitionNum,
         storageLevel, tol, resetProp, isWeight,
-        useBalancePartition, balancePartitionPercent, numBatch)
+        useBalancePartition, balancePartitionPercent, numBatch, batchSize)
     }
 
     GraphIO.save(ranks, output)
@@ -78,7 +79,7 @@ object PageRankExample {
                         storageLevel: StorageLevel,
                         tol: Float, resetProb: Float, isWeighted: Boolean,
                         useBalancePartition: Boolean,
-                        balancePartitionPercent: Float, numBatch: Int): DataFrame = {
+                        balancePartitionPercent: Float, numBatch: Int,  batchSize: Int): DataFrame = {
     val pageRank = new VertexCutPageRank()
       .setPartitionNum(partitionNum)
       .setStorageLevel(storageLevel)
@@ -89,6 +90,7 @@ object PageRankExample {
       .setUseBalancePartition(useBalancePartition)
       .setBalancePartitionPercent(balancePartitionPercent)
       .setNumBatch(numBatch)
+      .setBatchSize(batchSize)
     pageRank.transform(edges)
   }
 
@@ -97,7 +99,7 @@ object PageRankExample {
                       storageLevel: StorageLevel,
                       tol: Float, resetProb: Float, isWeighted: Boolean,
                       useBalancePartition: Boolean,
-                      balancePartitionPercent: Float, numBatch: Int): DataFrame = {
+                      balancePartitionPercent: Float, numBatch: Int, batchSize: Int): DataFrame = {
     val pageRank = new EdgeCutPageRank()
       .setPartitionNum(partitionNum)
       .setStorageLevel(storageLevel)
@@ -108,6 +110,7 @@ object PageRankExample {
       .setUseBalancePartition(useBalancePartition)
       .setBalancePartitionPercent(balancePartitionPercent)
       .setNumBatch(numBatch)
+      .setBatchSize(batchSize)
     pageRank.transform(edges)
   }
 
diff --git a/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/edgecut/PageRank.scala b/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/edgecut/PageRank.scala
index 8edebffd1..57b1677e8 100644
--- a/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/edgecut/PageRank.scala
+++ b/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/edgecut/PageRank.scala
@@ -33,7 +33,7 @@ class PageRank(override val uid: String) extends Transformer
   with HasSrcNodeIdCol with HasDstNodeIdCol with HasOutputNodeIdCol with HasOutputPageRankCol
   with HasStorageLevel with HasPartitionNum with HasPSPartitionNum
   with HasWeightCol with HasIsWeighted with HasUseBalancePartition
-  with HasUseEstimatePartition with HasBalancePartitionPercent {
+  with HasUseEstimatePartition with HasBalancePartitionPercent with HasBatchSize {
 
 
   final val tol = new FloatParam(this, "tol", "tol")
@@ -117,9 +117,25 @@ class PageRank(override val uid: String) extends Transformer
 
     model.normalizeRanks(numNodes)
 
-    val (partitionIds, ends) = PageRankOps.splitPartitionIds(model.matrixId, graph.getNumPartitions)
-    val retRDD = graph.flatMap(f => PageRankOps.save(f.getIndex, model, partitionIds, ends, $(numBatch)))
-      .flatMap(f => f._1.zip(f._2))
+    val batchs = math.max($(psPartitionNum) / $(numBatch), $(batchSize))
+    val retRDD = edges.flatMap(f => Array(f._1, f._2))
+      .map(key=> (key, 1))
+      .groupByKey($(partitionNum))
+      .mapPartitions{ iter =>
+        val batchIterator = iter.map(_._1).sliding(batchs, batchs)
+
+        new Iterator[Array[(Long, Float)]] with Serializable {
+          override def hasNext: Boolean = {
+            batchIterator.hasNext
+          }
+
+          override def next: Array[(Long, Float)] = {
+            val batch = batchIterator.next().toArray
+            val ranks = model.readRanks(batch)
+            batch.map(key => (key, ranks.get(key)))
+          }
+        }
+      }.flatMap(p=>p)
       .map { case (node, rank) => Row.fromSeq(Seq[Any](node, rank)) }
 
     val outputSchema = transformSchema(dataset.schema)
diff --git a/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/vertexcut/PageRank.scala b/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/vertexcut/PageRank.scala
index 155f4de51..3fd2cd510 100644
--- a/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/vertexcut/PageRank.scala
+++ b/spark-on-angel/graph/src/main/scala/com/tencent/angel/graph/rank/pagerank/vertexcut/PageRank.scala
@@ -33,7 +33,7 @@ class PageRank(override val uid: String) extends Transformer
   with HasStorageLevel with HasPartitionNum with HasPSPartitionNum
   with HasWeightCol with HasIsWeighted with HasUseBalancePartition
   with HasUseEstimatePartition with HasBalancePartitionPercent
-  with HasApproNodeNum {
+  with HasApproNodeNum with HasBatchSize {
 
 
   final val tol = new FloatParam(this, "tol", "tol")
@@ -115,9 +115,25 @@ class PageRank(override val uid: String) extends Transformer
 
     model.normalizeRanks(numNodes)
 
-    val (partitionIds, ends) = PageRankOps.splitPartitionIds(model.matrixId, graph.getNumPartitions)
-    val retRDD = graph.flatMap(f => PageRankOps.save(f.getIndex, model, partitionIds, ends, $(numBatch)))
-      .flatMap(f => f._1.zip(f._2))
+    val batchs = math.max($(psPartitionNum) / $(numBatch), $(batchSize))
+    val retRDD = edges.flatMap(f => Array(f._1, f._2))
+      .map(key=> (key, 1))
+      .groupByKey($(partitionNum))
+      .mapPartitions{ iter =>
+        val batchIterator = iter.map(_._1).sliding(batchs, batchs)
+
+        new Iterator[Array[(Long, Float)]] with Serializable {
+          override def hasNext: Boolean = {
+            batchIterator.hasNext
+          }
+
+          override def next: Array[(Long, Float)] = {
+            val batch = batchIterator.next().toArray
+            val ranks = model.readRanks(batch)
+            batch.map(key => (key, ranks.get(key)))
+          }
+        }
+      }.flatMap(p=>p)
       .map { case (node, rank) => Row.fromSeq(Seq[Any](node, rank)) }
 
     val outputSchema = transformSchema(dataset.schema)


Test patch: