From b780d88213e0db682169ec37b5410b8e7ad25efa Mon Sep 17 00:00:00 2001
From: Pavel Ershov <owner.mad.epa@gmail.com>
Date: Mon, 28 Aug 2023 16:13:28 +0300
Subject: [PATCH] Fix ordered iterator

Signed-off-by: Pavel Ershov <owner.mad.epa@gmail.com>
(cherry picked from commit e1ebeb73cd71c2049c0ed1dd36887680d870943c)

# Conflicts:
#	janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctOrderedIterator.java
---
 .../graphdb/JanusGraphIndexTest.java          |  1 +
 .../util/MultiDistinctOrderedIterator.java    | 51 ++++++++++---------
 2 files changed, 29 insertions(+), 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
index 0969d06177..faa7db4e64 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
@@ -1228,6 +1228,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());
         Metrics mMixedOr = tx.traversal().V().or(__.has("name", "bob"), __.has("age", 20))
             .profile().next().getMetrics(0);
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 ec60ddd4bf..ff6990e3e8 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,33 +23,37 @@
 
 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 long count = 0;
 
     public MultiDistinctOrderedIterator(final Integer lowLimit, final Integer highLimit, final List<Iterator<E>> iterators, final List<OrderEntry> orders) {
         this.limit = highLimit;
-        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();
@@ -58,12 +62,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();
@@ -72,24 +76,25 @@ 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);
                     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);
     }
-
 }