/**
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.solr.search;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.*;
import org.apache.solr.search.function.DocValues;
import org.apache.solr.search.function.ValueSource;
import java.io.IOException;
import java.util.*;
public class MultiCollector extends Collector {
final Collector[] collectors;
final boolean acceptsDocsOutOfOrder;
public static Collector wrap(List<? extends Collector> collectors) {
return collectors.size() == 1 ? collectors.get(0) : new MultiCollector(collectors);
}
public static Collector[] subCollectors(Collector collector) {
if (collector instanceof MultiCollector)
return ((MultiCollector)collector).collectors;
return new Collector[]{collector};
}
public MultiCollector(List<? extends Collector> collectors) {
this(collectors.toArray(new Collector[collectors.size()]));
}
public MultiCollector(Collector[] collectors) {
this.collectors = collectors;
boolean acceptsDocsOutOfOrder = true;
for (Collector collector : collectors) {
if (collector.acceptsDocsOutOfOrder() == false) {
acceptsDocsOutOfOrder = false;
break;
}
}
this.acceptsDocsOutOfOrder = acceptsDocsOutOfOrder;
}
@Override
public void setScorer(Scorer scorer) throws IOException {
for (Collector collector : collectors)
collector.setScorer(scorer);
}
@Override
public void collect(int doc) throws IOException {
for (Collector collector : collectors)
collector.collect(doc);
}
@Override
public void setNextReader(IndexReader reader, int docBase) throws IOException {
for (Collector collector : collectors)
collector.setNextReader(reader, docBase);
}
@Override
public boolean acceptsDocsOutOfOrder() {
return acceptsDocsOutOfOrder;
}
}
class SearchGroup {
public MutableValue groupValue;
int matches;
int topDoc;
// float topDocScore; // currently unused
int comparatorSlot;
// currently only used when sort != sort.group
FieldComparator[] sortGroupComparators;
int[] sortGroupReversed;
/***
@Override
public int hashCode() {
return super.hashCode();
}
@Override
public boolean equals(Object obj) {
return groupValue.equalsSameType(((SearchGroup)obj).groupValue);
}
***/
}
/** Finds the top set of groups, grouped by groupByVS when sort == group.sort */
class TopGroupCollector extends Collector {
final int nGroups;
final HashMap<MutableValue, SearchGroup> groupMap;
TreeSet<SearchGroup> orderedGroups;
final ValueSource vs;
final Map context;
final FieldComparator[] comparators;
final int[] reversed;
DocValues docValues;
DocValues.ValueFiller filler;
MutableValue mval;
Scorer scorer;
int docBase;
int spareSlot;
int matches;
public TopGroupCollector(ValueSource groupByVS, Map vsContext, Sort sort, int nGroups) throws IOException {
this.vs = groupByVS;
this.context = vsContext;
this.nGroups = nGroups;
SortField[] sortFields = sort.getSort();
this.comparators = new FieldComparator[sortFields.length];
this.reversed = new int[sortFields.length];
for (int i = 0; i < sortFields.length; i++) {
SortField sortField = sortFields[i];
reversed[i] = sortField.getReverse() ? -1 : 1;
// use nGroups + 1 so we have a spare slot to use for comparing (tracked by this.spareSlot)
comparators[i] = sortField.getComparator(nGroups + 1, i);
}
this.spareSlot = nGroups;
this.groupMap = new HashMap<MutableValue, SearchGroup>(nGroups);
}
@Override
public void setScorer(Scorer scorer) throws IOException {
this.scorer = scorer;
for (FieldComparator fc : comparators)
fc.setScorer(scorer);
}
@Override
public void collect(int doc) throws IOException {
matches++;
filler.fillValue(doc);
SearchGroup group = groupMap.get(mval);
if (group == null) {
int num = groupMap.size();
if (groupMap.size() < nGroups) {
SearchGroup sg = new SearchGroup();
sg.groupValue = mval.duplicate();
sg.comparatorSlot = num++;
sg.matches = 1;
sg.topDoc = docBase + doc;
// sg.topDocScore = scorer.score();
for (FieldComparator fc : comparators)
fc.copy(sg.comparatorSlot, doc);
groupMap.put(sg.groupValue, sg);
return;
}
if (orderedGroups == null) {
buildSet();
}
for (int i = 0;; i++) {
final int c = reversed[i] * comparators[i].compareBottom(doc);
if (c < 0) {
// Definitely not competitive.
return;
} else if (c > 0) {
// Definitely competitive.
break;
} else if (i == comparators.length - 1) {
// Here c=0. If we're at the last comparator, this doc is not
// competitive, since docs are visited in doc Id order, which means
// this doc cannot compete with any other document in the queue.
return;
}
}
// remove current smallest group
SearchGroup smallest = orderedGroups.pollLast();
groupMap.remove(smallest.groupValue);
// reuse the removed SearchGroup
smallest.groupValue.copy(mval);
smallest.matches = 1;
smallest.topDoc = docBase + doc;
// smallest.topDocScore = scorer.score();
for (FieldComparator fc : comparators)
fc.copy(smallest.comparatorSlot, doc);
groupMap.put(smallest.groupValue, smallest);
orderedGroups.add(smallest);
for (FieldComparator fc : comparators)
fc.setBottom(orderedGroups.last().comparatorSlot);
return;
}
//
// update existing group
//
group.matches++; // TODO: these aren't valid if the group is every discarded then re-added. keep track if there have been discards?
for (int i = 0;; i++) {
FieldComparator fc = comparators[i];
fc.copy(spareSlot, doc);
final int c = reversed[i] * fc.compare(group.comparatorSlot, spareSlot);
if (c < 0) {
// Definitely not competitive.
return;
} else if (c > 0) {
// Definitely competitive.
// Set remaining comparators
for (int j=i+1; j<comparators.length; j++)
comparators[j].copy(spareSlot, doc);
break;
} else if (i == comparators.length - 1) {
// Here c=0. If we're at the last comparator, this doc is not
// competitive, since docs are visited in doc Id order, which means
// this doc cannot compete with any other document in the queue.
return;
}
}
// remove before updating the group since lookup is done via comparators
// TODO: optimize this
if (orderedGroups != null)
orderedGroups.remove(group);
group.topDoc = docBase + doc;
// group.topDocScore = scorer.score();
int tmp = spareSlot; spareSlot = group.comparatorSlot; group.comparatorSlot=tmp; // swap slots
// re-add the changed group
if (orderedGroups != null)
orderedGroups.add(group);
}
void buildSet() {
Comparator<SearchGroup> comparator = new Comparator<SearchGroup>() {
public int compare(SearchGroup o1, SearchGroup o2) {
for (int i = 0;; i++) {
FieldComparator fc = comparators[i];
int c = reversed[i] * fc.compare(o1.comparatorSlot, o2.comparatorSlot);
if (c != 0) {
return c;
} else if (i == comparators.length - 1) {
return o1.topDoc - o2.topDoc;
}
}
}
};
orderedGroups = new TreeSet<SearchGroup>(comparator);
orderedGroups.addAll(groupMap.values());
if (orderedGroups.size() == 0) return;
for (FieldComparator fc : comparators)
fc.setBottom(orderedGroups.last().comparatorSlot);
}
@Override
public void setNextReader(IndexReader reader, int docBase) throws IOException {
this.docBase = docBase;
docValues = vs.getValues(context, reader);
filler = docValues.getValueFiller();
mval = filler.getValue();
for (FieldComparator fc : comparators)
fc.setNextReader(reader, docBase);
}
@Override
public boolean acceptsDocsOutOfOrder() {
return false;
}
public int getMatches() {
return matches;
}
}
/**
* This class allows a different sort within a group than what is used between groups.
* Sorting between groups is done by the sort value of the first (highest ranking)
* document in that group.
*/
class TopGroupSortCollector extends TopGroupCollector {
IndexReader reader;
Sort groupSort;
public TopGroupSortCollector(ValueSource groupByVS, Map vsContext, Sort sort, Sort groupSort, int nGroups) throws IOException {
super(groupByVS, vsContext, sort, nGroups);
this.groupSort = groupSort;
}
void constructComparators(FieldComparator[] comparators, int[] reversed, SortField[] sortFields, int size) throws IOException {
for (int i = 0; i < sortFields.length; i++) {
SortField sortField = sortFields[i];
reversed[i] = sortField.getReverse() ? -1 : 1;
comparators[i] = sortField.getComparator(size, i);
if (scorer != null) comparators[i].setScorer(scorer);
if (reader != null) comparators[i].setNextReader(reader, docBase);
}
}
@Override
public void setScorer(Scorer scorer) throws IOException {
super.setScorer(scorer);
for (SearchGroup searchGroup : groupMap.values()) {
for (FieldComparator fc : searchGroup.sortGroupComparators) {
fc.setScorer(scorer);
}
}
}
@Override
public void collect(int doc) throws IOException {
matches++;
filler.fillValue(doc);
SearchGroup group = groupMap.get(mval);
if (group == null) {
int num = groupMap.size();
if (groupMap.size() < nGroups) {
SearchGroup sg = new SearchGroup();
SortField[] sortGroupFields = groupSort.getSort();
sg.sortGroupComparators = new FieldComparator[sortGroupFields.length];
sg.sortGroupReversed = new int[sortGroupFields.length];
constructComparators(sg.sortGroupComparators, sg.sortGroupReversed, sortGroupFields, 1);
sg.groupValue = mval.duplicate();
sg.comparatorSlot = num++;
sg.matches = 1;
sg.topDoc = docBase + doc;
// sg.topDocScore = scorer.score();
for (FieldComparator fc : comparators)
fc.copy(sg.comparatorSlot, doc);
for (FieldComparator fc : sg.sortGroupComparators) {
fc.copy(0, doc);
fc.setBottom(0);
}
groupMap.put(sg.groupValue, sg);
return;
}
if (orderedGroups == null) {
buildSet();
}
SearchGroup leastSignificantGroup = orderedGroups.last();
for (int i = 0;; i++) {
final int c = leastSignificantGroup.sortGroupReversed[i] * leastSignificantGroup.sortGroupComparators[i].compareBottom(doc);
if (c < 0) {
// Definitely not competitive.
return;
} else if (c > 0) {
// Definitely competitive.
break;
} else if (i == leastSignificantGroup.sortGroupComparators.length - 1) {
// Here c=0. If we're at the last comparator, this doc is not
// competitive, since docs are visited in doc Id order, which means
// this doc cannot compete with any other document in the queue.
return;
}
}
// remove current smallest group
SearchGroup smallest = orderedGroups.pollLast();
groupMap.remove(smallest.groupValue);
// reuse the removed SearchGroup
smallest.groupValue.copy(mval);
smallest.matches = 1;
smallest.topDoc = docBase + doc;
// smallest.topDocScore = scorer.score();
for (FieldComparator fc : comparators)
fc.copy(smallest.comparatorSlot, doc);
for (FieldComparator fc : smallest.sortGroupComparators) {
fc.copy(0, doc);
fc.setBottom(0);
}
groupMap.put(smallest.groupValue, smallest);
orderedGroups.add(smallest);
for (FieldComparator fc : comparators)
fc.setBottom(orderedGroups.last().comparatorSlot);
for (FieldComparator fc : smallest.sortGroupComparators)
fc.setBottom(0);
return;
}
//
// update existing group
//
group.matches++; // TODO: these aren't valid if the group is every discarded then re-added. keep track if there have been discards?
for (int i = 0;; i++) {
FieldComparator fc = group.sortGroupComparators[i];
final int c = group.sortGroupReversed[i] * fc.compareBottom(doc);
if (c < 0) {
// Definitely not competitive.
return;
} else if (c > 0) {
// Definitely competitive.
// Set remaining comparators
for (int j = 0; j < group.sortGroupComparators.length; j++) {
group.sortGroupComparators[j].copy(0, doc);
group.sortGroupComparators[j].setBottom(0);
}
for (FieldComparator comparator : comparators) comparator.copy(spareSlot, doc);
break;
} else if (i == group.sortGroupComparators.length - 1) {
// Here c=0. If we're at the last comparator, this doc is not
// competitive, since docs are visited in doc Id order, which means
// this doc cannot compete with any other document in the queue.
return;
}
}
// remove before updating the group since lookup is done via comparators
// TODO: optimize this
if (orderedGroups != null)
orderedGroups.remove(group);
group.topDoc = docBase + doc;
// group.topDocScore = scorer.score();
int tmp = spareSlot; spareSlot = group.comparatorSlot; group.comparatorSlot=tmp; // swap slots
// re-add the changed group
if (orderedGroups != null)
orderedGroups.add(group);
}
@Override
public void setNextReader(IndexReader reader, int docBase) throws IOException {
super.setNextReader(reader, docBase);
this.reader = reader;
for (SearchGroup searchGroup : groupMap.values()) {
for (FieldComparator fc : searchGroup.sortGroupComparators) {
fc.setNextReader(reader, docBase);
}
}
}
}
class Phase2GroupCollector extends Collector {
final HashMap<MutableValue, SearchGroupDocs> groupMap;
final ValueSource vs;
final Map context;
DocValues docValues;
DocValues.ValueFiller filler;
MutableValue mval;
Scorer scorer;
int docBase;
// TODO: may want to decouple from the phase1 collector
public Phase2GroupCollector(TopGroupCollector topGroups, ValueSource groupByVS, Map vsContext, Sort sort, int docsPerGroup, boolean getScores) throws IOException {
boolean getSortFields = false;
groupMap = new HashMap<MutableValue, SearchGroupDocs>(topGroups.groupMap.size());
for (SearchGroup group : topGroups.groupMap.values()) {
SearchGroupDocs groupDocs = new SearchGroupDocs();
groupDocs.groupValue = group.groupValue;
groupDocs.collector = TopFieldCollector.create(sort, docsPerGroup, getSortFields, getScores, getScores, true);
groupMap.put(groupDocs.groupValue, groupDocs);
}
this.vs = groupByVS;
this.context = vsContext;
}
@Override
public void setScorer(Scorer scorer) throws IOException {
this.scorer = scorer;
for (SearchGroupDocs group : groupMap.values())
group.collector.setScorer(scorer);
}
@Override
public void collect(int doc) throws IOException {
filler.fillValue(doc);
SearchGroupDocs group = groupMap.get(mval);
if (group == null) return;
group.matches++;
group.collector.collect(doc);
}
@Override
public void setNextReader(IndexReader reader, int docBase) throws IOException {
this.docBase = docBase;
docValues = vs.getValues(context, reader);
filler = docValues.getValueFiller();
mval = filler.getValue();
for (SearchGroupDocs group : groupMap.values())
group.collector.setNextReader(reader, docBase);
}
@Override
public boolean acceptsDocsOutOfOrder() {
return false;
}
}
// TODO: merge with SearchGroup or not?
// ad: don't need to build a new hashmap
// disad: blows up the size of SearchGroup if we need many of them, and couples implementations
class SearchGroupDocs {
public MutableValue groupValue;
int matches;
TopFieldCollector collector;
}