From 91308f65cbd8506c4e6dc3a627e78c307c89f74c Mon Sep 17 00:00:00 2001
From: Pavel Ershov <owner.mad.epa@gmail.com>
Date: Thu, 15 Apr 2021 17:53:35 +0300
Subject: [PATCH] Do not store distinct elements when single iterator processed

Signed-off-by: Pavel Ershov <owner.mad.epa@gmail.com>
---
 .../graphdb/util/MultiDistinctOrderedIterator.java          | 6 +++++-
 .../graphdb/util/MultiDistinctUnorderedIterator.java        | 4 +++-
 2 files changed, 8 insertions(+), 2 deletions(-)

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 af2e0dede0..f3b9d00f24 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
@@ -38,10 +38,12 @@ public class MultiDistinctOrderedIterator<E extends Element> implements Closeabl
     private final TreeMap<E, Integer> currentElements;
     private final Set<Object> allElements = new HashSet<>();
     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);
@@ -73,7 +75,9 @@ public boolean hasNext() {
                 if (element != null) {
                     values.put(i, element);
                     currentElements.put(element, i);
-                    allElements.add(element.id());
+                    if (!singleIterator) {
+                        allElements.add(element.id());
+                    }
                 }
             }
         }
diff --git a/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctUnorderedIterator.java b/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctUnorderedIterator.java
index f9d0f28a0d..e58079cbf5 100644
--- a/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctUnorderedIterator.java
+++ b/janusgraph-core/src/main/java/org/janusgraph/graphdb/util/MultiDistinctUnorderedIterator.java
@@ -28,12 +28,14 @@ public class MultiDistinctUnorderedIterator<E extends Element> extends Closeable
     private final Set<Object> allElements = new HashSet<>();
     private final CloseableIterator<E> iterator;
     private final int limit;
+    private final boolean singleIterator;
     private long count;
 
     public MultiDistinctUnorderedIterator(final int lowLimit, final int highLimit, final List<Iterator<E>> iterators) {
         Objects.requireNonNull(iterators);
         iterator = CloseableIteratorUtils.concat(iterators);
         limit = highLimit;
+        singleIterator = iterators.size() == 1;
 
         long i = 0;
         while (i < lowLimit && hasNext()) {
@@ -47,7 +49,7 @@ protected E computeNext() {
         if (count < limit) {
             while (iterator.hasNext()) {
                 E elem = iterator.next();
-                if (allElements.add(elem.id())) {
+                if (singleIterator || allElements.add(elem.id())) {
                     count++;
                     return elem;
                 }