128 lines
4.2 KiB
Java
128 lines
4.2 KiB
Java
package it.cavallium.dbengine.lucene.collector;
|
|
|
|
import static it.cavallium.dbengine.lucene.searcher.PaginationInfo.ALLOW_UNSCORED_PAGINATION_MODE;
|
|
|
|
import it.cavallium.dbengine.lucene.LuceneUtils;
|
|
import it.cavallium.dbengine.lucene.collector.UnscoredCollector;
|
|
import java.io.IOException;
|
|
import java.util.Collection;
|
|
import java.util.List;
|
|
import org.apache.commons.lang3.NotImplementedException;
|
|
import org.apache.lucene.search.Collector;
|
|
import org.apache.lucene.search.CollectorManager;
|
|
import org.apache.lucene.search.FieldDoc;
|
|
import org.apache.lucene.search.IndexSearcher;
|
|
import org.apache.lucene.search.Query;
|
|
import org.apache.lucene.search.ScoreDoc;
|
|
import org.apache.lucene.search.ScoreMode;
|
|
import org.apache.lucene.search.Sort;
|
|
import org.apache.lucene.search.TopDocs;
|
|
import org.apache.lucene.search.TopDocsCollector;
|
|
import org.apache.lucene.search.TopFieldCollector;
|
|
import org.apache.lucene.search.TopFieldDocs;
|
|
import org.apache.lucene.search.TopScoreDocCollector;
|
|
import org.jetbrains.annotations.NotNull;
|
|
|
|
public class TopDocsCollectorMultiManager implements CollectorMultiManager<TopDocs, TopDocs> {
|
|
|
|
private final Sort luceneSort;
|
|
private final int limit;
|
|
private final ScoreDoc after;
|
|
private final int totalHitsThreshold;
|
|
private final boolean allowPagination;
|
|
private final boolean computeScores;
|
|
|
|
private final int topDocsOffset;
|
|
private final int topDocsCount;
|
|
|
|
public TopDocsCollectorMultiManager(Sort luceneSort,
|
|
int limit,
|
|
ScoreDoc after,
|
|
int totalHitsThreshold,
|
|
boolean allowPagination,
|
|
boolean computeScores,
|
|
int topDocsOffset,
|
|
int topDocsCount) {
|
|
this.luceneSort = luceneSort;
|
|
this.limit = limit;
|
|
this.after = after;
|
|
this.totalHitsThreshold = totalHitsThreshold;
|
|
this.allowPagination = allowPagination;
|
|
this.computeScores = computeScores;
|
|
|
|
this.topDocsOffset = topDocsOffset;
|
|
this.topDocsCount = topDocsCount;
|
|
}
|
|
|
|
public CollectorManager<TopDocsCollector<?>, TopDocs> get(@NotNull Query query, IndexSearcher indexSearcher) {
|
|
return new CollectorManager<>() {
|
|
@Override
|
|
public TopDocsCollector<?> newCollector() throws IOException {
|
|
TopDocsCollector<?> collector;
|
|
if (after != null && !allowPagination) {
|
|
throw new IllegalArgumentException("\"allowPagination\" is false, but \"after\" is set");
|
|
}
|
|
if (luceneSort == null) {
|
|
if (after == null) {
|
|
if (computeScores || allowPagination || !ALLOW_UNSCORED_PAGINATION_MODE) {
|
|
collector = TopScoreDocCollector.create(limit, totalHitsThreshold);
|
|
} else {
|
|
collector = new UnscoredCollector(limit);
|
|
}
|
|
} else {
|
|
collector = TopScoreDocCollector.create(limit, after, totalHitsThreshold);
|
|
}
|
|
} else {
|
|
if (after == null) {
|
|
collector = TopFieldCollector.create(luceneSort, limit, totalHitsThreshold);
|
|
} else if (after instanceof FieldDoc afterFieldDoc) {
|
|
collector = TopFieldCollector.create(luceneSort, limit, afterFieldDoc, totalHitsThreshold);
|
|
} else {
|
|
throw new UnsupportedOperationException("GetTopDocs with \"luceneSort\" != null requires \"after\" to be a FieldDoc");
|
|
}
|
|
}
|
|
return collector;
|
|
}
|
|
|
|
@Override
|
|
public TopDocs reduce(Collection<TopDocsCollector<?>> collectors) throws IOException {
|
|
TopDocs[] docsArray;
|
|
boolean needsSort = luceneSort != null;
|
|
boolean needsScores = luceneSort != null && luceneSort.needsScores();
|
|
if (needsSort) {
|
|
docsArray = new TopFieldDocs[collectors.size()];
|
|
} else {
|
|
docsArray = new TopDocs[collectors.size()];
|
|
}
|
|
int i = 0;
|
|
for (TopDocsCollector<?> collector : collectors) {
|
|
docsArray[i] = collector.topDocs();
|
|
i++;
|
|
}
|
|
var merged = LuceneUtils.mergeTopDocs(luceneSort, null, null, docsArray);
|
|
if (needsScores) {
|
|
TopFieldCollector.populateScores(merged.scoreDocs, indexSearcher, query);
|
|
}
|
|
return merged;
|
|
}
|
|
};
|
|
}
|
|
|
|
@Override
|
|
public ScoreMode scoreMode() {
|
|
throw new NotImplementedException();
|
|
}
|
|
|
|
@SuppressWarnings({"SuspiciousToArrayCall", "IfStatementWithIdenticalBranches"})
|
|
@Override
|
|
public TopDocs reduce(List<TopDocs> topDocs) {
|
|
TopDocs[] arr;
|
|
if (luceneSort != null) {
|
|
arr = topDocs.toArray(TopFieldDocs[]::new);
|
|
} else {
|
|
arr = topDocs.toArray(TopDocs[]::new);
|
|
}
|
|
return LuceneUtils.mergeTopDocs(luceneSort, topDocsOffset, topDocsCount, arr);
|
|
}
|
|
}
|