DocSliceResult.java

package zh.solr.se.searcher.relevance;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;

import org.apache.solr.search.DocIterator;
import org.apache.solr.search.DocSlice;
import org.apache.solr.search.SolrIndexSearcher;

import zh.solr.se.searcher.solr.SolrUtil;

public class DocSliceResult extends ArrayList<ScoredSolrDoc> implements SearchResult {	
	private int matches = 0;

	/**
	 * construct an empty SearchResult object
	 */
	public DocSliceResult() {
	}
	
	/**
	 * Construct a SearchResult object from a DocSlice object
	 * @param docSlice the DocSlice object
	 */
	public DocSliceResult(DocSlice docSlice) {
		this(docSlice, null, null);
	}
	
	/**
	 * Construct a SearchResult object from a DocSlice object
	 * @param docSlice the DocSlice object
	 * @param searcher Solr index searcher
	 * @param fieldSet the desired field names and whether they multi-valued
	 */
	public DocSliceResult(DocSlice docSlice, SolrIndexSearcher searcher, Set<String> fieldSet) {
		if (docSlice == null)
			throw new NullPointerException("The DocSlice object for constructing search result must not be null");
		
		if (fieldSet != null && searcher == null) {
			throw new NullPointerException("Searcher must not be null to retieve field values");
		}
		
		// get the matches
		this.setMatches(docSlice.matches());
		
		// retrieve all the documents
		DocIterator iterator = docSlice.iterator();
		while (iterator.hasNext()) {
			int docId = iterator.nextDoc();
			float score = 1.0f;			
			if (docSlice.hasScores())
				score = iterator.score();
			
			ScoredSolrDoc doc = new ScoredSolrDoc(docId, score);
			
			// retrieve requested field values
			if (fieldSet != null) {
				Map<String, String> valueMap = SolrUtil.retrieveFieldValues(searcher, docId, fieldSet);
				doc.setFieldValueMap(valueMap);
			}
			
			this.add(doc);
		}
	}
	
	public void addDoc(Object doc) {
		if (doc == null)
			return;
		
		super.add((ScoredSolrDoc)doc);
		matches++;
	}
	
	/**
	 * Sort the doc list by score from high to low, and only return the specified number of docs
	 * @param maxCount max number of docs to return
	 */
	public void sortByScore(int maxCount) {
		if (this.size() < 2)
			return;
		
		// creating the comparator that compares doc scores from high to low
		Comparator<ScoredSolrDoc> comparator = new Comparator<ScoredSolrDoc>() {
			public int compare(ScoredSolrDoc doc1, ScoredSolrDoc doc2) {
				return (int)Math.signum(doc2.getScore() - doc1.getScore());
			}
		};
		
		if (maxCount <= 0 || maxCount >= this.size()) {
			// full sorting
			Collections.sort(this, comparator);
		} else {
			// use priority queue for partial sorting
			PriorityQueue<ScoredSolrDoc> docQueue = new PriorityQueue<ScoredSolrDoc>(maxCount, comparator);
			for (ScoredSolrDoc doc : this) {
				docQueue.add(doc);
			}
			
			// replace the original list with the items in the queue
			this.clear();
			for (int i = 0; i < maxCount; i++) {
				this.add(docQueue.poll());
			}
		}
	}
	
	/**
	 * Convert to a DocSlice object
	 */
	public DocSlice toSolrSearchResult() {
		int length = size();
		int[] idList = new int[length];
		float[] scoreList = new float[length];
		float maxScore = 0f;
		for (int i = 0; i < length; i++) {
			ScoredSolrDoc solrDoc = get(i);
			idList[i] = solrDoc.getDocId();
			scoreList[i] = solrDoc.getScore();
			if (scoreList[i] > maxScore)
				maxScore = scoreList[i];
		}
		
		return new DocSlice(0, length, idList, scoreList, matches, maxScore);
	}
	
	public SearchResult merge(SearchResult result1, SearchResult result2, int maxCount) {
		if (maxCount <= 0)
			maxCount = DEFAULT_PAGE_SIZE;
		
		if (result1 == null)
			return result2;
		else if (result2 == null || ((DocSliceResult)result1).size() >= maxCount)
			return result1;
				
		DocSliceResult mergedResult = new DocSliceResult();
		
		// append the first result to the merged result
		mergedResult.addAll((DocSliceResult)result1);
		
		// append the second result to the merged result, excluding any duplicates
		HashSet<Object> docIdHashSet = result1.getDocIdHashSet();
		int duplicates = 0;
		for (int i = 0; (i < ((DocSliceResult)result2).size() && mergedResult.size() <= maxCount); i++) {
			ScoredSolrDoc doc = ((DocSliceResult)result2).get(i);
			if (docIdHashSet.contains(doc.getDocId())) {
				duplicates++;
			} else {
				mergedResult.add(doc);
				docIdHashSet.add(doc.getDocId());
			}
		}
		
		// exclude the duplicates from the total matches
		int totalMatches = result1.getMatches() + result2.getMatches() - duplicates;
		mergedResult.setMatches(totalMatches);
		
		return mergedResult;
	}

	public void mergeWith(SearchResult anotherResult, int maxCount) {
		if (anotherResult == null || ((DocSliceResult)anotherResult).size() == 0)
			return;
		
		if (maxCount <= 0)
			maxCount = DEFAULT_PAGE_SIZE;
		
		int totalMatches = this.getMatches() + anotherResult.getMatches();
		if (maxCount <= this.size()) {
			// we need to cut the current result, and we don't need the second result.
			this.removeRange(maxCount, this.size());
		} else {
			// append needed docs
			HashSet<Object> docIdHashSet = this.getDocIdHashSet();
			int duplicates = 0;
			for (int i = 0; (i < ((DocSliceResult)anotherResult).size() && this.size() <= maxCount); i++) {
				ScoredSolrDoc doc = ((DocSliceResult)anotherResult).get(i);
				if (docIdHashSet.contains(doc.getDocId())) {
					duplicates++;
				} else {
					this.add(doc);
					docIdHashSet.add(doc.getDocId());
				}
			}
			
			// remove the duplicates from the total matches
			totalMatches -= duplicates;
		}
		
		setMatches(totalMatches);
	}
	
	public void scaleScores(float scaleFactor) {
		if (scaleFactor < 0)
			return;
		for (ScoredSolrDoc doc : this) {
			doc.setScore(doc.getScore() * scaleFactor);
		}
	}
	
	public void discardLowScoreResults(float threshold) {
		// cut from the end, so that the earlier index does not change after the cut
		// please note that the scores don't necessarily in any order
		int origSize = this.size();
		for (int i = origSize - 1; i >= 0; i--) {
			if (this.get(i).getScore() < threshold) {
				this.remove(i);
			}
		}
		
		// if any of the search result on the first page is removed, then the rest on later pages
		// should be removed too. We should update the matches to what are left on the first page
		if (this.size() < origSize)
			setMatches(this.size());
	}
	
	public int getMatches() {
		return matches;
	}

	public void setMatches(int matches) {
		this.matches = matches;
	}

	public Float getMaxScore() {
		float maxScore = 0f;
		for (ScoredSolrDoc doc : this) {
			float tempScore = doc.getScore();
			if (tempScore > maxScore)
				maxScore = tempScore;
		}
		
		return maxScore;
	}
	
	public HashSet<Object> getDocIdHashSet() {
		HashSet<Object> idSet = new HashSet<Object>();
		for (ScoredSolrDoc doc : this) {
			idSet.add(doc.getDocId());
		}
		
		return idSet;
	}
	
	public void removeRange (int fromIndex, int toIndex) {
		super.removeRange(fromIndex, toIndex);
	}

}