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

package it.cavallium.dbengine.lucene.searcher;

import io.net5.buffer.api.Send;
import it.cavallium.dbengine.database.LLUtils;
import it.cavallium.dbengine.database.disk.LLIndexSearchers;
import it.cavallium.dbengine.lucene.searcher.LLSearchTransformer.TransformerInput;
import java.io.Closeable;
import java.io.IOException;
import reactor.core.publisher.Mono;

public class AdaptiveMultiSearcher implements MultiSearcher, Closeable {

	private static final MultiSearcher count
			= new UnsortedUnscoredSimpleMultiSearcher(new CountLocalSearcher());

	private static final MultiSearcher scoredPaged = new ScoredPagedMultiSearcher();

	private static final MultiSearcher unsortedUnscoredPaged
			= new UnsortedUnscoredSimpleMultiSearcher(new PagedLocalSearcher());

	private static final MultiSearcher unsortedUnscoredContinuous
			= new UnsortedUnscoredStreamingMultiSearcher();

	private final UnsortedScoredFullMultiSearcher unsortedScoredFull;

	private final SortedScoredFullMultiSearcher sortedScoredFull;

	public AdaptiveMultiSearcher() throws IOException {
		unsortedScoredFull = new UnsortedScoredFullMultiSearcher();
		sortedScoredFull = new SortedScoredFullMultiSearcher();
	}

	@Override
	public Mono<Send<LuceneSearchResult>> collectMulti(Mono<Send<LLIndexSearchers>> indexSearchersMono,
			LocalQueryParams queryParams,
			String keyFieldName,
			LLSearchTransformer transformer) {
		if (transformer == LLSearchTransformer.NO_TRANSFORMATION) {
			return transformedCollectMulti(indexSearchersMono, queryParams, keyFieldName, transformer);
		} else {
			return LLUtils.usingSendResource(indexSearchersMono, indexSearchers -> transformer
					.transform(Mono.fromCallable(() -> new TransformerInput(indexSearchers, queryParams)))
					.flatMap(queryParams2 -> this
							.transformedCollectMulti(indexSearchersMono, queryParams2, keyFieldName, LLSearchTransformer.NO_TRANSFORMATION)),
					true);
		}
	}

	public Mono<Send<LuceneSearchResult>> transformedCollectMulti(Mono<Send<LLIndexSearchers>> indexSearchersMono,
			LocalQueryParams queryParams,
			String keyFieldName,
			LLSearchTransformer transformer) {
		// offset + limit
		long realLimit = queryParams.offsetLong() + queryParams.limitLong();

		return LLUtils.usingSendResource(indexSearchersMono, indexSearchers -> {
			if (queryParams.limitLong() == 0) {
				return count.collectMulti(indexSearchersMono, queryParams, keyFieldName, transformer);
			} else if (queryParams.isSorted() || queryParams.needsScores()) {
				if (realLimit <= (long) queryParams.pageLimits().getPageLimit(0)) {
					return scoredPaged.collectMulti(indexSearchersMono, queryParams, keyFieldName, transformer);
				} else {
					if ((queryParams.isSorted() && !queryParams.isSortedByScore())) {
						return sortedScoredFull.collectMulti(indexSearchersMono, queryParams, keyFieldName, transformer);
					} else {
						return unsortedScoredFull.collectMulti(indexSearchersMono, queryParams, keyFieldName, transformer);
					}
				}
			} else if (realLimit <= (long) queryParams.pageLimits().getPageLimit(0)) {
				// Run single-page searches using the paged multi searcher
				return unsortedUnscoredPaged.collectMulti(indexSearchersMono, queryParams, keyFieldName, transformer);
			} else {
				// Run large/unbounded searches using the continuous multi searcher
				return unsortedUnscoredContinuous.collectMulti(indexSearchersMono, queryParams, keyFieldName, transformer);
			}
		}, true);
	}

	@Override
	public void close() throws IOException {
		sortedScoredFull.close();
		unsortedScoredFull.close();
	}

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