diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/tx/GraphTransaction.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/tx/GraphTransaction.java index 6d3f8de109..a2af8b0d3d 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/tx/GraphTransaction.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/tx/GraphTransaction.java @@ -72,6 +72,7 @@ import com.baidu.hugegraph.iterator.ExtendableIterator; import com.baidu.hugegraph.iterator.FilterIterator; import com.baidu.hugegraph.iterator.FlatMapperIterator; +import com.baidu.hugegraph.iterator.LimitIterator; import com.baidu.hugegraph.iterator.ListIterator; import com.baidu.hugegraph.iterator.MapperIterator; import com.baidu.hugegraph.iterator.WrappedIterator; @@ -1995,47 +1996,4 @@ private void traverseByLabel(SchemaLabel label, } while (page != null); } } - - // TODO: move to common module - public static class LimitIterator extends WrappedIterator { - - private final Iterator originIterator; - private final Function filterCallback; - - public LimitIterator(Iterator origin, Function filter) { - this.originIterator = origin; - this.filterCallback = filter; - } - - @Override - protected Iterator originIterator() { - return this.originIterator; - } - - @Override - protected final boolean fetch() { - while (this.originIterator.hasNext()) { - T next = this.originIterator.next(); - // Do filter - boolean reachLimit = this.filterCallback.apply(next); - if (reachLimit) { - this.closeOriginIterator(); - return false; - } - if (next != null) { - assert this.current == none(); - this.current = next; - return true; - } - } - return false; - } - - protected final void closeOriginIterator() { - if (this.originIterator == null) { - return; - } - close(this.originIterator); - } - } } diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/HugeTraverser.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/HugeTraverser.java index ef7f569523..ec1fb3ad1b 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/HugeTraverser.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/HugeTraverser.java @@ -48,6 +48,7 @@ import com.baidu.hugegraph.exception.NotFoundException; import com.baidu.hugegraph.iterator.ExtendableIterator; import com.baidu.hugegraph.iterator.FilterIterator; +import com.baidu.hugegraph.iterator.LimitIterator; import com.baidu.hugegraph.iterator.MapperIterator; import com.baidu.hugegraph.schema.SchemaLabel; import com.baidu.hugegraph.structure.HugeEdge; @@ -148,7 +149,8 @@ protected Set adjacentVertices(Set vertices, EdgeStep step, while (edges.hasNext()) { Id target = ((HugeEdge) edges.next()).id().otherVertexId(); KNode kNode = new KNode(target, (KNode) source); - if (excluded != null && excluded.contains(kNode)) { + if ((excluded != null && excluded.contains(kNode)) || + neighbors.contains(kNode)) { continue; } neighbors.add(kNode); @@ -182,10 +184,17 @@ protected Iterator edgesOfVertex(Id source, Directions dir, ExtendableIterator results = new ExtendableIterator<>(); for (Id label : labels.keySet()) { E.checkNotNull(label, "edge label"); - // TODO: limit should be applied to all labels results.extend(this.edgesOfVertex(source, dir, label, limit)); } - return results; + + if (limit == NO_LIMIT) { + return results; + } + + long[] count = new long[1]; + return new LimitIterator<>(results, e -> { + return count[0]++ >= limit; + }); } protected Iterator edgesOfVertex(Id source, EdgeStep edgeStep) { diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/KneighborTraverser.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/KneighborTraverser.java index 1fadd50147..bd306bf4df 100644 --- a/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/KneighborTraverser.java +++ b/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/KneighborTraverser.java @@ -86,7 +86,6 @@ public Set customizedKneighbor(Id source, EdgeStep step, int maxDepth, Node sourceV = new KNode(source, null); latest.add(sourceV); - all.add(sourceV); while (maxDepth-- > 0) { long remaining = limit == NO_LIMIT ? NO_LIMIT : limit - all.size();