You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by jp...@apache.org on 2025/05/22 09:49:54 UTC
(lucene) branch main updated: Speed up exhaustive evaluation. (#14679)
This is an automated email from the ASF dual-hosted git repository.
jpountz pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new 3dad852c1c0 Speed up exhaustive evaluation. (#14679)
3dad852c1c0 is described below
commit 3dad852c1c01a12e63078cc36a70a72d4471c90e
Author: Adrien Grand <jp...@gmail.com>
AuthorDate: Thu May 22 11:49:48 2025 +0200
Speed up exhaustive evaluation. (#14679)
This change helps speed up exhaustive evaluation of term queries, ie. calling
`DocIdSetIterator#nextDoc()` then `Scorer#score()` in a loop.
It helps in two ways:
- Iteration of matching doc IDs gets a bit more efficient, especially in the
case when a block of postings is encoded as a bit set.
- Computation of scores now gets (auto-)vectorized.
While this change doesn't help much when dynamic pruning kicks in, I'm hopeful
that we can improve this in the future.
---
lucene/CHANGES.txt | 3 ++
.../codecs/lucene103/Lucene103PostingsReader.java | 44 +++++++++++++++
.../java/org/apache/lucene/index/CheckIndex.java | 31 +++++++++++
.../java/org/apache/lucene/index/PostingsEnum.java | 45 ++++++++++++++++
.../org/apache/lucene/search/BooleanScorer.java | 29 +++++++---
.../org/apache/lucene/search/DocAndFreqBuffer.java | 62 ++++++++++++++++++++++
.../apache/lucene/search/DocAndScoreBuffer.java | 50 +++++++++++++++++
.../apache/lucene/search/MaxScoreBulkScorer.java | 27 +++++++---
.../src/java/org/apache/lucene/search/Scorer.java | 54 +++++++++++++++++++
.../java/org/apache/lucene/search/TermScorer.java | 55 +++++++++++++++++++
.../org/apache/lucene/search/TestTermScorer.java | 41 ++++++++++++++
.../lucene/tests/index/AssertingLeafReader.java | 15 ++++++
.../lucene/tests/index/RandomPostingsTester.java | 55 +++++++++++++++++++
.../lucene/tests/search/AssertingScorer.java | 17 ++++++
14 files changed, 513 insertions(+), 15 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 019d8adfc9c..c77fb1a1b61 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -113,6 +113,9 @@ Optimizations
* GITHUB#14609: Optimizes PointRangeQuery by rewriting to MatchAllDocsQuery/FieldExistsQuery/MatchNoDocsQuery if all
docs in index are contained or excluded (Elliott Bradshaw)
+* GITHUB#14679: Optimize exhaustive evaluation of disjunctive queries.
+ (Adrien Grand)
+
Bug Fixes
---------------------
* GITHUB#14654: ValueSource.fromDoubleValuesSource(dvs).getSortField() would throw errors when
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene103/Lucene103PostingsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene103/Lucene103PostingsReader.java
index 3ecaddef617..8b584368aa4 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene103/Lucene103PostingsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene103/Lucene103PostingsReader.java
@@ -46,6 +46,7 @@ import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.SegmentReadState;
import org.apache.lucene.internal.vectorization.PostingDecodingUtil;
import org.apache.lucene.internal.vectorization.VectorizationProvider;
+import org.apache.lucene.search.DocAndFreqBuffer;
import org.apache.lucene.store.ByteArrayDataInput;
import org.apache.lucene.store.ChecksumIndexInput;
import org.apache.lucene.store.DataInput;
@@ -1034,6 +1035,49 @@ public final class Lucene103PostingsReader extends PostingsReaderBase {
}
}
+ @Override
+ public void nextPostings(int upTo, DocAndFreqBuffer buffer) throws IOException {
+ assert needsRefilling == false;
+
+ if (needsFreq == false) {
+ super.nextPostings(upTo, buffer);
+ return;
+ }
+
+ buffer.size = 0;
+ if (doc >= upTo) {
+ return;
+ }
+
+ // Only return docs from the current block
+ buffer.growNoCopy(BLOCK_SIZE);
+ upTo = (int) Math.min(upTo, level0LastDocID + 1L);
+
+ // Frequencies are decoded lazily, calling freq() makes sure that the freq block is decoded
+ freq();
+
+ int start = docBufferUpto - 1;
+ switch (encoding) {
+ case PACKED:
+ int end = computeBufferEndBoundary(upTo);
+ buffer.size = end - start;
+ System.arraycopy(docBuffer, start, buffer.docs, 0, buffer.size);
+ break;
+ case UNARY:
+ docBitSet.forEach(
+ doc - docBitSetBase,
+ upTo - docBitSetBase,
+ docBitSetBase,
+ d -> buffer.docs[buffer.size++] = d);
+ break;
+ }
+
+ assert buffer.size > 0;
+ System.arraycopy(freqBuffer, start, buffer.freqs, 0, buffer.size);
+
+ advance(upTo);
+ }
+
private int computeBufferEndBoundary(int upTo) {
if (docBufferSize != 0 && docBuffer[docBufferSize - 1] < upTo) {
// All docs in the buffer are under upTo
diff --git a/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java b/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java
index c3369cfec39..e2d97b1355f 100644
--- a/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java
+++ b/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java
@@ -63,6 +63,7 @@ import org.apache.lucene.index.CheckIndex.Status.DocValuesStatus;
import org.apache.lucene.index.PointValues.IntersectVisitor;
import org.apache.lucene.index.PointValues.Relation;
import org.apache.lucene.internal.hppc.IntIntHashMap;
+import org.apache.lucene.search.DocAndFreqBuffer;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.FieldExistsQuery;
import org.apache.lucene.search.KnnCollector;
@@ -1454,6 +1455,7 @@ public final class CheckIndex implements Closeable {
int computedFieldCount = 0;
PostingsEnum postings = null;
+ PostingsEnum bulkPostings = null;
String lastField = null;
for (String field : fields) {
@@ -1649,6 +1651,10 @@ public final class CheckIndex implements Closeable {
sumDocFreq += docFreq;
postings = termsEnum.postings(postings, PostingsEnum.ALL);
+ bulkPostings = termsEnum.postings(bulkPostings, PostingsEnum.ALL);
+ bulkPostings.nextDoc();
+ DocAndFreqBuffer buffer = new DocAndFreqBuffer();
+ int bufferIndex = 0;
if (hasFreqs == false) {
if (termsEnum.totalTermFreq() != termsEnum.docFreq()) {
@@ -1697,6 +1703,31 @@ public final class CheckIndex implements Closeable {
throw new CheckIndexException(
"term " + term + ": doc " + doc + ": freq " + freq + " is out of bounds");
}
+
+ if (bufferIndex == buffer.size) {
+ bulkPostings.nextPostings(
+ (int) Math.min(Integer.MAX_VALUE, bulkPostings.docID() + 64L), buffer);
+ bufferIndex = 0;
+ }
+ if (bufferIndex >= buffer.size) {
+ throw new CheckIndexException("Doc " + doc + " not found by PostingsEnum#nextPostings");
+ }
+ if (doc != buffer.docs[bufferIndex]) {
+ throw new CheckIndexException(
+ "PostingsEnum#nextPostings returns "
+ + buffer.docs[bufferIndex]
+ + " as next doc while PostingsEnum#nextDoc returns "
+ + doc);
+ }
+ if (freq != buffer.freqs[bufferIndex]) {
+ throw new CheckIndexException(
+ "PostingsEnum#nextPostings returns "
+ + buffer.freqs[bufferIndex]
+ + " as term freq while PostingsEnum#freq returns "
+ + freq);
+ }
+ bufferIndex++;
+
if (hasFreqs == false) {
// When a field didn't index freq, it must
// consistently "lie" and pretend that freq was
diff --git a/lucene/core/src/java/org/apache/lucene/index/PostingsEnum.java b/lucene/core/src/java/org/apache/lucene/index/PostingsEnum.java
index 2cb0092aaa4..8e82e0a6d69 100644
--- a/lucene/core/src/java/org/apache/lucene/index/PostingsEnum.java
+++ b/lucene/core/src/java/org/apache/lucene/index/PostingsEnum.java
@@ -17,6 +17,7 @@
package org.apache.lucene.index;
import java.io.IOException;
+import org.apache.lucene.search.DocAndFreqBuffer;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.BytesRef;
@@ -97,4 +98,48 @@ public abstract class PostingsEnum extends DocIdSetIterator {
* anything (neither members of the returned BytesRef nor bytes in the byte[]).
*/
public abstract BytesRef getPayload() throws IOException;
+
+ /**
+ * Fill a buffer of doc IDs and frequencies with some number of doc IDs and their corresponding
+ * frequencies, starting at the current doc ID, and ending before {@code upTo}. Because it starts
+ * on the current doc ID, it is illegal to call this method if the {@link #docID() current doc ID}
+ * is {@code -1}.
+ *
+ * <p>An empty buffer after this method returns indicates that there are no postings left between
+ * the current doc ID and {@code upTo}.
+ *
+ * <p>Implementations should ideally fill the buffer with a number of entries comprised between 8
+ * and a couple hundreds, to keep heap requirements contained, while still being large enough to
+ * enable operations on the buffer to auto-vectorize efficiently.
+ *
+ * <p>The default implementation is provided below:
+ *
+ * <pre class="prettyprint">
+ * int batchSize = 16; // arbitrary
+ * buffer.growNoCopy(batchSize);
+ * int size = 0;
+ * for (int doc = docID(); doc < upTo && size < batchSize; doc = nextDoc()) {
+ * buffer.docs[size] = doc;
+ * buffer.freqs[size] = freq();
+ * ++size;
+ * }
+ * buffer.size = size;
+ * </pre>
+ *
+ * <p><b>NOTE</b>: The provided {@link DocAndFreqBuffer} should not hold references to internal
+ * data structures.
+ *
+ * @lucene.internal
+ */
+ public void nextPostings(int upTo, DocAndFreqBuffer buffer) throws IOException {
+ int batchSize = 16; // arbitrary
+ buffer.growNoCopy(batchSize);
+ int size = 0;
+ for (int doc = docID(); doc < upTo && size < batchSize; doc = nextDoc()) {
+ buffer.docs[size] = doc;
+ buffer.freqs[size] = freq();
+ ++size;
+ }
+ buffer.size = size;
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java b/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
index 0be42a2a3e7..1f4dd4f80ce 100644
--- a/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
+++ b/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
@@ -81,6 +81,7 @@ final class BooleanScorer extends BulkScorer {
final int minShouldMatch;
final long cost;
final boolean needsScores;
+ private final DocAndScoreBuffer docAndScoreBuffer = new DocAndScoreBuffer();
BooleanScorer(Collection<Scorer> scorers, int minShouldMatch, boolean needsScores) {
if (minShouldMatch < 1 || minShouldMatch > scorers.size()) {
@@ -135,23 +136,35 @@ final class BooleanScorer extends BulkScorer {
assert w.doc < max;
DocIdSetIterator it = w.iterator;
- int doc = w.doc;
- if (doc < min) {
- doc = it.advance(min);
+ if (w.doc < min) {
+ it.advance(min);
}
- if (buckets == null) {
+ if (buckets == null) { // means minShouldMatch=1 and scores are not needed
// This doesn't apply live docs, so we'll need to apply them later
it.intoBitSet(max, matching, base);
+ } else if (needsScores) {
+ for (w.scorer.nextDocsAndScores(max, acceptDocs, docAndScoreBuffer);
+ docAndScoreBuffer.size > 0;
+ w.scorer.nextDocsAndScores(max, acceptDocs, docAndScoreBuffer)) {
+ for (int index = 0; index < docAndScoreBuffer.size; ++index) {
+ final int doc = docAndScoreBuffer.docs[index];
+ final float score = docAndScoreBuffer.scores[index];
+ final int d = doc & MASK;
+ matching.set(d);
+ final Bucket bucket = buckets[d];
+ bucket.freq++;
+ bucket.score += score;
+ }
+ }
} else {
- for (; doc < max; doc = it.nextDoc()) {
+ // Scores are not needed but we need to keep track of freqs to know which hits match
+ assert minShouldMatch > 1;
+ for (int doc = it.docID(); doc < max; doc = it.nextDoc()) {
if (acceptDocs == null || acceptDocs.get(doc)) {
final int d = doc & MASK;
matching.set(d);
final Bucket bucket = buckets[d];
bucket.freq++;
- if (needsScores) {
- bucket.score += w.scorable.score();
- }
}
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/DocAndFreqBuffer.java b/lucene/core/src/java/org/apache/lucene/search/DocAndFreqBuffer.java
new file mode 100644
index 00000000000..3e45d3cf429
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/DocAndFreqBuffer.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.IntsRef;
+
+/**
+ * Wrapper around parallel arrays storing doc IDs and their corresponding frequencies.
+ *
+ * @lucene.internal
+ */
+public final class DocAndFreqBuffer {
+
+ /** Doc IDs */
+ public int[] docs = IntsRef.EMPTY_INTS;
+
+ /** Frequencies */
+ public int[] freqs = IntsRef.EMPTY_INTS;
+
+ /** Number of valid entries in the doc ID and frequency arrays. */
+ public int size;
+
+ /** Sole constructor. */
+ public DocAndFreqBuffer() {}
+
+ /** Grow both arrays to ensure that they can store at least the given number of entries. */
+ public void growNoCopy(int minSize) {
+ if (docs.length < minSize) {
+ docs = ArrayUtil.growNoCopy(docs, minSize);
+ freqs = new int[docs.length];
+ }
+ }
+
+ /** Remove entries from this buffer if their bit is unset in the given {@link Bits}. */
+ public void apply(Bits liveDocs) {
+ int newSize = 0;
+ for (int i = 0; i < size; ++i) {
+ if (liveDocs.get(docs[i])) {
+ docs[newSize] = docs[i];
+ freqs[newSize] = freqs[i];
+ newSize++;
+ }
+ }
+ this.size = newSize;
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/DocAndScoreBuffer.java b/lucene/core/src/java/org/apache/lucene/search/DocAndScoreBuffer.java
new file mode 100644
index 00000000000..5daef38c0ff
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/DocAndScoreBuffer.java
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.IntsRef;
+
+/**
+ * Wrapper around parallel arrays storing doc IDs and their corresponding scores.
+ *
+ * @lucene.internal
+ */
+public final class DocAndScoreBuffer {
+
+ private static final float[] EMPTY_FLOATS = new float[0];
+
+ /** Doc IDs */
+ public int[] docs = IntsRef.EMPTY_INTS;
+
+ /** Scores */
+ public float[] scores = EMPTY_FLOATS;
+
+ /** Number of valid entries in the doc ID and score arrays. */
+ public int size;
+
+ /** Sole constructor. */
+ public DocAndScoreBuffer() {}
+
+ /** Grow both arrays to ensure that they can store at least the given number of entries. */
+ public void growNoCopy(int minSize) {
+ if (docs.length < minSize) {
+ docs = ArrayUtil.growNoCopy(docs, minSize);
+ scores = new float[docs.length];
+ }
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/MaxScoreBulkScorer.java b/lucene/core/src/java/org/apache/lucene/search/MaxScoreBulkScorer.java
index 40e44b7c428..847fdb010a9 100644
--- a/lucene/core/src/java/org/apache/lucene/search/MaxScoreBulkScorer.java
+++ b/lucene/core/src/java/org/apache/lucene/search/MaxScoreBulkScorer.java
@@ -51,6 +51,8 @@ final class MaxScoreBulkScorer extends BulkScorer {
private final long[] windowMatches = new long[FixedBitSet.bits2words(INNER_WINDOW_SIZE)];
private final double[] windowScores = new double[INNER_WINDOW_SIZE];
+ private final DocAndScoreBuffer docAndScoreBuffer = new DocAndScoreBuffer();
+
MaxScoreBulkScorer(int maxDoc, List<Scorer> scorers, Scorer filter) throws IOException {
this.maxDoc = maxDoc;
this.filter = filter == null ? null : new DisiWrapper(filter, false);
@@ -218,12 +220,18 @@ final class MaxScoreBulkScorer extends BulkScorer {
// single essential clause in this window, we can iterate it directly and skip the bitset.
// this is a common case for 2-clauses queries
- for (int doc = top.doc; doc < upTo; doc = top.iterator.nextDoc()) {
- if (acceptDocs != null && acceptDocs.get(doc) == false) {
- continue;
+ for (top.scorer.nextDocsAndScores(upTo, acceptDocs, docAndScoreBuffer);
+ docAndScoreBuffer.size > 0;
+ top.scorer.nextDocsAndScores(upTo, acceptDocs, docAndScoreBuffer)) {
+ for (int i = 0; i < docAndScoreBuffer.size; ++i) {
+ scoreNonEssentialClauses(
+ collector,
+ docAndScoreBuffer.docs[i],
+ docAndScoreBuffer.scores[i],
+ firstEssentialScorer);
}
- scoreNonEssentialClauses(collector, doc, top.scorable.score(), firstEssentialScorer);
}
+
top.doc = top.iterator.docID();
essentialQueue.updateTop();
}
@@ -304,13 +312,18 @@ final class MaxScoreBulkScorer extends BulkScorer {
// Collect matches of essential clauses into a bitset
do {
- for (int doc = top.doc; doc < innerWindowMax; doc = top.iterator.nextDoc()) {
- if (acceptDocs == null || acceptDocs.get(doc)) {
+ for (top.scorer.nextDocsAndScores(innerWindowMax, acceptDocs, docAndScoreBuffer);
+ docAndScoreBuffer.size > 0;
+ top.scorer.nextDocsAndScores(innerWindowMax, acceptDocs, docAndScoreBuffer)) {
+ for (int index = 0; index < docAndScoreBuffer.size; ++index) {
+ final int doc = docAndScoreBuffer.docs[index];
+ final float score = docAndScoreBuffer.scores[index];
final int i = doc - innerWindowMin;
windowMatches[i >>> 6] |= 1L << i;
- windowScores[i] += top.scorable.score();
+ windowScores[i] += score;
}
}
+
top.doc = top.iterator.docID();
top = essentialQueue.updateTop();
} while (top.doc < innerWindowMax);
diff --git a/lucene/core/src/java/org/apache/lucene/search/Scorer.java b/lucene/core/src/java/org/apache/lucene/search/Scorer.java
index 7da2420207e..64db31ba9fd 100644
--- a/lucene/core/src/java/org/apache/lucene/search/Scorer.java
+++ b/lucene/core/src/java/org/apache/lucene/search/Scorer.java
@@ -17,6 +17,7 @@
package org.apache.lucene.search;
import java.io.IOException;
+import org.apache.lucene.util.Bits;
/**
* Expert: Common scoring functionality for different types of queries.
@@ -76,4 +77,57 @@ public abstract class Scorer extends Scorable {
* {@link #advanceShallow(int) shallow-advanced} to included and {@code upTo} included.
*/
public abstract float getMaxScore(int upTo) throws IOException;
+
+ /**
+ * Return a new batch of doc IDs and scores, starting at the current doc ID, and ending before
+ * {@code upTo}. Because it starts on the current doc ID, it is illegal to call this method if the
+ * {@link #docID() current doc ID} is {@code -1}.
+ *
+ * <p>An empty return value indicates that there are no postings left between the current doc ID
+ * and {@code upTo}.
+ *
+ * <p>Implementations should ideally fill the buffer with a number of entries comprised between 8
+ * and a couple hundreds, to keep heap requirements contained, while still being large enough to
+ * enable operations on the buffer to auto-vectorize efficiently.
+ *
+ * <p>The default implementation is provided below:
+ *
+ * <pre class="prettyprint">
+ * int batchSize = 16; // arbitrary
+ * buffer.growNoCopy(batchSize);
+ * int size = 0;
+ * DocIdSetIterator iterator = iterator();
+ * for (int doc = docID(); doc < upTo && size < batchSize; doc = iterator.nextDoc()) {
+ * if (liveDocs == null || liveDocs.get(doc)) {
+ * buffer.docs[size] = doc;
+ * buffer.scores[size] = score();
+ * ++size;
+ * }
+ * }
+ * buffer.size = size;
+ * </pre>
+ *
+ * <p><b>NOTE</b>: The provided {@link DocAndScoreBuffer} should not hold references to internal
+ * data structures.
+ *
+ * <p><b>NOTE</b>: In case this {@link Scorer} exposes a {@link #twoPhaseIterator()
+ * TwoPhaseIterator}, it should be positioned on a matching document before this method is called.
+ *
+ * @lucene.internal
+ */
+ public void nextDocsAndScores(int upTo, Bits liveDocs, DocAndScoreBuffer buffer)
+ throws IOException {
+ int batchSize = 16; // arbitrary
+ buffer.growNoCopy(batchSize);
+ int size = 0;
+ DocIdSetIterator iterator = iterator();
+ for (int doc = docID(); doc < upTo && size < batchSize; doc = iterator.nextDoc()) {
+ if (liveDocs == null || liveDocs.get(doc)) {
+ buffer.docs[size] = doc;
+ buffer.scores[size] = score();
+ ++size;
+ }
+ }
+ buffer.size = size;
+ }
}
diff --git a/lucene/core/src/java/org/apache/lucene/search/TermScorer.java b/lucene/core/src/java/org/apache/lucene/search/TermScorer.java
index 4b53788f233..ea272b6a4f3 100644
--- a/lucene/core/src/java/org/apache/lucene/search/TermScorer.java
+++ b/lucene/core/src/java/org/apache/lucene/search/TermScorer.java
@@ -17,11 +17,15 @@
package org.apache.lucene.search;
import java.io.IOException;
+import java.util.Arrays;
import org.apache.lucene.index.ImpactsEnum;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.SlowImpactsEnum;
import org.apache.lucene.search.similarities.Similarity.SimScorer;
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.LongsRef;
/**
* Expert: A <code>Scorer</code> for documents matching a <code>Term</code>.
@@ -35,6 +39,8 @@ public final class TermScorer extends Scorer {
private final NumericDocValues norms;
private final ImpactsDISI impactsDisi;
private final MaxScoreCache maxScoreCache;
+ private DocAndFreqBuffer docAndFreqBuffer;
+ private long[] normValues = LongsRef.EMPTY_LONGS;
/** Construct a {@link TermScorer} that will iterate all documents. */
public TermScorer(PostingsEnum postingsEnum, SimScorer scorer, NumericDocValues norms) {
@@ -120,4 +126,53 @@ public final class TermScorer extends Scorer {
impactsDisi.setMinCompetitiveScore(minScore);
}
}
+
+ @Override
+ public void nextDocsAndScores(int upTo, Bits liveDocs, DocAndScoreBuffer buffer)
+ throws IOException {
+ if (docAndFreqBuffer == null) {
+ docAndFreqBuffer = new DocAndFreqBuffer();
+ }
+
+ for (; ; ) {
+ postingsEnum.nextPostings(upTo, docAndFreqBuffer);
+ if (liveDocs != null && docAndFreqBuffer.size != 0) {
+ // An empty return value indicates that there are no more docs before upTo. We may be
+ // unlucky, and there are docs left, but all docs from the current batch happen to be marked
+ // as deleted. So we need to iterate until we find a batch that has at least one non-deleted
+ // doc.
+ docAndFreqBuffer.apply(liveDocs);
+ if (docAndFreqBuffer.size == 0) {
+ continue;
+ }
+ }
+ break;
+ }
+
+ int size = docAndFreqBuffer.size;
+ if (normValues.length < size) {
+ normValues = new long[ArrayUtil.oversize(size, Long.BYTES)];
+ if (norms == null) {
+ Arrays.fill(normValues, 1L);
+ }
+ }
+ if (norms != null) {
+ for (int i = 0; i < size; ++i) {
+ if (norms.advanceExact(docAndFreqBuffer.docs[i])) {
+ normValues[i] = norms.longValue();
+ } else {
+ normValues[i] = 1L;
+ }
+ }
+ }
+
+ buffer.growNoCopy(size);
+ buffer.size = size;
+ System.arraycopy(docAndFreqBuffer.docs, 0, buffer.docs, 0, size);
+ for (int i = 0; i < size; ++i) {
+ // Unless SimScorer#score is megamorphic, SimScorer#score should inline and (part of) score
+ // computations should auto-vectorize.
+ buffer.scores[i] = scorer.score(docAndFreqBuffer.freqs[i], normValues[i]);
+ }
+ }
}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestTermScorer.java b/lucene/core/src/test/org/apache/lucene/search/TestTermScorer.java
index 3823f2ea6d7..3264af9dc8b 100644
--- a/lucene/core/src/test/org/apache/lucene/search/TestTermScorer.java
+++ b/lucene/core/src/test/org/apache/lucene/search/TestTermScorer.java
@@ -39,6 +39,7 @@ import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.search.CheckHits;
import org.apache.lucene.tests.util.LuceneTestCase;
import org.apache.lucene.tests.util.TestUtil;
+import org.apache.lucene.util.Bits;
public class TestTermScorer extends LuceneTestCase {
protected Directory directory;
@@ -269,4 +270,44 @@ public class TestTermScorer extends LuceneTestCase {
reader.close();
dir.close();
}
+
+ public void testNextDocsAndScores() throws IOException {
+ Term allTerm = new Term(FIELD, "all");
+ TermQuery termQuery = new TermQuery(allTerm);
+ Weight weight = indexSearcher.createWeight(termQuery, ScoreMode.TOP_SCORES, 1f);
+ LeafReaderContext context = indexSearcher.getIndexReader().leaves().get(0);
+ Bits liveDocs = context.reader().getLiveDocs();
+ Scorer scorer1 = weight.scorer(context);
+ Scorer scorer2 = weight.scorer(context);
+ scorer1.iterator().nextDoc();
+ scorer2.iterator().nextDoc();
+ DocAndScoreBuffer buffer = new DocAndScoreBuffer();
+ while (true) {
+ int curDoc = scorer2.iterator().docID();
+ int upTo =
+ TestUtil.nextInt(random(), curDoc, (int) Math.min(Integer.MAX_VALUE, curDoc + 512L));
+ scorer1.nextDocsAndScores(upTo, liveDocs, buffer);
+ assertEquals(buffer.size == 0, curDoc >= upTo);
+
+ for (int i = 0; i < buffer.size; ++i) {
+ while (liveDocs != null && liveDocs.get(scorer2.iterator().docID()) == false) {
+ scorer2.iterator().nextDoc();
+ }
+ assertEquals(scorer2.iterator().docID(), buffer.docs[i]);
+ assertEquals(scorer2.score(), buffer.scores[i], 0f);
+ scorer2.iterator().nextDoc();
+ }
+
+ assertEquals(scorer2.iterator().docID(), scorer1.iterator().docID());
+ if (scorer1.iterator().docID() == DocIdSetIterator.NO_MORE_DOCS) {
+ break;
+ }
+ }
+
+ Scorer scorer3 = weight.scorer(context);
+ scorer3.iterator().nextDoc();
+ scorer3.nextDocsAndScores(
+ DocIdSetIterator.NO_MORE_DOCS, new Bits.MatchNoBits(context.reader().maxDoc()), buffer);
+ assertEquals(0, buffer.size);
+ }
}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java b/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java
index 6d4b6b56478..898d93fdc0f 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java
@@ -50,6 +50,7 @@ import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.internal.tests.IndexPackageAccess;
import org.apache.lucene.internal.tests.TestSecrets;
+import org.apache.lucene.search.DocAndFreqBuffer;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.BytesRef;
@@ -596,6 +597,20 @@ public class AssertingLeafReader extends FilterLeafReader {
positionCount = 0;
}
+ @Override
+ public void nextPostings(int upTo, DocAndFreqBuffer buffer) throws IOException {
+ assert state != DocsEnumState.START : "nextPostings() called before nextDoc()/advance()";
+ in.nextPostings(upTo, buffer);
+ doc = in.docID();
+ if (doc == DocIdSetIterator.NO_MORE_DOCS) {
+ state = DocsEnumState.FINISHED;
+ positionMax = 0;
+ } else {
+ state = DocsEnumState.ITERATING;
+ positionMax = super.freq();
+ }
+ }
+
void reset() {
state = DocsEnumState.START;
doc = in.docID();
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/index/RandomPostingsTester.java b/lucene/test-framework/src/java/org/apache/lucene/tests/index/RandomPostingsTester.java
index 2c27f55211c..d66d1667626 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/index/RandomPostingsTester.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/index/RandomPostingsTester.java
@@ -66,6 +66,7 @@ import org.apache.lucene.index.VectorEncoding;
import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.internal.tests.IndexPackageAccess;
import org.apache.lucene.internal.tests.TestSecrets;
+import org.apache.lucene.search.DocAndFreqBuffer;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FlushInfo;
@@ -114,6 +115,9 @@ public class RandomPostingsTester {
// Check DocIdSetIterator#intoBitSet
INTO_BIT_SET,
+ // Check PostingsEnum#nextDocsAndFreqs
+ NEXT_DOCS_AND_FREQS,
+
// Sometimes check payloads
PAYLOADS,
@@ -1410,6 +1414,57 @@ public class RandomPostingsTester {
set2.clear();
}
}
+
+ if (options.contains(Option.NEXT_DOCS_AND_FREQS)) {
+ int flags = PostingsEnum.FREQS;
+ if (doCheckPositions) {
+ flags |= PostingsEnum.POSITIONS;
+ if (doCheckOffsets) {
+ flags |= PostingsEnum.OFFSETS;
+ }
+ if (doCheckPayloads) {
+ flags |= PostingsEnum.PAYLOADS;
+ }
+ }
+
+ for (boolean impacts : new boolean[] {false, true}) {
+ PostingsEnum pe1;
+ if (impacts) {
+ pe1 = termsEnum.impacts(flags);
+ } else {
+ pe1 = termsEnum.postings(null, flags);
+ if (random.nextBoolean()) {
+ // test reuse
+ pe1.advance(maxDoc / 2);
+ pe1 = termsEnum.postings(pe1, flags);
+ }
+ }
+ PostingsEnum pe2 = termsEnum.postings(null, flags);
+
+ pe1.nextDoc();
+ pe2.nextDoc();
+ DocAndFreqBuffer buffer = new DocAndFreqBuffer();
+ while (true) {
+ int curDoc = pe1.docID();
+ int upTo =
+ TestUtil.nextInt(random, curDoc, (int) Math.min(Integer.MAX_VALUE, curDoc + 512L));
+ pe1.nextPostings(upTo, buffer);
+ assertEquals(buffer.size == 0, curDoc >= upTo);
+
+ for (int i = 0; i < buffer.size; ++i) {
+ assertTrue(buffer.docs[i] < upTo);
+ assertEquals(pe2.docID(), buffer.docs[i]);
+ assertEquals(pe2.freq(), buffer.freqs[i]);
+ pe2.nextDoc();
+ }
+
+ assertEquals(pe1.docID(), pe2.docID());
+ if (pe1.docID() == DocIdSetIterator.NO_MORE_DOCS) {
+ break;
+ }
+ }
+ }
+ }
}
private static class TestThread extends Thread {
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingScorer.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingScorer.java
index 9cb5e384b0d..1910a48d074 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingScorer.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingScorer.java
@@ -19,10 +19,12 @@ package org.apache.lucene.tests.search;
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
+import org.apache.lucene.search.DocAndScoreBuffer;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.FilterDocIdSetIterator;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.TwoPhaseIterator;
+import org.apache.lucene.util.Bits;
import org.apache.lucene.util.FixedBitSet;
/** Wraps a Scorer with additional checks */
@@ -282,4 +284,19 @@ public class AssertingScorer extends Scorer {
}
};
}
+
+ @Override
+ public void nextDocsAndScores(int upTo, Bits liveDocs, DocAndScoreBuffer buffer)
+ throws IOException {
+ assert doc != -1;
+ in.nextDocsAndScores(upTo, liveDocs, buffer);
+ if (doc != in.iterator().docID()) {
+ doc = in.iterator().docID();
+ if (doc == DocIdSetIterator.NO_MORE_DOCS) {
+ state = IteratorState.FINISHED;
+ } else {
+ state = IteratorState.ITERATING;
+ }
+ }
+ }
}