package org.elasticsearch.index.query.image; import net.semanticmetadata.lire.imageanalysis.LireFeature; import org.apache.lucene.index.*; import org.apache.lucene.search.*; import org.apache.lucene.search.similarities.DefaultSimilarity; import org.apache.lucene.util.Bits; import org.apache.lucene.util.ToStringUtils; import org.elasticsearch.common.lucene.search.Queries; import java.io.IOException; import java.util.Arrays; import java.util.BitSet; import java.util.Set; /** * Query by hash first and only calculate score for top n matches */ public class ImageHashLimitQuery extends Query { private String hashFieldName; private int[] hashes; private int maxResult; private String luceneFieldName; private LireFeature lireFeature; public ImageHashLimitQuery(String hashFieldName, int[] hashes, int maxResult, String luceneFieldName, LireFeature lireFeature, float boost) { this.hashFieldName = hashFieldName; this.hashes = hashes; this.maxResult = maxResult; this.luceneFieldName = luceneFieldName; this.lireFeature = lireFeature; setBoost(boost); } final class ImageHashScorer extends AbstractImageScorer { private int doc = -1; private final int maxDoc; private final int docBase; private final BitSet bitSet; private final Bits liveDocs; ImageHashScorer(Weight weight, BitSet bitSet, AtomicReaderContext context, Bits liveDocs) { super(weight, luceneFieldName, lireFeature, context.reader(), ImageHashLimitQuery.this.getBoost()); this.bitSet = bitSet; this.liveDocs = liveDocs; maxDoc = context.reader().maxDoc(); docBase = context.docBase; } @Override public int docID() { return doc; } @Override public int nextDoc() throws IOException { int d; do { d = bitSet.nextSetBit(docBase + doc + 1); if (d == -1 || d >= maxDoc + docBase) { doc = NO_MORE_DOCS; } else { doc = d - docBase; } } while (doc != NO_MORE_DOCS && d < maxDoc + docBase && liveDocs != null && !liveDocs.get(doc)); return doc; } @Override public int advance(int target) throws IOException { doc = target-1; return nextDoc(); } @Override public long cost() { return maxDoc; } } final class ImageHashLimitWeight extends Weight { private final BitSet bitSet; private final IndexSearcher searcher; public ImageHashLimitWeight(IndexSearcher searcher, BitSet bitSet) throws IOException { this.bitSet = bitSet; this.searcher = searcher; } @Override public String toString() { return "weight(" + ImageHashLimitQuery.this + ")"; } @Override public Query getQuery() { return ImageHashLimitQuery.this; } @Override public float getValueForNormalization() { return 1f; } @Override public void normalize(float queryNorm, float topLevelBoost) { } @Override public Scorer scorer(AtomicReaderContext context, Bits acceptDocs) throws IOException { return new ImageHashScorer(this, bitSet, context, acceptDocs); } @Override public Explanation explain(AtomicReaderContext context, int doc) throws IOException { Scorer scorer = scorer(context, context.reader().getLiveDocs()); if (scorer != null) { int newDoc = scorer.advance(doc); if (newDoc == doc) { float score = scorer.score(); ComplexExplanation result = new ComplexExplanation(); result.setDescription("ImageHashLimitQuery, product of:"); result.setValue(score); if (getBoost() != 1.0f) { result.addDetail(new Explanation(getBoost(),"boost")); score = score / getBoost(); } result.addDetail(new Explanation(score ,"image score (1/distance)")); result.setMatch(true); return result; } } return new ComplexExplanation(false, 0.0f, "no matching term"); } } @Override public Weight createWeight(IndexSearcher searcher) throws IOException { IndexSearcher indexSearcher = new IndexSearcher(searcher.getIndexReader()); indexSearcher.setSimilarity(new SimpleSimilarity()); BooleanQuery booleanQuery = new BooleanQuery(); for (int h : hashes) { booleanQuery.add(new BooleanClause(new TermQuery(new Term(hashFieldName, Integer.toString(h))), BooleanClause.Occur.SHOULD)); } TopDocs topDocs = indexSearcher.search(booleanQuery, maxResult); if (topDocs.scoreDocs.length == 0) { // no result find return Queries.newMatchNoDocsQuery().createWeight(searcher); } BitSet bitSet = new BitSet(topDocs.scoreDocs.length); for (ScoreDoc scoreDoc : topDocs.scoreDocs) { bitSet.set(scoreDoc.doc); } return new ImageHashLimitWeight(searcher, bitSet); } @Override public void extractTerms(Set<Term> terms) { } @Override public String toString(String field) { StringBuilder buffer = new StringBuilder(); buffer.append(hashFieldName); buffer.append(","); buffer.append(Arrays.toString(hashes)); buffer.append(","); buffer.append(maxResult); buffer.append(","); buffer.append(luceneFieldName); buffer.append(","); buffer.append(lireFeature.getClass().getSimpleName()); buffer.append(ToStringUtils.boost(getBoost())); return buffer.toString(); } @Override public boolean equals(Object o) { if (!(o instanceof ImageHashLimitQuery)) return false; ImageHashLimitQuery that = (ImageHashLimitQuery) o; if (maxResult != that.maxResult) return false; if (!hashFieldName.equals(that.hashFieldName)) return false; if (!Arrays.equals(hashes, that.hashes)) return false; if (!lireFeature.equals(that.lireFeature)) return false; if (!luceneFieldName.equals(that.luceneFieldName)) return false; return true; } @Override public int hashCode() { int result = super.hashCode(); result = 31 * result + hashFieldName.hashCode(); result = 31 * result + Arrays.hashCode(hashes); result = 31 * result + maxResult; result = 31 * result + luceneFieldName.hashCode(); result = 31 * result + lireFeature.hashCode(); return result; } final class SimpleSimilarity extends DefaultSimilarity{ @Override public float tf(float freq) { return 1; } @Override public float idf(long docFreq, long numDocs) { return 1; } @Override public float coord(int overlap, int maxOverlap) { return 1; } @Override public float queryNorm(float sumOfSquaredWeights) { return 1; } @Override public float lengthNorm(FieldInvertState state) { return 1; } @Override public float sloppyFreq(int distance) { return 1; } } }