package org.apache.lucene.search;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.apache.lucene.index.LeafReaderContext;
public class MultiCollector implements Collector {
public static Collector wrap(Collector... collectors) {
return wrap(Arrays.asList(collectors));
}
public static Collector wrap(Iterable<? extends Collector> collectors) {
int n = 0;
for (Collector c : collectors) {
if (c != null) {
n++;
}
}
if (n == 0) {
throw new IllegalArgumentException("At least 1 collector must not be null");
} else if (n == 1) {
Collector col = null;
for (Collector c : collectors) {
if (c != null) {
col = c;
break;
}
}
return col;
} else {
Collector[] colls = new Collector[n];
n = 0;
for (Collector c : collectors) {
if (c != null) {
colls[n++] = c;
}
}
return new MultiCollector(colls);
}
}
private final boolean cacheScores;
private final Collector[] collectors;
private MultiCollector(Collector... collectors) {
this.collectors = collectors;
int numNeedsScores = 0;
for (Collector collector : collectors) {
if (collector.scoreMode().needsScores()) {
numNeedsScores += 1;
}
}
this.cacheScores = numNeedsScores >= 2;
}
@Override
public ScoreMode scoreMode() {
ScoreMode scoreMode = null;
for (Collector collector : collectors) {
if (scoreMode == null) {
scoreMode = collector.scoreMode();
} else if (scoreMode != collector.scoreMode()) {
return ScoreMode.COMPLETE;
}
}
return scoreMode;
}
@Override
public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
final List<LeafCollector> leafCollectors = new ArrayList<>(collectors.length);
for (Collector collector : collectors) {
final LeafCollector leafCollector;
try {
leafCollector = collector.getLeafCollector(context);
} catch (CollectionTerminatedException e) {
continue;
}
leafCollectors.add(leafCollector);
}
switch (leafCollectors.size()) {
case 0:
throw new CollectionTerminatedException();
case 1:
return leafCollectors.get(0);
default:
return new MultiLeafCollector(leafCollectors, cacheScores, scoreMode() == ScoreMode.TOP_SCORES);
}
}
private static class MultiLeafCollector implements LeafCollector {
private final boolean cacheScores;
private final LeafCollector[] collectors;
private final float[] minScores;
private final boolean skipNonCompetitiveScores;
private MultiLeafCollector(List<LeafCollector> collectors, boolean cacheScores, boolean skipNonCompetitive) {
this.collectors = collectors.toArray(new LeafCollector[collectors.size()]);
this.cacheScores = cacheScores;
this.skipNonCompetitiveScores = skipNonCompetitive;
this.minScores = this.skipNonCompetitiveScores ? new float[this.collectors.length] : null;
}
@Override
public void setScorer(Scorable scorer) throws IOException {
if (cacheScores) {
scorer = new ScoreCachingWrappingScorer(scorer);
}
if (skipNonCompetitiveScores) {
for (int i = 0; i < collectors.length; ++i) {
final LeafCollector c = collectors[i];
if (c != null) {
c.setScorer(new MinCompetitiveScoreAwareScorable(scorer, i, minScores));
}
}
} else {
scorer = new FilterScorable(scorer) {
@Override
public void setMinCompetitiveScore(float minScore) throws IOException {
}
};
for (int i = 0; i < collectors.length; ++i) {
final LeafCollector c = collectors[i];
if (c != null) {
c.setScorer(scorer);
}
}
}
}
@Override
public void collect(int doc) throws IOException {
for (int i = 0; i < collectors.length; i++) {
final LeafCollector collector = collectors[i];
if (collector != null) {
try {
collector.collect(doc);
} catch (CollectionTerminatedException e) {
collectors[i] = null;
if (allCollectorsTerminated()) {
throw new CollectionTerminatedException();
}
}
}
}
}
private boolean allCollectorsTerminated() {
for (int i = 0; i < collectors.length; i++) {
if (collectors[i] != null) {
return false;
}
}
return true;
}
}
final static class MinCompetitiveScoreAwareScorable extends FilterScorable {
private final int idx;
private final float[] minScores;
MinCompetitiveScoreAwareScorable(Scorable in, int idx, float[] minScores) {
super(in);
this.idx = idx;
this.minScores = minScores;
}
@Override
public void setMinCompetitiveScore(float minScore) throws IOException {
if (minScore > minScores[idx]) {
minScores[idx] = minScore;
in.setMinCompetitiveScore(minScore());
}
}
private float minScore() {
float min = Float.MAX_VALUE;
for (int i = 0; i < minScores.length; i++) {
if (minScores[i] < min) {
min = minScores[i];
}
}
return min;
}
}
}