diff --git a/computer-algorithm/src/main/java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPath.java b/computer-algorithm/src/main/java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPath.java index c9b49cfde..b9b1f12b7 100644 --- a/computer-algorithm/src/main/java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPath.java +++ b/computer-algorithm/src/main/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 this.targetIdSet; + } } diff --git a/computer-algorithm/src/main/java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPathOutput.java b/computer-algorithm/src/main/java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPathOutput.java index 4a646cf41..943babb2e 100644 --- a/computer-algorithm/src/main/java/org/apache/hugegraph/computer/algorithm/path/shortest/SingleSourceShortestPathOutput.java +++ b/computer-algorithm/src/main/java/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 { @@ -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()); + } } diff --git a/computer-api/src/main/java/org/apache/hugegraph/computer/core/output/ComputerOutput.java b/computer-api/src/main/java/org/apache/hugegraph/computer/core/output/ComputerOutput.java index c2c2cc20d..3522853f9 100644 --- a/computer-api/src/main/java/org/apache/hugegraph/computer/core/output/ComputerOutput.java +++ b/computer-api/src/main/java/org/apache/hugegraph/computer/core/output/ComputerOutput.java @@ -19,6 +19,7 @@ import org.apache.hugegraph.computer.core.config.Config; import org.apache.hugegraph.computer.core.graph.vertex.Vertex; +import org.apache.hugegraph.computer.core.worker.Computation; /** * Computer output is used to output computer results. There is an output object @@ -37,6 +38,14 @@ public interface ComputerOutput { */ void write(Vertex vertex); + /** + * Write filter. + * True to commit the computation result, otherwise not to commit. + */ + default boolean filter(Config config, Computation computation, Vertex vertex) { + return true; + } + /** * Merge output files of multiple partitions if applicable. */ diff --git a/computer-core/src/main/java/org/apache/hugegraph/computer/core/compute/FileGraphPartition.java b/computer-core/src/main/java/org/apache/hugegraph/computer/core/compute/FileGraphPartition.java index 25312868e..a2e61a8a8 100644 --- a/computer-core/src/main/java/org/apache/hugegraph/computer/core/compute/FileGraphPartition.java +++ b/computer-core/src/main/java/org/apache/hugegraph/computer/core/compute/FileGraphPartition.java @@ -255,7 +255,9 @@ protected PartitionStat output() { Edges edges = this.edgesInput.edges(this.vertexInput.idPointer()); vertex.edges(edges); - output.write(vertex); + if (output.filter(this.context.config(), this.computation, vertex)) { + output.write(vertex); + } } try {