From 32bacddbd83d2d6a746e1575b544da8477d0e2ad Mon Sep 17 00:00:00 2001
From: houzhizhen <houzhizhen_101@163.com>
Date: Tue, 11 Aug 2020 08:54:18 +0800
Subject: [PATCH] add parameter top to print the top result in job result (#32)

* add parameter top to print the top result in job result
---
 .../job/algorithm/rank/PageRankAlgorithm.java | 50 +++++++++++++++++--
 1 file changed, 45 insertions(+), 5 deletions(-)

diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/rank/PageRankAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/rank/PageRankAlgorithm.java
index fa552e3997..651aa91259 100644
--- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/rank/PageRankAlgorithm.java
+++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/rank/PageRankAlgorithm.java
@@ -19,6 +19,8 @@
 
 package com.baidu.hugegraph.job.algorithm.rank;
 
+import com.baidu.hugegraph.traversal.algorithm.HugeTraverser;
+
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -61,6 +63,7 @@ public void checkParameters(Map<String, Object> parameters) {
         precision(parameters);
         degree(parameters);
         directionOutIn(parameters);
+        top(parameters);
     }
 
     @Override
@@ -70,7 +73,8 @@ public Object call(UserJob<Object> job, Map<String, Object> parameters) {
                                       times(parameters),
                                       precision(parameters),
                                       degree(parameters),
-                                      directionOutIn(parameters));
+                                      directionOutIn(parameters),
+                                      top(parameters));
         } catch (Throwable e) {
             job.graph().tx().rollback();
             throw e;
@@ -90,11 +94,15 @@ public Traverser(UserJob<Object> job) {
             this.vertexRankMap = new HashMap<>();
         }
 
-        public Object pageRank(double alpha,
+        /**
+         * If topN > 0, then return topN elements with rank value in json.
+         */
+        private Object pageRank(double alpha,
                                int maxTimes,
                                double precision,
                                long degree,
-                               Directions direction) {
+                               Directions direction,
+                               long topN) {
             this.initSchema();
 
             int times;
@@ -146,12 +154,32 @@ public Object pageRank(double alpha,
 
             this.writeBackRankValues();
 
+            if (topN > 0) {
+                Object topNJson = this.getTopRank(topN);
+                return ImmutableMap.of("alpha", alpha,
+                        "iteration_times", times,
+                        "last_changed_rank", changedRank,
+                        "times", maxTimes,
+                        "top", topNJson);
+            }
             return ImmutableMap.of("alpha", alpha,
                                    "iteration_times", times,
                                    "last_changed_rank", changedRank,
                                    "times", maxTimes);
         }
 
+        private Object getTopRank(long topN) {
+            JsonMap jsonMap = new JsonMap();
+            jsonMap.startObject();
+            Map<Id, DoublePair> topNMap =
+                    HugeTraverser.topN(this.vertexRankMap, true, topN);
+            for (Map.Entry<Id, DoublePair> e : topNMap.entrySet()) {
+                jsonMap.append(e.getKey().toString(), e.getValue().left);
+            }
+            jsonMap.endObject();
+            return jsonMap.asJson();
+        }
+
         private long initRankMap() {
             long vertexCount = 0;
             Iterator<Vertex> vertices = this.vertices();
@@ -239,12 +267,12 @@ private double computeRank(double alpha, long numOfVertices) {
         }
     }
 
-    public static class DoublePair {
+    public static class DoublePair implements Comparable<DoublePair> {
 
         private double left;
         private double right;
 
-        public DoublePair(double left, double right) {
+        private DoublePair(double left, double right) {
             this.left = left;
             this.right = right;
         }
@@ -294,5 +322,17 @@ public boolean equals(Object obj) {
         public int hashCode() {
             return Double.hashCode(this.left) ^ Double.hashCode(this.right);
         }
+
+        // only left saves the rank value.
+        @Override
+        public int compareTo(DoublePair o) {
+            double result = this.left - o.left;
+            if (result > 0.0) {
+                return 1;
+            } else if (result < 0.0) {
+                return -1;
+            }
+            return 0;
+        }
     }
 }