package org.apache.lucene.queries.function.valuesource;
import java.io.IOException;
import java.util.Map;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.docvalues.FloatDocValues;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Weight;
import org.apache.lucene.util.mutable.MutableValue;
import org.apache.lucene.util.mutable.MutableValueFloat;
public class QueryValueSource extends ValueSource {
final Query q;
final float defVal;
public QueryValueSource(Query q, float defVal) {
this.q = q;
this.defVal = defVal;
}
public Query getQuery() { return q; }
public float getDefaultValue() { return defVal; }
@Override
public String description() {
return "query(" + q + ",def=" + defVal + ")";
}
@Override
public FunctionValues getValues(Map fcontext, LeafReaderContext readerContext) throws IOException {
return new QueryDocValues(this, readerContext, fcontext);
}
@Override
public int hashCode() {
return q.hashCode() * 29;
}
@Override
public boolean equals(Object o) {
if (QueryValueSource.class != o.getClass()) return false;
QueryValueSource other = (QueryValueSource)o;
return this.q.equals(other.q) && this.defVal==other.defVal;
}
@Override
public void createWeight(Map context, IndexSearcher searcher) throws IOException {
Query rewritten = searcher.rewrite(q);
Weight w = searcher.createWeight(rewritten, ScoreMode.COMPLETE, 1);
context.put(this, w);
}
}
class QueryDocValues extends FloatDocValues {
final LeafReaderContext readerContext;
final Weight weight;
final float defVal;
final Map fcontext;
final Query q;
Scorer scorer;
DocIdSetIterator it;
int scorerDoc;
boolean noMatches=false;
int lastDocRequested=Integer.MAX_VALUE;
public QueryDocValues(QueryValueSource vs, LeafReaderContext readerContext, Map fcontext) throws IOException {
super(vs);
this.readerContext = readerContext;
this.defVal = vs.defVal;
this.q = vs.q;
this.fcontext = fcontext;
Weight w = fcontext==null ? null : (Weight)fcontext.get(vs);
if (w == null) {
IndexSearcher weightSearcher;
if(fcontext == null) {
weightSearcher = new IndexSearcher(ReaderUtil.getTopLevelContext(readerContext));
} else {
weightSearcher = (IndexSearcher)fcontext.get("searcher");
if (weightSearcher == null) {
weightSearcher = new IndexSearcher(ReaderUtil.getTopLevelContext(readerContext));
}
}
vs.createWeight(fcontext, weightSearcher);
w = (Weight)fcontext.get(vs);
}
weight = w;
}
@Override
public float floatVal(int doc) {
try {
if (doc < lastDocRequested) {
if (noMatches) return defVal;
scorer = weight.scorer(readerContext);
if (scorer==null) {
noMatches = true;
return defVal;
}
it = scorer.iterator();
scorerDoc = -1;
}
lastDocRequested = doc;
if (scorerDoc < doc) {
scorerDoc = it.advance(doc);
}
if (scorerDoc > doc) {
return defVal;
}
return scorer.score();
} catch (IOException e) {
throw new RuntimeException("caught exception in QueryDocVals("+q+") doc="+doc, e);
}
}
@Override
public boolean exists(int doc) {
try {
if (doc < lastDocRequested) {
if (noMatches) return false;
scorer = weight.scorer(readerContext);
scorerDoc = -1;
if (scorer==null) {
noMatches = true;
return false;
}
it = scorer.iterator();
}
lastDocRequested = doc;
if (scorerDoc < doc) {
scorerDoc = it.advance(doc);
}
if (scorerDoc > doc) {
return false;
}
return true;
} catch (IOException e) {
throw new RuntimeException("caught exception in QueryDocVals("+q+") doc="+doc, e);
}
}
@Override
public Object objectVal(int doc) {
try {
return exists(doc) ? scorer.score() : null;
} catch (IOException e) {
throw new RuntimeException("caught exception in QueryDocVals("+q+") doc="+doc, e);
}
}
@Override
public ValueFiller getValueFiller() {
return new ValueFiller() {
private final MutableValueFloat mval = new MutableValueFloat();
@Override
public MutableValue getValue() {
return mval;
}
@Override
public void fillValue(int doc) {
try {
if (noMatches) {
mval.value = defVal;
mval.exists = false;
return;
}
scorer = weight.scorer(readerContext);
scorerDoc = -1;
if (scorer==null) {
noMatches = true;
mval.value = defVal;
mval.exists = false;
return;
}
it = scorer.iterator();
lastDocRequested = doc;
if (scorerDoc < doc) {
scorerDoc = it.advance(doc);
}
if (scorerDoc > doc) {
mval.value = defVal;
mval.exists = false;
return;
}
mval.value = scorer.score();
mval.exists = true;
} catch (IOException e) {
throw new RuntimeException("caught exception in QueryDocVals("+q+") doc="+doc, e);
}
}
};
}
@Override
public String toString(int doc) {
return "query(" + q + ",def=" + defVal + ")=" + floatVal(doc);
}
}