78 lines
2.2 KiB
Java
78 lines
2.2 KiB
Java
package com.twitter.search.common.query;
|
|
|
|
import java.io.IOException;
|
|
|
|
import com.google.common.annotations.VisibleForTesting;
|
|
import com.google.common.base.Preconditions;
|
|
|
|
import org.apache.lucene.index.IndexReader;
|
|
import org.apache.lucene.search.IndexSearcher;
|
|
import org.apache.lucene.search.Query;
|
|
import org.apache.lucene.search.ScoreMode;
|
|
import org.apache.lucene.search.Weight;
|
|
|
|
/**
|
|
* Query implementation adds attribute collection support for an underlying query.
|
|
*/
|
|
public class IdentifiableQuery extends Query {
|
|
protected final Query inner;
|
|
private final FieldRankHitInfo queryId;
|
|
private final HitAttributeCollector attrCollector;
|
|
|
|
public IdentifiableQuery(Query inner, FieldRankHitInfo queryId,
|
|
HitAttributeCollector attrCollector) {
|
|
this.inner = Preconditions.checkNotNull(inner);
|
|
this.queryId = queryId;
|
|
this.attrCollector = Preconditions.checkNotNull(attrCollector);
|
|
}
|
|
|
|
@Override
|
|
public Weight createWeight(
|
|
IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
|
|
Weight innerWeight = inner.createWeight(searcher, scoreMode, boost);
|
|
return new IdentifiableQueryWeight(this, innerWeight, queryId, attrCollector);
|
|
}
|
|
|
|
@Override
|
|
public Query rewrite(IndexReader reader) throws IOException {
|
|
Query rewritten = inner.rewrite(reader);
|
|
if (rewritten != inner) {
|
|
return new IdentifiableQuery(rewritten, queryId, attrCollector);
|
|
}
|
|
return this;
|
|
}
|
|
|
|
@Override
|
|
public int hashCode() {
|
|
return inner.hashCode() * 13 + (queryId == null ? 0 : queryId.hashCode());
|
|
}
|
|
|
|
@Override
|
|
public boolean equals(Object obj) {
|
|
if (!(obj instanceof IdentifiableQuery)) {
|
|
return false;
|
|
}
|
|
|
|
IdentifiableQuery identifiableQuery = IdentifiableQuery.class.cast(obj);
|
|
return inner.equals(identifiableQuery.inner)
|
|
&& (queryId == null
|
|
? identifiableQuery.queryId == null
|
|
: queryId.equals(identifiableQuery.queryId));
|
|
}
|
|
|
|
@Override
|
|
public String toString(String field) {
|
|
return inner.toString(field);
|
|
}
|
|
|
|
@VisibleForTesting
|
|
public Query getQueryForTest() {
|
|
return inner;
|
|
}
|
|
|
|
@VisibleForTesting
|
|
public FieldRankHitInfo getQueryIdForTest() {
|
|
return queryId;
|
|
}
|
|
}
|