/*
* 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.lucene.search.grouping;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NavigableSet;
import java.util.TreeSet;
import org.apache.lucene.search.FieldComparator;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
/**
* Represents a group that is found during the first pass search.
*
* @lucene.experimental
*/
public class SearchGroup<T> {
/** The value that defines this group */
public T groupValue;
/** The sort values used during sorting. These are the
* groupSort field values of the highest rank document
* (by the groupSort) within the group. Can be
* <code>null</code> if <code>fillFields=false</code> had
* been passed to {@link FirstPassGroupingCollector#getTopGroups} */
public Object[] sortValues;
@Override
public String toString() {
return("SearchGroup(groupValue=" + groupValue + " sortValues=" + Arrays.toString(sortValues) + ")");
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
SearchGroup<?> that = (SearchGroup<?>) o;
if (groupValue == null) {
if (that.groupValue != null) {
return false;
}
} else if (!groupValue.equals(that.groupValue)) {
return false;
}
return true;
}
@Override
public int hashCode() {
return groupValue != null ? groupValue.hashCode() : 0;
}
private static class ShardIter<T> {
public final Iterator<SearchGroup<T>> iter;
public final int shardIndex;
public ShardIter(Collection<SearchGroup<T>> shard, int shardIndex) {
this.shardIndex = shardIndex;
iter = shard.iterator();
assert iter.hasNext();
}
public SearchGroup<T> next() {
assert iter.hasNext();
final SearchGroup<T> group = iter.next();
if (group.sortValues == null) {
throw new IllegalArgumentException("group.sortValues is null; you must pass fillFields=true to the first pass collector");
}
return group;
}
@Override
public String toString() {
return "ShardIter(shard=" + shardIndex + ")";
}
}
// Holds all shards currently on the same group
private static class MergedGroup<T> {
// groupValue may be null!
public final T groupValue;
public Object[] topValues;
public final List<ShardIter<T>> shards = new ArrayList<>();
public int minShardIndex;
public boolean processed;
public boolean inQueue;
public MergedGroup(T groupValue) {
this.groupValue = groupValue;
}
// Only for assert
private boolean neverEquals(Object _other) {
if (_other instanceof MergedGroup) {
MergedGroup<?> other = (MergedGroup<?>) _other;
if (groupValue == null) {
assert other.groupValue != null;
} else {
assert !groupValue.equals(other.groupValue);
}
}
return true;
}
@Override
public boolean equals(Object _other) {
// We never have another MergedGroup instance with
// same groupValue
assert neverEquals(_other);
if (_other instanceof MergedGroup) {
MergedGroup<?> other = (MergedGroup<?>) _other;
if (groupValue == null) {
return other == null;
} else {
return groupValue.equals(other);
}
} else {
return false;
}
}
@Override
public int hashCode() {
if (groupValue == null) {
return 0;
} else {
return groupValue.hashCode();
}
}
}
private static class GroupComparator<T> implements Comparator<MergedGroup<T>> {
@SuppressWarnings("rawtypes")
public final FieldComparator[] comparators;
public final int[] reversed;
@SuppressWarnings({"unchecked", "rawtypes"})
public GroupComparator(Sort groupSort) {
final SortField[] sortFields = groupSort.getSort();
comparators = new FieldComparator[sortFields.length];
reversed = new int[sortFields.length];
for (int compIDX = 0; compIDX < sortFields.length; compIDX++) {
final SortField sortField = sortFields[compIDX];
comparators[compIDX] = sortField.getComparator(1, compIDX);
reversed[compIDX] = sortField.getReverse() ? -1 : 1;
}
}
@Override
@SuppressWarnings({"unchecked","rawtypes"})
public int compare(MergedGroup<T> group, MergedGroup<T> other) {
if (group == other) {
return 0;
}
//System.out.println("compare group=" + group + " other=" + other);
final Object[] groupValues = group.topValues;
final Object[] otherValues = other.topValues;
//System.out.println(" groupValues=" + groupValues + " otherValues=" + otherValues);
for (int compIDX = 0;compIDX < comparators.length; compIDX++) {
final int c = reversed[compIDX] * comparators[compIDX].compareValues(groupValues[compIDX],
otherValues[compIDX]);
if (c != 0) {
return c;
}
}
// Tie break by min shard index:
assert group.minShardIndex != other.minShardIndex;
return group.minShardIndex - other.minShardIndex;
}
}
private static class GroupMerger<T> {
private final GroupComparator<T> groupComp;
private final NavigableSet<MergedGroup<T>> queue;
private final Map<T,MergedGroup<T>> groupsSeen;
public GroupMerger(Sort groupSort) {
groupComp = new GroupComparator<>(groupSort);
queue = new TreeSet<>(groupComp);
groupsSeen = new HashMap<>();
}
@SuppressWarnings({"unchecked","rawtypes"})
private void updateNextGroup(int topN, ShardIter<T> shard) {
while(shard.iter.hasNext()) {
final SearchGroup<T> group = shard.next();
MergedGroup<T> mergedGroup = groupsSeen.get(group.groupValue);
final boolean isNew = mergedGroup == null;
//System.out.println(" next group=" + (group.groupValue == null ? "null" : ((BytesRef) group.groupValue).utf8ToString()) + " sort=" + Arrays.toString(group.sortValues));
if (isNew) {
// Start a new group:
//System.out.println(" new");
mergedGroup = new MergedGroup<>(group.groupValue);
mergedGroup.minShardIndex = shard.shardIndex;
assert group.sortValues != null;
mergedGroup.topValues = group.sortValues;
groupsSeen.put(group.groupValue, mergedGroup);
mergedGroup.inQueue = true;
queue.add(mergedGroup);
} else if (mergedGroup.processed) {
// This shard produced a group that we already
// processed; move on to next group...
continue;
} else {
//System.out.println(" old");
boolean competes = false;
for(int compIDX=0;compIDX<groupComp.comparators.length;compIDX++) {
final int cmp = groupComp.reversed[compIDX] * groupComp.comparators[compIDX].compareValues(group.sortValues[compIDX],
mergedGroup.topValues[compIDX]);
if (cmp < 0) {
// Definitely competes
competes = true;
break;
} else if (cmp > 0) {
// Definitely does not compete
break;
} else if (compIDX == groupComp.comparators.length-1) {
if (shard.shardIndex < mergedGroup.minShardIndex) {
competes = true;
}
}
}
//System.out.println(" competes=" + competes);
if (competes) {
// Group's sort changed -- remove & re-insert
if (mergedGroup.inQueue) {
queue.remove(mergedGroup);
}
mergedGroup.topValues = group.sortValues;
mergedGroup.minShardIndex = shard.shardIndex;
queue.add(mergedGroup);
mergedGroup.inQueue = true;
}
}
mergedGroup.shards.add(shard);
break;
}
// Prune un-competitive groups:
while(queue.size() > topN) {
final MergedGroup<T> group = queue.pollLast();
//System.out.println("PRUNE: " + group);
group.inQueue = false;
}
}
public Collection<SearchGroup<T>> merge(List<Collection<SearchGroup<T>>> shards, int offset, int topN) {
final int maxQueueSize = offset + topN;
//System.out.println("merge");
// Init queue:
for(int shardIDX=0;shardIDX<shards.size();shardIDX++) {
final Collection<SearchGroup<T>> shard = shards.get(shardIDX);
if (!shard.isEmpty()) {
//System.out.println(" insert shard=" + shardIDX);
updateNextGroup(maxQueueSize, new ShardIter<>(shard, shardIDX));
}
}
// Pull merged topN groups:
final List<SearchGroup<T>> newTopGroups = new ArrayList<>(topN);
int count = 0;
while(!queue.isEmpty()) {
final MergedGroup<T> group = queue.pollFirst();
group.processed = true;
//System.out.println(" pop: shards=" + group.shards + " group=" + (group.groupValue == null ? "null" : (((BytesRef) group.groupValue).utf8ToString())) + " sortValues=" + Arrays.toString(group.topValues));
if (count++ >= offset) {
final SearchGroup<T> newGroup = new SearchGroup<>();
newGroup.groupValue = group.groupValue;
newGroup.sortValues = group.topValues;
newTopGroups.add(newGroup);
if (newTopGroups.size() == topN) {
break;
}
//} else {
// System.out.println(" skip < offset");
}
// Advance all iters in this group:
for(ShardIter<T> shardIter : group.shards) {
updateNextGroup(maxQueueSize, shardIter);
}
}
if (newTopGroups.isEmpty()) {
return null;
} else {
return newTopGroups;
}
}
}
/** Merges multiple collections of top groups, for example
* obtained from separate index shards. The provided
* groupSort must match how the groups were sorted, and
* the provided SearchGroups must have been computed
* with fillFields=true passed to {@link
* FirstPassGroupingCollector#getTopGroups}.
*
* <p>NOTE: this returns null if the topGroups is empty.
*/
public static <T> Collection<SearchGroup<T>> merge(List<Collection<SearchGroup<T>>> topGroups, int offset, int topN, Sort groupSort) {
if (topGroups.isEmpty()) {
return null;
} else {
return new GroupMerger<T>(groupSort).merge(topGroups, offset, topN);
}
}
}