package org.apache.lucene.spatial.prefix;
import java.io.IOException;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.ConstantScoreScorer;
import org.apache.lucene.search.ConstantScoreWeight;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryVisitor;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Weight;
import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.DocIdSetBuilder;
import org.locationtech.spatial4j.shape.Shape;
public abstract class AbstractPrefixTreeQuery extends Query {
protected final Shape queryShape;
protected final String fieldName;
protected final SpatialPrefixTree grid;
protected final int detailLevel;
public AbstractPrefixTreeQuery(Shape queryShape, String fieldName, SpatialPrefixTree grid, int detailLevel) {
this.queryShape = queryShape;
this.fieldName = fieldName;
this.grid = grid;
this.detailLevel = detailLevel;
}
@Override
public boolean equals(Object o) {
return sameClassAs(o) &&
equalsTo(getClass().cast(o));
}
private boolean equalsTo(AbstractPrefixTreeQuery other) {
return detailLevel == other.detailLevel &&
fieldName.equals(other.fieldName) &&
queryShape.equals(other.queryShape);
}
@Override
public int hashCode() {
int result = classHash();
result = 31 * result + queryShape.hashCode();
result = 31 * result + fieldName.hashCode();
result = 31 * result + detailLevel;
return result;
}
@Override
public void visit(QueryVisitor visitor) {
if (visitor.acceptField(fieldName)) {
visitor.visitLeaf(this);
}
}
@Override
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
return new ConstantScoreWeight(this, boost) {
@Override
public Scorer scorer(LeafReaderContext context) throws IOException {
DocIdSet docSet = getDocIdSet(context);
if (docSet == null) {
return null;
}
DocIdSetIterator disi = docSet.iterator();
if (disi == null) {
return null;
}
return new ConstantScoreScorer(this, score(), scoreMode, disi);
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
return true;
}
};
}
protected abstract DocIdSet getDocIdSet(LeafReaderContext context) throws IOException;
public abstract class BaseTermsEnumTraverser {
protected final LeafReaderContext context;
protected final int maxDoc;
protected final Terms terms;
protected final TermsEnum termsEnum;
protected PostingsEnum postingsEnum;
public BaseTermsEnumTraverser(LeafReaderContext context) throws IOException {
this.context = context;
LeafReader reader = context.reader();
this.maxDoc = reader.maxDoc();
terms = reader.terms(fieldName);
if (terms != null) {
this.termsEnum = terms.iterator();
} else {
this.termsEnum = null;
}
}
protected void collectDocs(BitSet bitSet) throws IOException {
assert termsEnum != null;
postingsEnum = termsEnum.postings(postingsEnum, PostingsEnum.NONE);
bitSet.or(postingsEnum);
}
protected void collectDocs(DocIdSetBuilder docSetBuilder) throws IOException {
assert termsEnum != null;
postingsEnum = termsEnum.postings(postingsEnum, PostingsEnum.NONE);
docSetBuilder.add(postingsEnum);
}
}
}