93 lines
2.7 KiB
Java
93 lines
2.7 KiB
Java
package it.cavallium.dbengine.lucene.collector;
|
|
|
|
import static it.cavallium.dbengine.database.LLUtils.mapList;
|
|
|
|
import it.cavallium.dbengine.lucene.IntSmear;
|
|
import it.unimi.dsi.fastutil.ints.IntHash;
|
|
import java.io.IOException;
|
|
import java.util.Collection;
|
|
import org.apache.lucene.facet.FacetsCollectorManager;
|
|
import org.apache.lucene.index.LeafReaderContext;
|
|
import org.apache.lucene.search.CollectorManager;
|
|
import org.apache.lucene.search.DocIdSetIterator;
|
|
import org.apache.lucene.search.LeafCollector;
|
|
import org.apache.lucene.search.Scorable;
|
|
import org.apache.lucene.search.ScoreMode;
|
|
|
|
public class FastFacetsCollectorManager implements CollectorManager<FacetsCollector, FacetsCollector> {
|
|
|
|
private final int collectionRate;
|
|
private final IntHash.Strategy hash;
|
|
private final FacetsCollectorManager facetsCollectorManager;
|
|
|
|
public FastFacetsCollectorManager(int collectionRate) {
|
|
this.collectionRate = collectionRate;
|
|
this.hash = new IntSmear();
|
|
this.facetsCollectorManager = new FacetsCollectorManager();
|
|
}
|
|
|
|
@Override
|
|
public FacetsCollector newCollector() {
|
|
return new FastFacetsCollector(collectionRate, hash);
|
|
}
|
|
|
|
@Override
|
|
public FacetsCollector reduce(Collection<FacetsCollector> collectors) throws IOException {
|
|
return FacetsCollector.wrap(facetsCollectorManager.reduce(mapList(collectors,
|
|
facetsCollector -> facetsCollector.getLuceneFacetsCollector()
|
|
)));
|
|
}
|
|
|
|
private static class FastFacetsCollector implements FacetsCollector {
|
|
|
|
private final org.apache.lucene.facet.FacetsCollector collector;
|
|
private final int collectionRate;
|
|
private final IntHash.Strategy hash;
|
|
|
|
public FastFacetsCollector(int collectionRate, IntHash.Strategy hash) {
|
|
this.collectionRate = collectionRate;
|
|
this.hash = hash;
|
|
this.collector = new org.apache.lucene.facet.FacetsCollector(false) {
|
|
@Override
|
|
public ScoreMode scoreMode() {
|
|
return ScoreMode.COMPLETE_NO_SCORES;
|
|
}
|
|
};
|
|
}
|
|
|
|
|
|
@Override
|
|
public org.apache.lucene.facet.FacetsCollector getLuceneFacetsCollector() {
|
|
return collector;
|
|
}
|
|
|
|
@Override
|
|
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
|
|
var leafCollector = collector.getLeafCollector(context);
|
|
return new LeafCollector() {
|
|
@Override
|
|
public void setScorer(Scorable scorer) throws IOException {
|
|
leafCollector.setScorer(scorer);
|
|
}
|
|
|
|
@Override
|
|
public void collect(int doc) throws IOException {
|
|
if (collectionRate == 1 || hash.hashCode(doc) % collectionRate == 0) {
|
|
leafCollector.collect(doc);
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public DocIdSetIterator competitiveIterator() throws IOException {
|
|
return leafCollector.competitiveIterator();
|
|
}
|
|
};
|
|
}
|
|
|
|
@Override
|
|
public ScoreMode scoreMode() {
|
|
return collector.scoreMode();
|
|
}
|
|
}
|
|
}
|