Skip to content

Commit

Permalink
refact: let kout degree applied to all labels & remove source vertex …
Browse files Browse the repository at this point in the history
…from kneighbor default results (apache#1459)
  • Loading branch information
imbajin authored May 21, 2021
1 parent 0fbe4d1 commit 718daa6
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 47 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1995,47 +1996,4 @@ private <T> void traverseByLabel(SchemaLabel label,
} while (page != null);
}
}

// TODO: move to common module
public static class LimitIterator<T> extends WrappedIterator<T> {

private final Iterator<T> originIterator;
private final Function<T, Boolean> filterCallback;

public LimitIterator(Iterator<T> origin, Function<T, Boolean> filter) {
this.originIterator = origin;
this.filterCallback = filter;
}

@Override
protected Iterator<T> 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);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -148,7 +149,8 @@ protected Set<Node> adjacentVertices(Set<Node> 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);
Expand Down Expand Up @@ -182,10 +184,17 @@ protected Iterator<Edge> edgesOfVertex(Id source, Directions dir,
ExtendableIterator<Edge> 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<Edge> edgesOfVertex(Id source, EdgeStep edgeStep) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,6 @@ public Set<Node> 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();
Expand Down

0 comments on commit 718daa6

Please sign in to comment.