package it.cavallium.dbengine.client; import io.net5.buffer.api.Send; import it.cavallium.dbengine.client.query.ClientQueryParams; import it.cavallium.dbengine.client.query.current.data.Query; import it.cavallium.dbengine.client.query.current.data.QueryParams; import it.cavallium.dbengine.client.query.current.data.TotalHitsCount; import it.cavallium.dbengine.database.LLLuceneIndex; import it.cavallium.dbengine.database.LLSearchResultShard; import it.cavallium.dbengine.database.LLSnapshot; import it.cavallium.dbengine.database.LLTerm; import it.cavallium.dbengine.lucene.collector.Buckets; import it.cavallium.dbengine.lucene.searcher.BucketParams; import it.unimi.dsi.fastutil.doubles.DoubleArrayList; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Set; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.util.function.Tuple2; public class LuceneIndexImpl implements LuceneIndex { private final LLLuceneIndex luceneIndex; private final Indicizer indicizer; public LuceneIndexImpl(LLLuceneIndex luceneIndex, Indicizer indicizer) { this.luceneIndex = luceneIndex; this.indicizer = indicizer; } private LLSnapshot resolveSnapshot(CompositeSnapshot snapshot) { if (snapshot == null) { return null; } else { return snapshot.getSnapshot(luceneIndex); } } @Override public Mono addDocument(T key, U value) { return indicizer .toDocument(key, value) .flatMap(doc -> luceneIndex.addDocument(indicizer.toIndex(key), doc)); } @Override public Mono addDocuments(Flux> entries) { return luceneIndex .addDocuments(entries .flatMap(entry -> indicizer .toDocument(entry.getKey(), entry.getValue()) .map(doc -> Map.entry(indicizer.toIndex(entry.getKey()), doc))) ); } @Override public Mono deleteDocument(T key) { LLTerm id = indicizer.toIndex(key); return luceneIndex.deleteDocument(id); } @Override public Mono updateDocument(T key, @NotNull U value) { return indicizer .toIndexRequest(key, value) .flatMap(doc -> luceneIndex.update(indicizer.toIndex(key), doc)); } @Override public Mono updateDocuments(Flux> entries) { return luceneIndex .updateDocuments(entries .flatMap(entry -> indicizer .toDocument(entry.getKey(), entry.getValue()) .map(doc -> Map.entry(indicizer.toIndex(entry.getKey()), doc))) .collectMap(Entry::getKey, Entry::getValue) ); } @Override public Mono deleteAll() { return luceneIndex.deleteAll(); } @Override public Mono>> moreLikeThis(ClientQueryParams queryParams, T key, U mltDocumentValue) { var mltDocumentFields = indicizer.getMoreLikeThisDocumentFields(key, mltDocumentValue); return luceneIndex .moreLikeThis(resolveSnapshot(queryParams.snapshot()), queryParams.toQueryParams(), indicizer.getKeyFieldName(), mltDocumentFields ) .map(this::mapResults) .single(); } @Override public Mono>> search(ClientQueryParams queryParams) { return luceneIndex .search(resolveSnapshot(queryParams.snapshot()), queryParams.toQueryParams(), indicizer.getKeyFieldName() ) .map(this::mapResults) .single(); } @Override public Mono computeBuckets(@Nullable CompositeSnapshot snapshot, @NotNull List query, @Nullable Query normalizationQuery, BucketParams bucketParams) { return luceneIndex .computeBuckets(resolveSnapshot(snapshot), query, normalizationQuery, bucketParams ) .single(); } private Hits> mapResults(LLSearchResultShard llSearchResult) { var scoresWithKeysFlux = llSearchResult .results() .map(hit -> new HitKey<>(indicizer.getKey(hit.key()), hit.score())); return new Hits<>(scoresWithKeysFlux, llSearchResult.totalHitsCount(), llSearchResult::close); } @Override public Mono count(@Nullable CompositeSnapshot snapshot, Query query) { return this .search(ClientQueryParams.builder().snapshot(snapshot).query(query).limit(0).build()) .single() .map(searchResultKeys -> { try (searchResultKeys) { return searchResultKeys.totalHitsCount(); } }); } @Override public boolean isLowMemoryMode() { return luceneIndex.isLowMemoryMode(); } @Override public Mono close() { return luceneIndex.close(); } /** * Flush writes to disk */ @Override public Mono flush() { return luceneIndex.flush(); } /** * Refresh index searcher */ @Override public Mono refresh(boolean force) { return luceneIndex.refresh(force); } @Override public Mono takeSnapshot() { return luceneIndex.takeSnapshot(); } @Override public Mono releaseSnapshot(LLSnapshot snapshot) { return luceneIndex.releaseSnapshot(snapshot); } }