From 5f6dfac1da1c61bc48a89462f1afe34f10881a9e Mon Sep 17 00:00:00 2001
From: Andrea Cavalli <andrea@cavallium.it>
Date: Mon, 21 Mar 2022 15:25:26 +0100
Subject: [PATCH] Respect limit

---
 .../dbengine/client/LuceneIndexImpl.java      | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/main/java/it/cavallium/dbengine/client/LuceneIndexImpl.java b/src/main/java/it/cavallium/dbengine/client/LuceneIndexImpl.java
index d3c26ee..27e2b22 100644
--- a/src/main/java/it/cavallium/dbengine/client/LuceneIndexImpl.java
+++ b/src/main/java/it/cavallium/dbengine/client/LuceneIndexImpl.java
@@ -105,7 +105,7 @@ public class LuceneIndexImpl<T, U> implements LuceneIndex<T, U> {
 						mltDocumentFields
 				)
 				.collectList()
-				.flatMap(LuceneIndexImpl::mergeResults)
+				.flatMap(shards -> mergeResults(queryParams, shards))
 				.map(this::mapResults)
 				.single();
 	}
@@ -118,7 +118,7 @@ public class LuceneIndexImpl<T, U> implements LuceneIndex<T, U> {
 						indicizer.getKeyFieldName()
 				)
 				.collectList()
-				.flatMap(LuceneIndexImpl::mergeResults)
+				.flatMap(shards -> mergeResults(queryParams, shards))
 				.map(this::mapResults)
 				.single();
 	}
@@ -193,7 +193,8 @@ public class LuceneIndexImpl<T, U> implements LuceneIndex<T, U> {
 		return luceneIndex.releaseSnapshot(snapshot);
 	}
 
-	private static Mono<LLSearchResultShard> mergeResults(List<LLSearchResultShard> shards) {
+	private static Mono<LLSearchResultShard> mergeResults(ClientQueryParams queryParams,
+			List<LLSearchResultShard> shards) {
 		if (shards.size() == 0) {
 			return Mono.empty();
 		} else if (shards.size() == 1) {
@@ -209,17 +210,19 @@ public class LuceneIndexImpl<T, U> implements LuceneIndex<T, U> {
 				} else {
 					count = LuceneUtils.sum(count, shard.totalHitsCount());
 				}
-				results.add(shard.results());
+				var maxLimit = queryParams.offset() + queryParams.limit();
+				results.add(shard.results().take(maxLimit, true));
 				resources.add(shard);
 			}
 			Objects.requireNonNull(count);
-			var resultsFlux = Flux.zip(results, parts -> {
-				var arr = new ArrayList<LLKeyScore>(parts.length);
-				for (Object part : parts) {
-					arr.add((LLKeyScore) part);
-				}
-				return arr;
-			}).concatMapIterable(list -> list);
+			Flux<LLKeyScore> resultsFlux;
+			if (results.size() == 0) {
+				resultsFlux = Flux.empty();
+			} else if (results.size() == 1) {
+				resultsFlux = results.get(0);
+			} else {
+				resultsFlux = Flux.merge(results);
+			}
 			return new LLSearchResultShard(resultsFlux, count, () -> {
 				for (Resource<?> resource : resources) {
 					resource.close();