65 lines
2.1 KiB
Java
65 lines
2.1 KiB
Java
package it.cavallium.dbengine.lucene.collector;
|
|
|
|
import it.cavallium.dbengine.lucene.LuceneUtils;
|
|
import it.cavallium.dbengine.lucene.collector.TotalHitCountCollectorManager.TimeLimitingTotalHitCountCollector;
|
|
import java.io.IOException;
|
|
import java.time.Duration;
|
|
import java.util.Collection;
|
|
import org.apache.lucene.index.LeafReaderContext;
|
|
import org.apache.lucene.search.Collector;
|
|
import org.apache.lucene.search.CollectorManager;
|
|
import org.apache.lucene.search.LeafCollector;
|
|
import org.apache.lucene.search.ScoreMode;
|
|
import org.apache.lucene.search.TotalHitCountCollector;
|
|
|
|
public class TotalHitCountCollectorManager implements CollectorManager<TimeLimitingTotalHitCountCollector, Long> {
|
|
|
|
private final Duration timeout;
|
|
|
|
public TotalHitCountCollectorManager(Duration timeout) {
|
|
this.timeout = timeout;
|
|
}
|
|
|
|
@Override
|
|
public TimeLimitingTotalHitCountCollector newCollector() {
|
|
var totalHitCountCollector = new TotalHitCountCollector();
|
|
var timeLimitingCollector = LuceneUtils.withTimeout(totalHitCountCollector, timeout);
|
|
return new TimeLimitingTotalHitCountCollector(totalHitCountCollector, timeLimitingCollector);
|
|
}
|
|
|
|
@Override
|
|
public Long reduce(Collection<TimeLimitingTotalHitCountCollector> collectors) {
|
|
long totalHits = 0;
|
|
for (var collector : collectors) {
|
|
totalHits += collector.getTotalHits();
|
|
}
|
|
return totalHits;
|
|
}
|
|
|
|
public static final class TimeLimitingTotalHitCountCollector implements Collector {
|
|
|
|
private final TotalHitCountCollector totalHitCountCollector;
|
|
private final Collector timeLimitingCollector;
|
|
|
|
private TimeLimitingTotalHitCountCollector(TotalHitCountCollector totalHitCountCollector,
|
|
Collector timeLimitingCollector) {
|
|
this.totalHitCountCollector = totalHitCountCollector;
|
|
this.timeLimitingCollector = timeLimitingCollector;
|
|
}
|
|
|
|
@Override
|
|
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
|
|
return timeLimitingCollector.getLeafCollector(context);
|
|
}
|
|
|
|
@Override
|
|
public ScoreMode scoreMode() {
|
|
return timeLimitingCollector.scoreMode();
|
|
}
|
|
|
|
public long getTotalHits() {
|
|
return totalHitCountCollector.getTotalHits();
|
|
}
|
|
}
|
|
}
|