CavalliumDBEngine/src/main/java/it/cavallium/dbengine/lucene/searcher/StandardSearcher.java

package it.cavallium.dbengine.lucene.searcher;

import static java.util.Objects.requireNonNull;

import it.cavallium.dbengine.client.query.current.data.TotalHitsCount;
import it.cavallium.dbengine.database.LLKeyScore;
import it.cavallium.dbengine.database.LLUtils;
import it.cavallium.dbengine.database.disk.LLIndexSearchers;
import it.cavallium.dbengine.lucene.LuceneCloseable;
import it.cavallium.dbengine.lucene.LuceneUtils;
import it.cavallium.dbengine.utils.DBException;
import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.stream.Stream;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.ScoreDoc;
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.Nullable;

public class StandardSearcher implements MultiSearcher {

	protected static final Logger LOG = LogManager.getLogger(StandardSearcher.class);

	public StandardSearcher() {
	}

	@Override
	public LuceneSearchResult collectMulti(LLIndexSearchers indexSearchers,
			LocalQueryParams queryParams,
			@Nullable String keyFieldName,
			GlobalQueryRewrite transformer) {
		if (transformer != GlobalQueryRewrite.NO_REWRITE) {
			return LuceneUtils.rewriteMulti(this, indexSearchers, queryParams, keyFieldName, transformer);
		}
		// Search results
		var fullDocs = this.search(indexSearchers.shards(), queryParams);
		// Compute the results
		return this.computeResults(fullDocs, indexSearchers, keyFieldName, queryParams);
	}

	/**
	 * Search effectively the raw results
	 */
	@SuppressWarnings({"rawtypes"})
	private TopDocs search(Collection<IndexSearcher> indexSearchers, LocalQueryParams queryParams) {
		var totalHitsThreshold = queryParams.getTotalHitsThresholdInt();
		CollectorManager<? extends TopDocsCollector<?>, ? extends TopDocs> sharedManager;
		if (queryParams.isSorted() && !queryParams.isSortedByScore()) {
			sharedManager = TopFieldCollector.createSharedManager(queryParams.sort(),
					queryParams.limitInt(), null, totalHitsThreshold);
		} else {
			sharedManager = TopScoreDocCollector.createSharedManager(queryParams.limitInt(), null, totalHitsThreshold);
		};
		var collectors = indexSearchers.stream().map(shard -> {
			try {
				TopDocsCollector<?> collector;
				collector = sharedManager.newCollector();
				assert queryParams.computePreciseHitsCount() == null || (queryParams.computePreciseHitsCount() == collector
						.scoreMode()
						.isExhaustive());

				shard.search(queryParams.query(), LuceneUtils.withTimeout(collector, queryParams.timeout()));
				return collector;
			} catch (IOException e) {
				throw new DBException(e);
			}
		}).toList();

		try {
			if (collectors.size() <= 1) {
				//noinspection unchecked
				return sharedManager.reduce((List) collectors);
			} else if (queryParams.isSorted() && !queryParams.isSortedByScore()) {
				final TopFieldDocs[] topDocs = new TopFieldDocs[collectors.size()];
				int i = 0;
				for (var collector : collectors) {
					var topFieldDocs = ((TopFieldCollector) collector).topDocs();
					for (ScoreDoc scoreDoc : topFieldDocs.scoreDocs) {
						scoreDoc.shardIndex = i;
					}
					topDocs[i++] = topFieldDocs;
				}
				return TopDocs.merge(requireNonNull(queryParams.sort()), 0, queryParams.limitInt(), topDocs);
			} else {
				final TopDocs[] topDocs = new TopDocs[collectors.size()];
				int i = 0;
				for (var collector : collectors) {
					var topScoreDocs = collector.topDocs();
					for (ScoreDoc scoreDoc : topScoreDocs.scoreDocs) {
						scoreDoc.shardIndex = i;
					}
					topDocs[i++] = topScoreDocs;
				}
				return TopDocs.merge(0, queryParams.limitInt(), topDocs);
			}
		} catch (IOException ex) {
			throw new DBException(ex);
		}
	}

	/**
	 * Compute the results, extracting useful data
	 */
	private LuceneSearchResult computeResults(TopDocs data,
			LLIndexSearchers indexSearchers,
			String keyFieldName,
			LocalQueryParams queryParams) {
		var totalHitsCount = LuceneUtils.convertTotalHitsCount(data.totalHits);

		Stream<LLKeyScore> hitsStream = LuceneUtils
				.convertHits(Stream.of(data.scoreDocs),
						indexSearchers.shards(), keyFieldName
				)
				.skip(queryParams.offsetLong())
				.limit(queryParams.limitLong());

		return new MyLuceneSearchResult(totalHitsCount, hitsStream, indexSearchers);
	}

	@Override
	public String getName() {
		return "standard";
	}

	private static class MyLuceneSearchResult extends LuceneSearchResult implements LuceneCloseable {

		private final LLIndexSearchers indexSearchers;

		public MyLuceneSearchResult(TotalHitsCount totalHitsCount,
				Stream<LLKeyScore> hitsStream,
				LLIndexSearchers indexSearchers) {
			super(totalHitsCount, hitsStream);
			this.indexSearchers = indexSearchers;
		}

		@Override
		protected void onClose() {
			try {
				indexSearchers.close();
			} catch (Throwable e) {
				LOG.error("Can't close index searchers", e);
			}
			super.onClose();
		}
	}
}