Fix ordered iterator

Signed-off-by: Pavel Ershov <owner.mad.epa@gmail.com>
JanusGraph · Sep 12, 2023 · e1ebeb7 · e1ebeb7
1 parent 0663c29
commit e1ebeb7
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 23 deletions.
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
@@ -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
@@ -23,35 +23,39 @@
 
 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<E extends Element> extends CloseableAbstractIterator<E> {
 
-public class MultiDistinctOrderedIterator<E extends Element> implements CloseableIterator<E> {
-
-    private final Map<Integer, Iterator<E>> iterators = new LinkedHashMap<>();
-    private final Map<Integer, E> values = new LinkedHashMap<>();
-    private final TreeMap<E, Integer> currentElements;
+    private final Map<Integer, Iterator<E>> iterators;
+    private final PriorityQueue<E> queue;
     private final Set<Object> allElements = new HashSet<>();
+    private final List<E> iteratorValues;
+    private final Map<E, Integer> iteratorIdx;
     private final Integer limit;
     private final boolean singleIterator;
     private long count = 0;
 
     public MultiDistinctOrderedIterator(final Integer lowLimit, final Integer highLimit, final List<Iterator<E>> iterators, final List<OrderEntry> orders) {
         this.limit = highLimit;
         this.singleIterator = iterators.size() == 1;
-        final List<Comparator<E>> comp = new ArrayList<>();
-        orders.forEach(o -> comp.add(new ElementValueComparator(o.key, o.order)));
-        Comparator<E> 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<Comparator<E>> 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);
     }
-
 }