Reactive collector multi manager

This commit is contained in:
Andrea Cavalli 2021-11-08 23:53:53 +01:00
parent 7c0a4b2af9
commit 1ba4a1866f
5 changed files with 65 additions and 77 deletions

View File

@ -0,0 +1,8 @@
package it.cavallium.dbengine.lucene.collector;
import org.apache.lucene.search.ScoreMode;
public interface CollectorMultiManager {
ScoreMode scoreMode();
}

View File

@ -1,33 +0,0 @@
package it.cavallium.dbengine.lucene.collector;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.LeafCollector;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.ScoreMode;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.Sinks.Many;
public class ReactiveCollector implements Collector {
private final FluxSink<ScoreDoc> scoreDocsSink;
private int shardIndex;
public ReactiveCollector(FluxSink<ScoreDoc> scoreDocsSink) {
this.scoreDocsSink = scoreDocsSink;
}
@Override
public LeafCollector getLeafCollector(LeafReaderContext leafReaderContext) {
return new ReactiveLeafCollector(leafReaderContext, scoreDocsSink, shardIndex);
}
@Override
public ScoreMode scoreMode() {
return ScoreMode.COMPLETE_NO_SCORES;
}
public void setShardIndex(int shardIndex) {
this.shardIndex = shardIndex;
}
}

View File

@ -1,27 +0,0 @@
package it.cavallium.dbengine.lucene.collector;
import java.util.Collection;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.ScoreDoc;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.Sinks.Many;
public class ReactiveCollectorManager implements CollectorManager<Collector, Void> {
private final FluxSink<ScoreDoc> scoreDocsSink;
public ReactiveCollectorManager(FluxSink<ScoreDoc> scoreDocsSink) {
this.scoreDocsSink = scoreDocsSink;
}
@Override
public ReactiveCollector newCollector() {
return new ReactiveCollector(scoreDocsSink);
}
@Override
public Void reduce(Collection<Collector> collection) {
throw new UnsupportedOperationException();
}
}

View File

@ -0,0 +1,52 @@
package it.cavallium.dbengine.lucene.collector;
import java.io.IOException;
import java.util.Collection;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.LeafCollector;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.ScoreMode;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.Sinks.Many;
public class ReactiveCollectorMultiManager implements CollectorMultiManager {
private final FluxSink<ScoreDoc> scoreDocsSink;
public ReactiveCollectorMultiManager(FluxSink<ScoreDoc> scoreDocsSink) {
this.scoreDocsSink = scoreDocsSink;
}
public CollectorManager<Collector, Void> get(int shardIndex) {
return new CollectorManager<>() {
@Override
public Collector newCollector() {
return new Collector() {
@Override
public LeafCollector getLeafCollector(LeafReaderContext leafReaderContext) {
return new ReactiveLeafCollector(leafReaderContext, scoreDocsSink, shardIndex);
}
@Override
public ScoreMode scoreMode() {
return ReactiveCollectorMultiManager.this.scoreMode();
}
};
}
@Override
public Void reduce(Collection<Collector> collection) {
return null;
}
};
}
@Override
public ScoreMode scoreMode() {
return ScoreMode.COMPLETE_NO_SCORES;
}
}

View File

@ -6,28 +6,18 @@ import it.cavallium.dbengine.database.LLKeyScore;
import it.cavallium.dbengine.database.LLUtils;
import it.cavallium.dbengine.database.disk.LLIndexSearchers;
import it.cavallium.dbengine.lucene.LuceneUtils;
import it.cavallium.dbengine.lucene.collector.ReactiveCollectorManager;
import it.cavallium.dbengine.lucene.collector.ReactiveCollectorMultiManager;
import it.cavallium.dbengine.lucene.searcher.LLSearchTransformer.TransformerInput;
import java.util.Queue;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.LockSupport;
import java.util.function.Supplier;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.ScoreDoc;
import org.warp.commonutils.type.ShortNamedThreadFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink.OverflowStrategy;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.publisher.Sinks.EmitResult;
import reactor.core.publisher.Sinks.Many;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;
import reactor.util.concurrent.Queues;
public class UnsortedUnscoredStreamingMultiSearcher implements MultiSearcher {
@ -66,7 +56,7 @@ public class UnsortedUnscoredStreamingMultiSearcher implements MultiSearcher {
var shards = indexSearchers.shards();
Flux<ScoreDoc> scoreDocsFlux = Flux.create(scoreDocsSink -> {
var cm = new ReactiveCollectorManager(scoreDocsSink);
var cmm = new ReactiveCollectorMultiManager(scoreDocsSink);
AtomicInteger runningTasks = new AtomicInteger(0);
@ -76,15 +66,13 @@ public class UnsortedUnscoredStreamingMultiSearcher implements MultiSearcher {
int shardIndex = mutableShardIndex++;
EXECUTOR_SERVICE.execute(() -> {
try {
var collector = cm.newCollector();
assert queryParams.complete() == collector.scoreMode().isExhaustive();
var collector = cmm.get(shardIndex);
assert queryParams.complete() == cmm.scoreMode().isExhaustive();
assert queryParams
.getScoreModeOptional()
.map(scoreMode -> scoreMode == collector.scoreMode())
.map(scoreMode -> scoreMode == cmm.scoreMode())
.orElse(true);
collector.setShardIndex(shardIndex);
shard.search(localQueryParams.query(), collector);
} catch (Throwable e) {
if (!(e instanceof CancellationException)) {