improvement: sssp algorithm output filter

apache · Mar 13, 2024 · 724fbd1 · 724fbd1
1 parent e295b4d
commit 724fbd1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/.../java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPath.java b/.../java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPath.java
@@ -278,4 +278,8 @@ private boolean isAllTargetsReached(Vertex vertex) {
         }
         return false;
     }
+
+    public IdSet getTargetIdSet() {
+        return targetIdSet;
+    }
 }
diff --git a/...org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPathOutput.java b/...org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPathOutput.java
@@ -20,8 +20,10 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import org.apache.hugegraph.computer.core.config.Config;
 import org.apache.hugegraph.computer.core.graph.vertex.Vertex;
 import org.apache.hugegraph.computer.core.output.hg.HugeGraphOutput;
+import org.apache.hugegraph.computer.core.worker.Computation;
 import org.apache.hugegraph.util.JsonUtil;
 
 public class SingleSourceShortestPathOutput extends HugeGraphOutput<String> {
@@ -45,4 +47,10 @@ protected String value(Vertex vertex) {
         map.put("total_weight", value.totalWeight());
         return JsonUtil.toJson(map);
     }
+
+    @Override
+    public boolean filter(Config config, Computation computation, Vertex vertex) {
+        SingleSourceShortestPath sssp = (SingleSourceShortestPath) computation;
+        return sssp.getTargetIdSet() == null || sssp.getTargetIdSet().contains(vertex.id());
+    }
 }