diff --git a/janusgraph-backend-testutils/src/main/java/org/janusgraph/graphdb/JanusGraphIndexTest.java b/janusgraph-backend-testutils/src/main/java/org/janusgraph/graphdb/JanusGraphIndexTest.java index 3247440c26..7f32787df4 100644 --- a/janusgraph-backend-testutils/src/main/java/org/janusgraph/graphdb/JanusGraphIndexTest.java +++ b/janusgraph-backend-testutils/src/main/java/org/janusgraph/graphdb/JanusGraphIndexTest.java @@ -1255,6 +1255,7 @@ public void testGraphCentricQueryProfiling() { // satisfied by two graph-centric queries, one is mixed index query and the other one is full scan newTx(); assertEquals(3, tx.traversal().V().or(__.has("name", "bob"), __.has("age", 20)).count().next()); + assertEquals(3, tx.traversal().V().or(__.has("name", "bob"), __.has("age", 20)).order().by("age").count().next()); assertEquals(3, tx.traversal().V().or(__.has("name", "bob"), __.has("age", 20)).toList().size()); assertEquals(190, tx.traversal().V().or(__.has("name", "bob"), __.has("age", 20)).values("height").max().next()); assertEquals(170, tx.traversal().V().or(__.has("name", "bob"), __.has("age", 20)).values("height").min().next()); diff --git a/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctOrderedIterator.java b/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctOrderedIterator.java index 0717b2c20f..dc4aae455b 100644 --- a/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctOrderedIterator.java +++ b/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctOrderedIterator.java @@ -23,21 +23,22 @@ import java.util.ArrayList; import java.util.Comparator; +import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.PriorityQueue; import java.util.Set; -import java.util.TreeMap; +public class MultiDistinctOrderedIterator extends CloseableAbstractIterator { -public class MultiDistinctOrderedIterator implements CloseableIterator { - - private final Map> iterators = new LinkedHashMap<>(); - private final Map values = new LinkedHashMap<>(); - private final TreeMap currentElements; + private final Map> iterators; + private final PriorityQueue queue; private final Set allElements = new HashSet<>(); + private final List iteratorValues; + private final Map iteratorIdx; private final Integer limit; private final boolean singleIterator; private long count = 0; @@ -45,13 +46,16 @@ public class MultiDistinctOrderedIterator implements Closeabl public MultiDistinctOrderedIterator(final Integer lowLimit, final Integer highLimit, final List> iterators, final List orders) { this.limit = highLimit; this.singleIterator = iterators.size() == 1; - final List> comp = new ArrayList<>(); - orders.forEach(o -> comp.add(new ElementValueComparator(o.key, o.order))); - Comparator comparator = new MultiComparator<>(comp); + this.iterators = new LinkedHashMap<>(iterators.size()); + this.iteratorValues = new ArrayList<>(iterators.size()); + this.iteratorIdx = new HashMap<>(iterators.size()); for (int i = 0; i < iterators.size(); i++) { this.iterators.put(i, iterators.get(i)); + this.iteratorValues.add(null); } - currentElements = new TreeMap<>(comparator); + final List> comp = new ArrayList<>(); + orders.forEach(o -> comp.add(new ElementValueComparator(o.key, o.order))); + this.queue = new PriorityQueue<>(iterators.size(), new MultiComparator<>(comp)); long i = 0; while (i < lowLimit && this.hasNext()) { this.next(); @@ -60,12 +64,12 @@ public MultiDistinctOrderedIterator(final Integer lowLimit, final Integer highLi } @Override - public boolean hasNext() { + protected E computeNext() { if (limit != null && limit != Query.NO_LIMIT && count >= limit) { - return false; + return endOfData(); } for (int i = 0; i < iterators.size(); i++) { - if (!values.containsKey(i) && iterators.get(i).hasNext()){ + if (iteratorValues.get(i) == null && iterators.get(i).hasNext()) { E element = null; do { element = iterators.get(i).next(); @@ -74,26 +78,27 @@ public boolean hasNext() { } } while (element == null && iterators.get(i).hasNext()); if (element != null) { - values.put(i, element); - currentElements.put(element, i); + iteratorValues.set(i, element); + iteratorIdx.put(element, i); + queue.add(element); if (!singleIterator) { allElements.add(element.id()); } } } } - return !values.isEmpty(); - } - - @Override - public E next() { - count++; - return values.remove(currentElements.remove(currentElements.firstKey())); + if (queue.isEmpty()) { + return endOfData(); + } else { + count++; + E result = queue.remove(); + iteratorValues.set(iteratorIdx.remove(result), null); + return result; + } } @Override public void close() { iterators.values().forEach(CloseableIterator::closeIterator); } - }