package org.apache.lucene.queries.function;
import java.io.IOException;
import java.util.Objects;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.DoubleValues;
import org.apache.lucene.search.DoubleValuesSource;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LongValues;
import org.apache.lucene.search.LongValuesSource;
public final class IndexReaderFunctions {
private IndexReaderFunctions() {}
public static DoubleValuesSource docFreq(Term term) {
return new IndexReaderDoubleValuesSource(r -> (double) r.docFreq(term), "docFreq(" + term.toString() + ")");
}
public static DoubleValuesSource maxDoc() {
return new IndexReaderDoubleValuesSource(IndexReader::maxDoc, "maxDoc()");
}
public static DoubleValuesSource numDocs() {
return new IndexReaderDoubleValuesSource(IndexReader::numDocs, "numDocs()");
}
public static DoubleValuesSource numDeletedDocs() {
return new IndexReaderDoubleValuesSource(IndexReader::numDeletedDocs, "numDeletedDocs()");
}
public static LongValuesSource sumTotalTermFreq(String field) {
return new SumTotalTermFreqValuesSource(field);
}
private static class SumTotalTermFreqValuesSource extends LongValuesSource {
private final String field;
private SumTotalTermFreqValuesSource(String field) {
this.field = field;
}
@Override
public LongValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
throw new UnsupportedOperationException("IndexReaderFunction must be rewritten before use");
}
@Override
public boolean needsScores() {
return false;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
SumTotalTermFreqValuesSource that = (SumTotalTermFreqValuesSource) o;
return Objects.equals(field, that.field);
}
@Override
public int hashCode() {
return Objects.hash(field);
}
@Override
public LongValuesSource rewrite(IndexSearcher searcher) throws IOException {
return new NoCacheConstantLongValuesSource(searcher.getIndexReader().getSumTotalTermFreq(field), this);
}
@Override
public String toString() {
return "sumTotalTermFreq(" + field + ")";
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
return false;
}
}
private static class NoCacheConstantLongValuesSource extends LongValuesSource {
final long value;
final LongValuesSource parent;
private NoCacheConstantLongValuesSource(long value, LongValuesSource parent) {
this.value = value;
this.parent = parent;
}
@Override
public LongValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
return new LongValues() {
@Override
public long longValue() throws IOException {
return value;
}
@Override
public boolean advanceExact(int doc) throws IOException {
return true;
}
};
}
@Override
public boolean needsScores() {
return false;
}
@Override
public LongValuesSource rewrite(IndexSearcher reader) throws IOException {
return this;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (!(o instanceof NoCacheConstantLongValuesSource)) return false;
NoCacheConstantLongValuesSource that = (NoCacheConstantLongValuesSource) o;
return value == that.value &&
Objects.equals(parent, that.parent);
}
@Override
public int hashCode() {
return Objects.hash(value, parent);
}
@Override
public String toString() {
return parent.toString();
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
return false;
}
}
public static DoubleValuesSource termFreq(Term term) {
return new TermFreqDoubleValuesSource(term);
}
private static class TermFreqDoubleValuesSource extends DoubleValuesSource {
private final Term term;
private TermFreqDoubleValuesSource(Term term) {
this.term = term;
}
@Override
public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
Terms terms = ctx.reader().terms(term.field());
TermsEnum te = terms == null ? null : terms.iterator();
if (te == null || te.seekExact(term.bytes()) == false) {
return DoubleValues.EMPTY;
}
final PostingsEnum pe = te.postings(null);
assert pe != null;
return new DoubleValues() {
@Override
public double doubleValue() throws IOException {
return pe.freq();
}
@Override
public boolean advanceExact(int doc) throws IOException {
if (pe.docID() > doc)
return false;
return pe.docID() == doc || pe.advance(doc) == doc;
}
};
}
@Override
public boolean needsScores() {
return false;
}
@Override
public DoubleValuesSource rewrite(IndexSearcher searcher) throws IOException {
return this;
}
@Override
public String toString() {
return "termFreq(" + term.toString() + ")";
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
TermFreqDoubleValuesSource that = (TermFreqDoubleValuesSource) o;
return Objects.equals(term, that.term);
}
@Override
public int hashCode() {
return Objects.hash(term);
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
return true;
}
}
public static DoubleValuesSource totalTermFreq(Term term) {
return new IndexReaderDoubleValuesSource(r -> r.totalTermFreq(term), "totalTermFreq(" + term.toString() + ")");
}
public static DoubleValuesSource sumDocFreq(String field) {
return new IndexReaderDoubleValuesSource(r -> r.getSumDocFreq(field), "sumDocFreq(" + field + ")");
}
public static DoubleValuesSource docCount(String field) {
return new IndexReaderDoubleValuesSource(r -> r.getDocCount(field), "docCount(" + field + ")");
}
@FunctionalInterface
private interface ReaderFunction {
double apply(IndexReader reader) throws IOException;
}
private static class IndexReaderDoubleValuesSource extends DoubleValuesSource {
private final ReaderFunction func;
private final String description;
private IndexReaderDoubleValuesSource(ReaderFunction func, String description) {
this.func = func;
this.description = description;
}
@Override
public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
throw new UnsupportedOperationException("IndexReaderFunction must be rewritten before use");
}
@Override
public boolean needsScores() {
return false;
}
@Override
public DoubleValuesSource rewrite(IndexSearcher searcher) throws IOException {
return new NoCacheConstantDoubleValuesSource(func.apply(searcher.getIndexReader()), this);
}
@Override
public String toString() {
return description;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
IndexReaderDoubleValuesSource that = (IndexReaderDoubleValuesSource) o;
return Objects.equals(description, that.description) && Objects.equals(func, that.func);
}
@Override
public int hashCode() {
return Objects.hash(description, func);
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
return false;
}
}
private static class NoCacheConstantDoubleValuesSource extends DoubleValuesSource {
final double value;
final DoubleValuesSource parent;
private NoCacheConstantDoubleValuesSource(double value, DoubleValuesSource parent) {
this.value = value;
this.parent = parent;
}
@Override
public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
return new DoubleValues() {
@Override
public double doubleValue() throws IOException {
return value;
}
@Override
public boolean advanceExact(int doc) throws IOException {
return true;
}
};
}
@Override
public boolean needsScores() {
return false;
}
@Override
public DoubleValuesSource rewrite(IndexSearcher reader) throws IOException {
return this;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (!(o instanceof NoCacheConstantDoubleValuesSource)) return false;
NoCacheConstantDoubleValuesSource that = (NoCacheConstantDoubleValuesSource) o;
return Double.compare(that.value, value) == 0 &&
Objects.equals(parent, that.parent);
}
@Override
public int hashCode() {
return Objects.hash(value, parent);
}
@Override
public String toString() {
return parent.toString();
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
return false;
}
}
}