diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/AbstractAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/AbstractAlgorithm.java index ffe55e8f6f..064b1e344c 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/AbstractAlgorithm.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/AbstractAlgorithm.java @@ -23,6 +23,7 @@ import java.util.Iterator; import java.util.Map; import java.util.Map.Entry; +import java.util.NoSuchElementException; import java.util.Set; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; @@ -454,6 +455,16 @@ protected V execute(GraphTraversal traversal, } } + protected Number tryNext(GraphTraversal iter) { + return this.execute(iter, () -> { + try { + return iter.next(); + } catch (NoSuchElementException e) { + return 0; + } + }); + } + protected void commitIfNeeded() { // commit if needed Transaction tx = this.graph().tx(); diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/AbstractCentAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/AbstractCentAlgorithm.java index 19da8e968e..fd8453fff3 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/AbstractCentAlgorithm.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/AbstractCentAlgorithm.java @@ -19,7 +19,9 @@ package com.baidu.hugegraph.job.algorithm.cent; +import java.util.Collections; import java.util.HashMap; +import java.util.Iterator; import java.util.List; import java.util.Map; @@ -32,15 +34,21 @@ import org.apache.tinkerpop.gremlin.structure.Column; import org.apache.tinkerpop.gremlin.structure.Direction; import org.apache.tinkerpop.gremlin.structure.Vertex; +import org.slf4j.Logger; import com.baidu.hugegraph.backend.id.Id; +import com.baidu.hugegraph.iterator.MapperIterator; import com.baidu.hugegraph.job.UserJob; import com.baidu.hugegraph.job.algorithm.AbstractAlgorithm; import com.baidu.hugegraph.structure.HugeElement; +import com.baidu.hugegraph.structure.HugeVertex; import com.baidu.hugegraph.type.define.Directions; +import com.baidu.hugegraph.util.Log; public abstract class AbstractCentAlgorithm extends AbstractAlgorithm { + private static final Logger LOG = Log.logger(AbstractCentAlgorithm.class); + @Override public String category() { return CATEGORY_CENT; @@ -161,6 +169,33 @@ protected GraphTraversal filterNonShortestPath( }); } + protected GraphTraversal substractPath( + GraphTraversal t, + boolean withBoundary) { + // t.select(Pop.all, "v").unfold().id() + return t.select(Pop.all, "v").flatMap(it -> { + List path = (List) it.get(); + if (withBoundary) { + @SuppressWarnings("unchecked") + Iterator items = (Iterator) + path.iterator(); + return new MapperIterator<>(items, v -> v.id()); + } + int len = path.size(); + if (len < 3) { + return Collections.emptyIterator(); + } + + LOG.debug("CentAlgorithm substract path: {}", path); + path.remove(path.size() -1); + path.remove(0); + @SuppressWarnings("unchecked") + Iterator items = (Iterator) + path.iterator(); + return new MapperIterator<>(items, v -> v.id()); + }); + } + protected GraphTraversal topN(GraphTraversal t, long topN) { if (topN > 0L || topN == NO_LIMIT) { diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/BetweenessCentralityAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/BetweenessCentralityAlgorithm.java index 465c6f96c9..968f636bae 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/BetweenessCentralityAlgorithm.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/BetweenessCentralityAlgorithm.java @@ -22,21 +22,29 @@ import java.util.Map; import org.apache.tinkerpop.gremlin.process.traversal.P; -import org.apache.tinkerpop.gremlin.process.traversal.Pop; import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal; import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.__; import org.apache.tinkerpop.gremlin.structure.Vertex; import com.baidu.hugegraph.job.UserJob; import com.baidu.hugegraph.type.define.Directions; +import com.baidu.hugegraph.util.ParameterUtil; public class BetweenessCentralityAlgorithm extends AbstractCentAlgorithm { + public static final String KEY_WITH_BOUNDARY = "with_boundary"; + @Override public String name() { return "betweeness_centrality"; } + @Override + public void checkParameters(Map parameters) { + super.checkParameters(parameters); + withBoundary(parameters); + } + @Override public Object call(UserJob job, Map parameters) { try (Traverser traverser = new Traverser(job)) { @@ -45,6 +53,7 @@ public Object call(UserJob job, Map parameters) { depth(parameters), degree(parameters), sample(parameters), + withBoundary(parameters), sourceLabel(parameters), sourceSample(parameters), sourceCLabel(parameters), @@ -52,6 +61,13 @@ public Object call(UserJob job, Map parameters) { } } + protected static boolean withBoundary(Map parameters) { + if (!parameters.containsKey(KEY_WITH_BOUNDARY)) { + return false; + } + return ParameterUtil.parameterBoolean(parameters, KEY_WITH_BOUNDARY); + } + private static class Traverser extends AbstractCentAlgorithm.Traverser { public Traverser(UserJob job) { @@ -63,6 +79,7 @@ public Object betweenessCentrality(Directions direction, int depth, long degree, long sample, + boolean withBoundary, String sourceLabel, long sourceSample, String sourceCLabel, @@ -79,8 +96,8 @@ public Object betweenessCentrality(Directions direction, t = t.emit().until(__.loops().is(P.gte(depth))); t = filterNonShortestPath(t, false); - GraphTraversal tg = t.select(Pop.all, "v") - .unfold().id().groupCount(); + GraphTraversal tg = this.substractPath(t, withBoundary) + .groupCount(); GraphTraversal tLimit = topN(tg, topN); return this.execute(tLimit, () -> tLimit.next()); diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LouvainTraverser.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LouvainTraverser.java index c6cd24fab6..42ac3e9906 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LouvainTraverser.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LouvainTraverser.java @@ -28,7 +28,6 @@ import java.util.List; import java.util.Map; import java.util.Map.Entry; -import java.util.NoSuchElementException; import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -649,16 +648,6 @@ private double modularity(String label) { return q; } - private Number tryNext(GraphTraversal iter) { - return this.execute(iter, () -> { - try { - return iter.next(); - } catch (NoSuchElementException e) { - return 0; - } - }); - } - public Collection showCommunity(String community) { final String C_PASS0 = labelOfPassN(0); Collection comms = Arrays.asList(community); diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LpaAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LpaAlgorithm.java index 8b54241fea..973e93f7b2 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LpaAlgorithm.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/comm/LpaAlgorithm.java @@ -117,9 +117,9 @@ public Object lpa(String sourceLabel, String edgeLabel, } } - long communities = this.graph().traversal().V().limit(100000L) - .groupCount().by(C_LABEL) - .count(Scope.local).next(); + Number communities = tryNext(this.graph().traversal().V() + .groupCount().by(C_LABEL) + .count(Scope.local)); return ImmutableMap.of("iteration_times", times, "last_precision", changedPercent, "times", maxTimes,