From e2b1a0f4db922aec5e07c4219a987d54f321f3f6 Mon Sep 17 00:00:00 2001 From: Jermy Li Date: Wed, 15 Apr 2020 16:02:44 +0800 Subject: [PATCH] add direction and label parameter for centrality algorithms (#13) Change-Id: I20b72ea0da673359e2bd21888010290efca81441 --- .../job/algorithm/AbstractAlgorithm.java | 10 ++++- .../algorithm/cent/AbstractCentAlgorithm.java | 25 +++++++++-- .../cent/BetweenessCentralityAlgorithm.java | 12 ++++-- .../cent/ClosenessCentralityAlgorithm.java | 12 ++++-- .../cent/DegreeCentralityAlgorithm.java | 43 ++++++++++++++----- .../cent/EigenvectorCentralityAlgorithm.java | 12 ++++-- 6 files changed, 89 insertions(+), 25 deletions(-) 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 e77473668c..248a92bdb1 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 @@ -49,6 +49,7 @@ import com.baidu.hugegraph.type.define.Directions; import com.baidu.hugegraph.type.define.HugeKeys; import com.baidu.hugegraph.util.Bytes; +import com.baidu.hugegraph.util.CollectionUtil; import com.baidu.hugegraph.util.E; import com.baidu.hugegraph.util.JsonUtil; @@ -119,6 +120,9 @@ protected static String edgeLabel(Map parameters) { } protected static Directions direction(Map parameters) { + if (!parameters.containsKey(KEY_DIRECTION)) { + return Directions.BOTH; + } Object direction = parameter(parameters, KEY_DIRECTION); return parseDirection(direction); } @@ -437,7 +441,11 @@ public void put(K key, long value) { } public Set> entrySet() { - this.shrinkIfNeeded(this.topN); + if (this.tops.size() <= this.topN) { + this.tops = CollectionUtil.sortByValue(this.tops, false); + } else { + this.shrinkIfNeeded(this.topN); + } return this.tops.entrySet(); } 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 c36743176c..fb0c33d503 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 @@ -27,12 +27,14 @@ 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.Direction; import org.apache.tinkerpop.gremlin.structure.Vertex; import com.baidu.hugegraph.backend.id.Id; import com.baidu.hugegraph.job.Job; import com.baidu.hugegraph.job.algorithm.AbstractAlgorithm; import com.baidu.hugegraph.structure.HugeElement; +import com.baidu.hugegraph.type.define.Directions; public abstract class AbstractCentAlgorithm extends AbstractAlgorithm { @@ -46,6 +48,8 @@ public void checkParameters(Map parameters) { depth(parameters); degree(parameters); sample(parameters); + direction(parameters); + edgeLabel(parameters); sourceSample(parameters); sourceLabel(parameters); sourceCLabel(parameters); @@ -83,9 +87,11 @@ protected GraphTraversal constructSource( } protected GraphTraversal constructPath( - GraphTraversal t, long degree, - long sample, String sourceLabel, String sourceCLabel) { - GraphTraversal unit = constructPathUnit(degree, sample, + GraphTraversal t, Directions dir, + String label, long degree, long sample, + String sourceLabel, String sourceCLabel) { + GraphTraversal unit = constructPathUnit(dir, label, + degree, sample, sourceLabel, sourceCLabel); t = t.as("v").repeat(__.local(unit).simplePath().as("v")); @@ -94,10 +100,21 @@ protected GraphTraversal constructPath( } protected GraphTraversal constructPathUnit( + Directions dir, String label, long degree, long sample, String sourceLabel, String sourceCLabel) { - GraphTraversal unit = __.both(); + if (dir == null) { + dir = Directions.BOTH; + } + Direction direction = dir.direction(); + + String[] labels = {}; + if (label != null) { + labels = new String[]{label}; + } + + GraphTraversal unit = __.to(direction, labels); if (sourceLabel != null) { unit = unit.hasLabel(sourceLabel); } 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 9a72d2f626..12e3acba02 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 @@ -31,6 +31,7 @@ import org.apache.tinkerpop.gremlin.structure.Vertex; import com.baidu.hugegraph.job.Job; +import com.baidu.hugegraph.type.define.Directions; public class BetweenessCentralityAlgorithm extends AbstractCentAlgorithm { @@ -42,7 +43,9 @@ public String name() { @Override public Object call(Job job, Map parameters) { Traverser traverser = new Traverser(job); - return traverser.betweenessCentrality(depth(parameters), + return traverser.betweenessCentrality(direction(parameters), + edgeLabel(parameters), + depth(parameters), degree(parameters), sample(parameters), sourceLabel(parameters), @@ -57,7 +60,9 @@ public Traverser(Job job) { super(job); } - public Object betweenessCentrality(int depth, + public Object betweenessCentrality(Directions direction, + String label, + int depth, long degree, long sample, String sourceLabel, @@ -71,7 +76,8 @@ public Object betweenessCentrality(int depth, GraphTraversal t = constructSource(sourceLabel, sourceSample, sourceCLabel); - t = constructPath(t, degree, sample, sourceLabel, sourceCLabel); + t = constructPath(t, direction, label, degree, sample, + sourceLabel, sourceCLabel); t = t.emit().until(__.loops().is(P.gte(depth))); t = filterNonShortestPath(t); diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/ClosenessCentralityAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/ClosenessCentralityAlgorithm.java index 96e9709fef..cb64bd8bc5 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/ClosenessCentralityAlgorithm.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/ClosenessCentralityAlgorithm.java @@ -32,6 +32,7 @@ import org.apache.tinkerpop.gremlin.structure.Vertex; import com.baidu.hugegraph.job.Job; +import com.baidu.hugegraph.type.define.Directions; public class ClosenessCentralityAlgorithm extends AbstractCentAlgorithm { @@ -51,7 +52,9 @@ public void checkParameters(Map parameters) { @Override public Object call(Job job, Map parameters) { Traverser traverser = new Traverser(job); - return traverser.closenessCentrality(depth(parameters), + return traverser.closenessCentrality(direction(parameters), + edgeLabel(parameters), + depth(parameters), degree(parameters), sample(parameters), sourceLabel(parameters), @@ -66,7 +69,9 @@ public Traverser(Job job) { super(job); } - public Object closenessCentrality(int depth, + public Object closenessCentrality(Directions direction, + String label, + int depth, long degree, long sample, String sourceLabel, @@ -80,7 +85,8 @@ public Object closenessCentrality(int depth, GraphTraversal t = constructSource(sourceLabel, sourceSample, sourceCLabel); - t = constructPath(t, degree, sample, sourceLabel, sourceCLabel); + t = constructPath(t, direction, label, degree, sample, + sourceLabel, sourceCLabel); t = t.emit().until(__.loops().is(P.gte(depth))); t = filterNonShortestPath(t); diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/DegreeCentralityAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/DegreeCentralityAlgorithm.java index 5f6781b21b..a19c098229 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/DegreeCentralityAlgorithm.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/DegreeCentralityAlgorithm.java @@ -22,9 +22,9 @@ import java.util.Iterator; import java.util.Map; -import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource; import org.apache.tinkerpop.gremlin.structure.Edge; import org.apache.tinkerpop.gremlin.structure.Vertex; +import org.apache.tinkerpop.gremlin.util.iterator.IteratorUtils; import com.baidu.hugegraph.backend.id.Id; import com.baidu.hugegraph.job.Job; @@ -41,6 +41,7 @@ public String name() { @Override public void checkParameters(Map parameters) { direction(parameters); + edgeLabel(parameters); top(parameters); } @@ -48,6 +49,7 @@ public void checkParameters(Map parameters) { public Object call(Job job, Map parameters) { Traverser traverser = new Traverser(job); return traverser.degreeCentrality(direction(parameters), + edgeLabel(parameters), top(parameters)); } @@ -57,9 +59,11 @@ public Traverser(Job job) { super(job); } - public Object degreeCentrality(Directions direction, long topN) { + public Object degreeCentrality(Directions direction, + String label, + long topN) { if (direction == null || direction == Directions.BOTH) { - return degreeCentrality(topN); + return degreeCentrality(label, topN); } assert direction == Directions.OUT || direction == Directions.IN; assert topN >= 0L; @@ -69,6 +73,7 @@ public Object degreeCentrality(Directions direction, long topN) { JsonMap degrees = new JsonMap(); TopMap tops = new TopMap<>(topN); Id vertex = null; + Id labelId = this.getEdgeLabelId(label); long degree = 0L; long total = 0L; @@ -77,12 +82,20 @@ public Object degreeCentrality(Directions direction, long topN) { HugeEdge edge = (HugeEdge) edges.next(); this.updateProgress(++total); + Id schemaLabel = edge.schemaLabel().id(); + if (labelId != null && !labelId.equals(schemaLabel)) { + continue; + } + Id source = edge.ownerVertex().id(); if (source.equals(vertex)) { + // edges belong to same source vertex degree++; continue; } + if (vertex != null) { + // next vertex found if (topN <= 0L) { degrees.append(vertex, degree); } else { @@ -107,25 +120,26 @@ public Object degreeCentrality(Directions direction, long topN) { return degrees.asJson(); } - protected Object degreeCentrality(long topN) { + protected Object degreeCentrality(String label, long topN) { assert topN >= 0L; long total = 0L; JsonMap degrees = new JsonMap(); TopMap tops = new TopMap<>(topN); - GraphTraversalSource traversal = this.graph().traversal(); Iterator vertices = this.vertices(); degrees.startObject(); while (vertices.hasNext()) { - Vertex source = vertices.next(); + Id source = (Id) vertices.next().id(); this.updateProgress(++total); - Long degree = traversal.V(source).bothE().count().next(); - if (topN <= 0L) { - degrees.append(source.id(), degree); - } else { - tops.put((Id) source.id(), degree); + long degree = this.degree(source, label); + if (degree > 0L) { + if (topN <= 0L) { + degrees.append(source, degree); + } else { + tops.put(source, degree); + } } } @@ -136,5 +150,12 @@ protected Object degreeCentrality(long topN) { return degrees.asJson(); } + + private long degree(Id source, String label) { + Id labelId = this.getEdgeLabelId(label); + Iterator edges = this.edgesOfVertex(source, Directions.BOTH, + labelId, NO_LIMIT); + return IteratorUtils.count(edges); + } } } diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/EigenvectorCentralityAlgorithm.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/EigenvectorCentralityAlgorithm.java index d87fc79316..ce47417c4d 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/EigenvectorCentralityAlgorithm.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/job/algorithm/cent/EigenvectorCentralityAlgorithm.java @@ -30,6 +30,7 @@ import org.apache.tinkerpop.gremlin.structure.Vertex; import com.baidu.hugegraph.job.Job; +import com.baidu.hugegraph.type.define.Directions; public class EigenvectorCentralityAlgorithm extends AbstractCentAlgorithm { @@ -44,7 +45,9 @@ public String name() { @Override public Object call(Job job, Map parameters) { Traverser traverser = new Traverser(job); - return traverser.eigenvectorCentrality(depth(parameters), + return traverser.eigenvectorCentrality(direction(parameters), + edgeLabel(parameters), + depth(parameters), degree(parameters), sample(parameters), sourceLabel(parameters), @@ -59,7 +62,9 @@ public Traverser(Job job) { super(job); } - public Object eigenvectorCentrality(int depth, + public Object eigenvectorCentrality(Directions direction, + String label, + int depth, long degree, long sample, String sourceLabel, @@ -83,7 +88,8 @@ public Object eigenvectorCentrality(int depth, GraphTraversal t = constructSource(sourceLabel, sourceSample, sourceCLabel); - GraphTraversal unit = constructPathUnit(degree, sample, + GraphTraversal unit = constructPathUnit(direction, label, + degree, sample, sourceLabel, sourceCLabel); t = t.repeat(__.groupCount("m").by(T.id)