diff --git a/lucene/core/src/java/org/apache/lucene/search/ExactPhraseMatcher.java b/lucene/core/src/java/org/apache/lucene/search/ExactPhraseMatcher.java index ca125e7f0f95..3d911ff5afc6 100644 --- a/lucene/core/src/java/org/apache/lucene/search/ExactPhraseMatcher.java +++ b/lucene/core/src/java/org/apache/lucene/search/ExactPhraseMatcher.java @@ -269,6 +269,7 @@ protected boolean lessThan(SubIterator a, SubIterator b) { boolean hasImpacts = false; List onlyImpactList = null; + List subIterators = new ArrayList<>(); for (int i = 0; i < impacts.length; ++i) { int impactsLevel = getLevel(impacts[i], docIdUpTo); if (impactsLevel == -1) { @@ -284,7 +285,7 @@ protected boolean lessThan(SubIterator a, SubIterator b) { } SubIterator subIterator = new SubIterator(impactList); - pq.add(subIterator); + subIterators.add(subIterator); if (hasImpacts == false) { hasImpacts = true; onlyImpactList = impactList; @@ -308,6 +309,7 @@ protected boolean lessThan(SubIterator a, SubIterator b) { // We walk impacts in parallel through a PQ ordered by freq. At any time, // the competitive impact consists of the lowest freq among all entries of // the PQ (the top) and the highest norm (tracked separately). + pq.addAll(subIterators); List mergedImpacts = new ArrayList<>(); SubIterator top = pq.top(); int currentFreq = top.current.freq; diff --git a/lucene/core/src/java/org/apache/lucene/util/PriorityQueue.java b/lucene/core/src/java/org/apache/lucene/util/PriorityQueue.java index fbd06abfe0d3..fe376cf2f01c 100644 --- a/lucene/core/src/java/org/apache/lucene/util/PriorityQueue.java +++ b/lucene/core/src/java/org/apache/lucene/util/PriorityQueue.java @@ -16,6 +16,7 @@ */ package org.apache.lucene.util; +import java.util.Collection; import java.util.Iterator; import java.util.NoSuchElementException; import java.util.function.Supplier; @@ -107,6 +108,32 @@ public PriorityQueue(int maxSize, Supplier sentinelObjectSupplier) { } } + /** + * Adds all elements of the collection into the queue. This method should be preferred over + * calling add() in loop if all elements are known in advance as it builds queue faster. + */ + public void addAll(Collection elements) { + if (this.size + elements.size() > this.maxSize) { + throw new IllegalArgumentException( + "Cannot add " + + elements.size() + + " elements to a queue with remaining capacity: " + + (maxSize - size)); + } + // Heap with size S always takes first S elements of the array, + // and thus it's safe to fill array further - no actual non-sentinel value will be overwritten. + Iterator iterator = elements.iterator(); + while (iterator.hasNext()) { + this.heap[size + 1] = iterator.next(); + this.size++; + } + + // The loop goes down to 1 as heap is 1-based not 0-based. + for (int i = (size >>> 1); i >= 1; i--) { + downHeap(i); + } + } + /** * Determines the ordering of objects in this priority queue. Subclasses must define this one * method. diff --git a/lucene/core/src/test/org/apache/lucene/util/TestPriorityQueue.java b/lucene/core/src/test/org/apache/lucene/util/TestPriorityQueue.java index 31525214dda1..567e761b0302 100644 --- a/lucene/core/src/test/org/apache/lucene/util/TestPriorityQueue.java +++ b/lucene/core/src/test/org/apache/lucene/util/TestPriorityQueue.java @@ -163,6 +163,43 @@ public void testInsertWithOverflow() { assertEquals((Integer) 2, pq.top()); } + public void testAddAllToEmptyQueue() { + int size = 10; + List list = new ArrayList<>(); + for (int i = 0; i < size; i++) { + list.add(random().nextInt()); + } + IntegerQueue pq = new IntegerQueue(size); + pq.addAll(list); + pq.checkValidity(); + } + + public void testAddAllToPartiallyFilledQueue() { + IntegerQueue pq = new IntegerQueue(20); + List list = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + list.add(random().nextInt()); + pq.add(random().nextInt()); + } + pq.addAll(list); + pq.checkValidity(); + } + + public void testAddAllDontFitIntoQueue() { + IntegerQueue pq = new IntegerQueue(20); + List list = new ArrayList<>(); + + for (int i = 0; i < 11; i++) { + list.add(random().nextInt()); + pq.add(random().nextInt()); + } + + assertThrows( + "Cannot add 11 elements to a queue with remaining capacity: 9", + IllegalArgumentException.class, + () -> pq.addAll(list)); + } + public void testRemovalsAndInsertions() { Random random = random(); int numDocsInPQ = TestUtil.nextInt(random, 1, 100);