+++ /dev/null
-package org.apache.lucene.search;
-
-/**
- * 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.
- */
-
-import java.io.IOException;
-
-import org.apache.lucene.util.PriorityQueue;
-
-/** Represents hits returned by {@link
- * Searcher#search(Query,Filter,int)} and {@link
- * Searcher#search(Query,int)}. */
-public class TopDocs implements java.io.Serializable {
-
- /** The total number of hits for the query. */
- public int totalHits;
-
- /** The top hits for the query. */
- public ScoreDoc[] scoreDocs;
-
- /** Stores the maximum score value encountered, needed for normalizing. */
- private float maxScore;
-
- /**
- * Returns the maximum score value encountered. Note that in case
- * scores are not tracked, this returns {@link Float#NaN}.
- */
- public float getMaxScore() {
- return maxScore;
- }
-
- /** Sets the maximum score value encountered. */
- public void setMaxScore(float maxScore) {
- this.maxScore=maxScore;
- }
-
- /** Constructs a TopDocs with a default maxScore=Float.NaN. */
- TopDocs(int totalHits, ScoreDoc[] scoreDocs) {
- this(totalHits, scoreDocs, Float.NaN);
- }
-
- public TopDocs(int totalHits, ScoreDoc[] scoreDocs, float maxScore) {
- this.totalHits = totalHits;
- this.scoreDocs = scoreDocs;
- this.maxScore = maxScore;
- }
-
- // Refers to one hit:
- private static class ShardRef {
- // Which shard (index into shardHits[]):
- final int shardIndex;
-
- // Which hit within the shard:
- int hitIndex;
-
- public ShardRef(int shardIndex) {
- this.shardIndex = shardIndex;
- }
-
- @Override
- public String toString() {
- return "ShardRef(shardIndex=" + shardIndex + " hitIndex=" + hitIndex + ")";
- }
- };
-
- // Specialized MergeSortQueue that just merges by
- // relevance score, descending:
- private static class ScoreMergeSortQueue extends PriorityQueue<ShardRef> {
- final ScoreDoc[][] shardHits;
-
- public ScoreMergeSortQueue(TopDocs[] shardHits) {
- initialize(shardHits.length);
- this.shardHits = new ScoreDoc[shardHits.length][];
- for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
- this.shardHits[shardIDX] = shardHits[shardIDX].scoreDocs;
- }
- }
-
- // Returns true if first is < second
- public boolean lessThan(ShardRef first, ShardRef second) {
- assert first != second;
- final float firstScore = shardHits[first.shardIndex][first.hitIndex].score;
- final float secondScore = shardHits[second.shardIndex][second.hitIndex].score;
-
- if (firstScore < secondScore) {
- return false;
- } else if (firstScore > secondScore) {
- return true;
- } else {
- // Tie break: earlier shard wins
- if (first.shardIndex < second.shardIndex) {
- return true;
- } else if (first.shardIndex > second.shardIndex) {
- return false;
- } else {
- // Tie break in same shard: resolve however the
- // shard had resolved it:
- assert first.hitIndex != second.hitIndex;
- return first.hitIndex < second.hitIndex;
- }
- }
- }
- }
-
- private static class MergeSortQueue extends PriorityQueue<ShardRef> {
- // These are really FieldDoc instances:
- final ScoreDoc[][] shardHits;
- final FieldComparator[] comparators;
- final int[] reverseMul;
-
- public MergeSortQueue(Sort sort, TopDocs[] shardHits) throws IOException {
- initialize(shardHits.length);
- this.shardHits = new ScoreDoc[shardHits.length][];
- for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
- final ScoreDoc[] shard = shardHits[shardIDX].scoreDocs;
- //System.out.println(" init shardIdx=" + shardIDX + " hits=" + shard);
- if (shard != null) {
- this.shardHits[shardIDX] = shard;
- // Fail gracefully if API is misused:
- for(int hitIDX=0;hitIDX<shard.length;hitIDX++) {
- final ScoreDoc sd = shard[hitIDX];
- if (!(sd instanceof FieldDoc)) {
- throw new IllegalArgumentException("shard " + shardIDX + " was not sorted by the provided Sort (expected FieldDoc but got ScoreDoc)");
- }
- final FieldDoc fd = (FieldDoc) sd;
- if (fd.fields == null) {
- throw new IllegalArgumentException("shard " + shardIDX + " did not set sort field values (FieldDoc.fields is null); you must pass fillFields=true to IndexSearcher.search on each shard");
- }
- }
- }
- }
-
- final SortField[] sortFields = sort.getSort();
- comparators = new FieldComparator[sortFields.length];
- reverseMul = new int[sortFields.length];
- for(int compIDX=0;compIDX<sortFields.length;compIDX++) {
- final SortField sortField = sortFields[compIDX];
- comparators[compIDX] = sortField.getComparator(1, compIDX);
- reverseMul[compIDX] = sortField.getReverse() ? -1 : 1;
- }
- }
-
- // Returns true if first is < second
- @SuppressWarnings("unchecked")
- public boolean lessThan(ShardRef first, ShardRef second) {
- assert first != second;
- final FieldDoc firstFD = (FieldDoc) shardHits[first.shardIndex][first.hitIndex];
- final FieldDoc secondFD = (FieldDoc) shardHits[second.shardIndex][second.hitIndex];
- //System.out.println(" lessThan:\n first=" + first + " doc=" + firstFD.doc + " score=" + firstFD.score + "\n second=" + second + " doc=" + secondFD.doc + " score=" + secondFD.score);
-
- for(int compIDX=0;compIDX<comparators.length;compIDX++) {
- final FieldComparator comp = comparators[compIDX];
- //System.out.println(" cmp idx=" + compIDX + " cmp1=" + firstFD.fields[compIDX] + " cmp2=" + secondFD.fields[compIDX] + " reverse=" + reverseMul[compIDX]);
-
- final int cmp = reverseMul[compIDX] * comp.compareValues(firstFD.fields[compIDX], secondFD.fields[compIDX]);
-
- if (cmp != 0) {
- //System.out.println(" return " + (cmp < 0));
- return cmp < 0;
- }
- }
-
- // Tie break: earlier shard wins
- if (first.shardIndex < second.shardIndex) {
- //System.out.println(" return tb true");
- return true;
- } else if (first.shardIndex > second.shardIndex) {
- //System.out.println(" return tb false");
- return false;
- } else {
- // Tie break in same shard: resolve however the
- // shard had resolved it:
- //System.out.println(" return tb " + (first.hitIndex < second.hitIndex));
- assert first.hitIndex != second.hitIndex;
- return first.hitIndex < second.hitIndex;
- }
- }
- }
-
- /** Returns a new TopDocs, containing topN results across
- * the provided TopDocs, sorting by the specified {@link
- * Sort}. Each of the TopDocs must have been sorted by
- * the same Sort, and sort field values must have been
- * filled (ie, <code>fillFields=true</code> must be
- * passed to {@link
- * TopFieldCollector#create}.
- *
- * <p>Pass sort=null to merge sort by score descending.
- *
- * @lucene.experimental */
- public static TopDocs merge(Sort sort, int topN, TopDocs[] shardHits) throws IOException {
-
- final PriorityQueue<ShardRef> queue;
- if (sort == null) {
- queue = new ScoreMergeSortQueue(shardHits);
- } else {
- queue = new MergeSortQueue(sort, shardHits);
- }
-
- int totalHitCount = 0;
- int availHitCount = 0;
- float maxScore = Float.MIN_VALUE;
- for(int shardIDX=0;shardIDX<shardHits.length;shardIDX++) {
- final TopDocs shard = shardHits[shardIDX];
- if (shard.scoreDocs != null && shard.scoreDocs.length > 0) {
- totalHitCount += shard.totalHits;
- availHitCount += shard.scoreDocs.length;
- queue.add(new ShardRef(shardIDX));
- maxScore = Math.max(maxScore, shard.getMaxScore());
- //System.out.println(" maxScore now " + maxScore + " vs " + shard.getMaxScore());
- }
- }
-
- final ScoreDoc[] hits = new ScoreDoc[Math.min(topN, availHitCount)];
-
- int hitUpto = 0;
- while(hitUpto < hits.length) {
- assert queue.size() > 0;
- ShardRef ref = queue.pop();
- final ScoreDoc hit = shardHits[ref.shardIndex].scoreDocs[ref.hitIndex++];
- if (sort == null) {
- hits[hitUpto] = new ScoreDoc(hit.doc, hit.score, ref.shardIndex);
- } else {
- hits[hitUpto] = new FieldDoc(hit.doc, hit.score, ((FieldDoc) hit).fields, ref.shardIndex);
- }
-
- //System.out.println(" hitUpto=" + hitUpto);
- //System.out.println(" doc=" + hits[hitUpto].doc + " score=" + hits[hitUpto].score);
-
- hitUpto++;
-
- if (ref.hitIndex < shardHits[ref.shardIndex].scoreDocs.length) {
- // Not done with this these TopDocs yet:
- queue.add(ref);
- }
- }
-
- if (sort == null) {
- return new TopDocs(totalHitCount, hits, maxScore);
- } else {
- return new TopFieldDocs(totalHitCount, hits, sort.getSort(), maxScore);
- }
- }
-}
fnp / pylucene.git / blobdiff