diff --git a/lucene/core/src/java/org/apache/lucene/search/BooleanWeight.java b/lucene/core/src/java/org/apache/lucene/search/BooleanWeight.java index 778cb63..58baf1e 100644 --- a/lucene/core/src/java/org/apache/lucene/search/BooleanWeight.java +++ b/lucene/core/src/java/org/apache/lucene/search/BooleanWeight.java @@ -280,6 +280,11 @@ final class BooleanWeight extends Weight { @Override public BulkScorer bulkScorer(LeafReaderContext context) throws IOException { + // nocommit: hack + if (query.getMinimumNumberShouldMatch() <= 1 && + query.getClauses(Occur.SHOULD).size() == query.clauses().size()) { + return super.bulkScorer(context); + } final BulkScorer bulkScorer = booleanScorer(context); if (bulkScorer != null) { // bulk scoring is applicable, use it @@ -292,6 +297,24 @@ final class BooleanWeight extends Weight { @Override public Scorer scorer(LeafReaderContext context) throws IOException { + // nocommit: hack + if (query.getMinimumNumberShouldMatch() <= 1 && + query.getClauses(Occur.SHOULD).size() == query.clauses().size()) { + List scorers = new ArrayList<>(); + for (Weight weight : weights) { + Scorer scorer = weight.scorer(context); + if (scorer != null) { + scorers.add(scorer); + } + } + if (scorers.isEmpty()) { + return null; + } else if (scorers.size() == 1) { + return scorers.get(0); + } + return new MaxScoreScorer(this, scorers); + } + ScorerSupplier scorerSupplier = scorerSupplier(context); if (scorerSupplier == null) { return null; diff --git a/lucene/core/src/java/org/apache/lucene/search/DisiWrapper.java b/lucene/core/src/java/org/apache/lucene/search/DisiWrapper.java index f254340..a95ef1d 100644 --- a/lucene/core/src/java/org/apache/lucene/search/DisiWrapper.java +++ b/lucene/core/src/java/org/apache/lucene/search/DisiWrapper.java @@ -30,6 +30,7 @@ public class DisiWrapper { public final float matchCost; // the match cost for two-phase iterators, 0 otherwise public int doc; // the current doc, used for comparison public DisiWrapper next; // reference to a next element, see #topList + public final float maxScore; // An approximation of the iterator, or the iterator itself if it does not // support two-phase iteration @@ -58,6 +59,7 @@ public class DisiWrapper { approximation = iterator; matchCost = 0f; } + this.maxScore = scorer.maxScore(); } public DisiWrapper(Spans spans) { @@ -77,6 +79,7 @@ public class DisiWrapper { } this.lastApproxNonMatchDoc = -2; this.lastApproxMatchDoc = -2; + this.maxScore = Float.POSITIVE_INFINITY; } } diff --git a/lucene/core/src/java/org/apache/lucene/search/MaxScoreScorer.java b/lucene/core/src/java/org/apache/lucene/search/MaxScoreScorer.java new file mode 100644 index 0000000..1a113a0 --- /dev/null +++ b/lucene/core/src/java/org/apache/lucene/search/MaxScoreScorer.java @@ -0,0 +1,380 @@ +/* + * 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 static org.apache.lucene.search.DisiPriorityQueue.leftNode; +import static org.apache.lucene.search.DisiPriorityQueue.parentNode; +import static org.apache.lucene.search.DisiPriorityQueue.rightNode; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; + +final class MaxScoreScorer extends Scorer { + + private double minCompetitiveScore; + + // list of scorers which 'lead' the iteration and are currently + // positioned on 'doc' + DisiWrapper lead; + int doc; // current doc ID of the leads + int freq; // number of scorers on the desired doc ID + double leadMaxScore; // sum of the max scores of scorers in 'lead' + + // priority queue of scorers that are too advanced compared to the current + // doc. Ordered by doc ID. + final DisiPriorityQueue head; + + // priority queue of scorers which are behind the current doc. + // Ordered by maxScore. + final DisiWrapper[] tail; + double tailMaxScore; // sum of the max scores of scorers in 'tail' + int tailSize; + + final long cost; + + MaxScoreScorer(Weight weight, Collection scorers) { + super(weight); + + this.minCompetitiveScore = Double.NEGATIVE_INFINITY; + this.doc = -1; + + head = new DisiPriorityQueue(scorers.size()); + // there can be at most minShouldMatch - 1 scorers beyond the current position + tail = new DisiWrapper[scorers.size() - 1]; + + for (Scorer scorer : scorers) { + addLead(new DisiWrapper(scorer)); + } + + long cost = 0; + for (DisiWrapper w = lead; w != null; w = w.next) { + cost += w.cost; + } + + this.cost = cost; + } + + @Override + public void setMinCompetitiveScore(float minScore) { + this.minCompetitiveScore = minScore; + } + + @Override + public final Collection getChildren() throws IOException { + List matchingChildren = new ArrayList<>(); + updateFreq(); + for (DisiWrapper s = lead; s != null; s = s.next) { + matchingChildren.add(new ChildScorer(s.scorer, "SHOULD")); + } + return matchingChildren; + } + + @Override + public DocIdSetIterator iterator() { + return TwoPhaseIterator.asDocIdSetIterator(twoPhaseIterator()); + } + + @Override + public TwoPhaseIterator twoPhaseIterator() { + DocIdSetIterator approximation = new DocIdSetIterator() { + + @Override + public int docID() { + assert doc == lead.doc; + return doc; + } + + @Override + public int nextDoc() throws IOException { + return advance(doc + 1); + } + + @Override + public int advance(int target) throws IOException { + // Move 'lead' iterators back to the tail + for (DisiWrapper s = lead; s != null; s = s.next) { + final DisiWrapper evicted = insertTailWithOverFlow(s); + if (evicted != null) { + if (evicted.doc == target - 1) { + evicted.doc = evicted.iterator.nextDoc(); + } else { + evicted.doc = evicted.iterator.advance(target); + } + head.add(evicted); + } + } + + // Advance 'head' as well + DisiWrapper headTop = head.top(); + while (headTop.doc < target) { + final DisiWrapper evicted = insertTailWithOverFlow(headTop); + if (evicted != null) { + evicted.doc = evicted.iterator.advance(target); + headTop = head.updateTop(evicted); + } + } + + // Make sure there cannot be a match in the 'tail' or we could miss matches + ensureTailIsNotCompetitive(target); + + // Pop the new 'lead' from the 'head' + setDocAndFreq(); + // Advance to the next possible match + return doNextCandidate(); + } + + @Override + public long cost() { + return cost; + } + }; + return new TwoPhaseIterator(approximation) { + + @Override + public boolean matches() throws IOException { + while (leadMaxScore < minCompetitiveScore) { + if (leadMaxScore + tailMaxScore >= minCompetitiveScore) { + // a match on doc is still possible, try to + // advance scorers from the tail + advanceTail(); + } else { + return false; + } + } + return true; + } + + @Override + public float matchCost() { + // maximum number of scorer that matches() might advance + return tail.length; + } + + }; + } + + private void addLead(DisiWrapper lead) { + lead.next = this.lead; + this.lead = lead; + freq += 1; + leadMaxScore += lead.maxScore; + } + + private void pushBackLeads() throws IOException { + for (DisiWrapper s = lead; s != null; s = s.next) { + addTail(s); + } + } + + private void ensureTailIsNotCompetitive(int target) throws IOException { + while (tailSize > 0 && tailMaxScore >= minCompetitiveScore) { + DisiWrapper w = popTail(); + if (w.doc == target - 1) { + w.doc = w.iterator.nextDoc(); + } else { + w.doc = w.iterator.advance(target); + } + head.add(w); + } + } + + private void advanceTail(DisiWrapper top) throws IOException { + top.doc = top.iterator.advance(doc); + if (top.doc == doc) { + addLead(top); + } else { + head.add(top); + } + } + + private void advanceTail() throws IOException { + final DisiWrapper top = popTail(); + advanceTail(top); + } + + /** Reinitializes head, freq and doc from 'head' */ + private void setDocAndFreq() { + assert head.size() > 0; + + // The top of `head` defines the next potential match + // pop all documents which are on this doc + lead = head.pop(); + lead.next = null; + freq = 1; + leadMaxScore = lead.maxScore; + doc = lead.doc; + while (head.size() > 0 && head.top().doc == doc) { + addLead(head.pop()); + } + } + + /** Move iterators to the tail until the cumulated size of lead+tail is + * greater than or equal to minShouldMath */ + private int doNextCandidate() throws IOException { + while (leadMaxScore + tailMaxScore < minCompetitiveScore) { + // no match on doc is possible, move to the next potential match + if (head.size() == 0) { + // special case: the total max score is less than the min competitive score, there are no more matches + return doc = DocIdSetIterator.NO_MORE_DOCS; + } + pushBackLeads(); + ensureTailIsNotCompetitive(doc + 1); + setDocAndFreq(); + } + + return doc; + } + + /** Advance all entries from the tail to know about all matches on the + * current doc. */ + private void updateFreq() throws IOException { + assert leadMaxScore >= minCompetitiveScore; + // we return the next doc when there are minShouldMatch matching clauses + // but some of the clauses in 'tail' might match as well + // in general we want to advance least-costly clauses first in order to + // skip over non-matching documents as fast as possible. However here, + // we are advancing everything anyway so iterating over clauses in + // (roughly) cost-descending order might help avoid some permutations in + // the head heap + for (int i = tailSize - 1; i >= 0; --i) { + advanceTail(tail[i]); + } + tailSize = 0; + tailMaxScore = 0; + } + + @Override + public int freq() throws IOException { + // we need to know about all matches + updateFreq(); + return freq; + } + + @Override + public float score() throws IOException { + // we need to know about all matches + updateFreq(); + double score = 0; + for (DisiWrapper s = lead; s != null; s = s.next) { + score += s.scorer.score(); + } + return (float) score; + } + + @Override + public int docID() { + assert doc == lead.doc; + return doc; + } + + /** Insert an entry in 'tail' and evict the least-costly scorer if full. */ + private DisiWrapper insertTailWithOverFlow(DisiWrapper s) { + if (tailSize < tail.length) { + addTail(s); + return null; + } else if (tail.length >= 1) { + final DisiWrapper top = tail[0]; + if (top.cost < s.cost) { + tail[0] = s; + downHeapMaxScore(tail, tailSize); + return top; + } + } + return s; + } + + /** Add an entry to 'tail'. Fails if over capacity. */ + private void addTail(DisiWrapper s) { + tail[tailSize] = s; + upHeapMaxScore(tail, tailSize); + tailSize += 1; + tailMaxScore += s.maxScore; + } + + /** Pop the least-costly scorer from 'tail'. */ + private DisiWrapper popTail() { + assert tailSize > 0; + final DisiWrapper result = tail[0]; + tail[0] = tail[--tailSize]; + downHeapMaxScore(tail, tailSize); + // We don't subtract max scores as it could accumulate floating-point + // arithmetic errors or even return NaN if there are infinities. So we + // add all max scores from the tail again. + //tailMaxScore -= result.maxScore; + tailMaxScore = 0; + for (int i = 0; i < tailSize; ++i) { + tailMaxScore += tail[i].maxScore; + } + return result; + } + + /** Heap helpers */ + + private static void upHeapMaxScore(DisiWrapper[] heap, int i) { + final DisiWrapper node = heap[i]; + final double nodeMaxScore = node.maxScore; + int j = parentNode(i); + while (j >= 0 && nodeMaxScore > heap[j].maxScore) { + heap[i] = heap[j]; + i = j; + j = parentNode(j); + } + heap[i] = node; + } + + private static void downHeapMaxScore(DisiWrapper[] heap, int size) { + int i = 0; + final DisiWrapper node = heap[0]; + int j = leftNode(i); + if (j < size) { + int k = rightNode(j); + if (k < size && tailLessThan(heap[k], heap[j])) { + j = k; + } + if (tailLessThan(heap[j], node)) { + do { + heap[i] = heap[j]; + i = j; + j = leftNode(i); + k = rightNode(j); + if (k < size && tailLessThan(heap[k], heap[j])) { + j = k; + } + } while (j < size && tailLessThan(heap[j], node)); + heap[i] = node; + } + } + } + + /** + * In the tail, we want to get first entries that produce the maximum scores + * and in case of ties (eg. constant-score queries), those that have the least + * cost so that they are likely to advance further. + */ + private static boolean tailLessThan(DisiWrapper w1, DisiWrapper w2) { + if (w1.maxScore > w2.maxScore) { + return true; + } else if (w1.maxScore < w2.maxScore) { + return false; + } else { + return w1.cost < w2.cost; + } + } + +} 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 7ceed33..03d2b34 100644 --- a/lucene/core/src/java/org/apache/lucene/search/Scorer.java +++ b/lucene/core/src/java/org/apache/lucene/search/Scorer.java @@ -146,4 +146,17 @@ public abstract class Scorer { public TwoPhaseIterator twoPhaseIterator() { return null; } + + /** Tell the scorer that its iterator may safely ignore all documents whose + * score is less than the given {@code minScore}. This is a no-op by + * default. */ + public void setMinCompetitiveScore(float minScore) { + // no-op by default + } + + /** Return the maximum score that this scorer may produce. The default + * implementation returns +Infinity. */ + public float maxScore() { + return Float.POSITIVE_INFINITY; + } } 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 6c823f1..eeeb669 100644 --- a/lucene/core/src/java/org/apache/lucene/search/TermScorer.java +++ b/lucene/core/src/java/org/apache/lucene/search/TermScorer.java @@ -66,6 +66,11 @@ final class TermScorer extends Scorer { return docScorer.score(postingsEnum.docID(), postingsEnum.freq()); } + @Override + public float maxScore() { + return docScorer.maxScore(); + } + /** Returns a string representation of this TermScorer. */ @Override public String toString() { return "scorer(" + weight + ")[" + super.toString() + "]"; } diff --git a/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java b/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java index a840b82..431c78e 100644 --- a/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/TopScoreDocCollector.java @@ -49,8 +49,11 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { private static class SimpleTopScoreDocCollector extends TopScoreDocCollector { - SimpleTopScoreDocCollector(int numHits) { + private final boolean needsTotalHits; + + SimpleTopScoreDocCollector(int numHits, boolean needsTotalHits) { super(numHits); + this.needsTotalHits = needsTotalHits; } @Override @@ -60,6 +63,16 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { return new ScorerLeafCollector() { @Override + public void setScorer(Scorer scorer) throws IOException { + super.setScorer(scorer); + if (needsTotalHits == false && pqTop != null) { + // since we tie-break on doc id and collect in doc id order, we can require + // the next float + scorer.setMinCompetitiveScore(Math.nextUp(pqTop.score)); + } + } + + @Override public void collect(int doc) throws IOException { float score = scorer.score(); @@ -77,6 +90,11 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { pqTop.doc = doc + docBase; pqTop.score = score; pqTop = pq.updateTop(); + if (needsTotalHits == false) { + // since we tie-break on doc id and collect in doc id order, we can require + // the next float + scorer.setMinCompetitiveScore(Math.nextUp(score)); + } } }; @@ -171,12 +189,17 @@ public abstract class TopScoreDocCollector extends TopDocsCollector { } if (after == null) { - return new SimpleTopScoreDocCollector(numHits); + return new SimpleTopScoreDocCollector(numHits, false); // nocommit: should be true } else { return new PagingTopScoreDocCollector(numHits, after); } } + // nocommit: better API + public static TopScoreDocCollector create(int numHits, boolean needsTotalHits) { + return new SimpleTopScoreDocCollector(numHits, needsTotalHits); + } + ScoreDoc pqTop; // prevents instantiation diff --git a/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java b/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java index 35554e2..6501488d 100644 --- a/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java +++ b/lucene/core/src/java/org/apache/lucene/search/similarities/BM25Similarity.java @@ -238,7 +238,12 @@ public class BM25Similarity extends Similarity { } return weightValue * freq / (freq + norm); } - + + @Override + public float maxScore() { + return weightValue; // nocommit: we can do better + } + @Override public Explanation explain(int doc, Explanation freq) throws IOException { return explainScore(doc, freq, stats, norms, lengthCache); diff --git a/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java b/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java index 7f0f27c..ece0d5a 100644 --- a/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java +++ b/lucene/core/src/java/org/apache/lucene/search/similarities/Similarity.java @@ -168,7 +168,13 @@ public abstract class Similarity { /** Calculate a scoring factor based on the data in the payload. */ public abstract float computePayloadFactor(int doc, int start, int end, BytesRef payload); - + + /** Return the maximum score that this scorer may produce. */ + // nocommit: make abstract + public float maxScore() { + return Float.POSITIVE_INFINITY; + } + /** * Explain the score for a single document * @param doc document id within the inverted index segment