package org.infinispan.query.clustered;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map.Entry;
import java.util.NoSuchElementException;
import java.util.UUID;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopFieldDocs;
import org.infinispan.AdvancedCache;
import org.infinispan.query.ResultIterator;
import org.infinispan.query.logging.Log;
import org.infinispan.util.logging.LogFactory;
/**
* DistributedIterator.
*
* Iterates on a distributed query.
*
* @author Israel Lacerra <israeldl@gmail.com>
* @author <a href="mailto:mluksa@redhat.com">Marko Luksa</a>
* @author Sanne Grinovero
* @since 5.1
*/
public class DistributedIterator<E> implements ResultIterator<E> {
private static final Log log = LogFactory.getLog(DistributedIterator.class, Log.class);
protected final AdvancedCache<?, ?> cache;
private int currentIndex = -1;
private final int fetchSize;
private final int resultSize;
private final int maxResults;
private final int firstResult;
private final ClusteredTopDocs[] partialResults;
private final int[] partialPositionNext;
private final TopDocs mergedResults;
public DistributedIterator(Sort sort, int fetchSize, int resultSize, int maxResults, int firstResult, HashMap<UUID, ClusteredTopDocs> topDocsResponses, AdvancedCache<?, ?> cache) {
this.fetchSize = fetchSize;
this.resultSize = resultSize;
this.maxResults = maxResults;
this.firstResult = firstResult;
this.cache = cache;
final int parallels = topDocsResponses.size();
this.partialResults = new ClusteredTopDocs[parallels];
TopDocs[] partialTopDocs = sort != null ? new TopFieldDocs[parallels] : new TopDocs[parallels];
this.partialPositionNext = new int[parallels];
int i=0;
for (Entry<UUID, ClusteredTopDocs> entry : topDocsResponses.entrySet()) {
partialResults[i] = entry.getValue();
partialTopDocs[i] = partialResults[i].getNodeTopDocs().topDocs;
i++;
}
try {
if (sort != null) {
mergedResults = TopDocs.merge(sort, firstResult, maxResults, (TopFieldDocs[]) partialTopDocs);
} else {
mergedResults = TopDocs.merge(firstResult, maxResults, partialTopDocs);
}
} catch (IOException e) {
throw log.unexpectedIOException(e);
}
}
@Override
public void close() {
// Nothing to do (extension point)
}
@Override
public E next() {
if (!hasNext())
throw new NoSuchElementException("Out of boundaries");
currentIndex++;
// fetch and return the value
ScoreDoc scoreDoc = mergedResults.scoreDocs[currentIndex];
int index = scoreDoc.shardIndex;
int specificPosition = partialPositionNext[index];
partialPositionNext[index]++;
return fetchValue(specificPosition, partialResults[index]);
}
protected E fetchValue(int scoreIndex, ClusteredTopDocs topDoc) {
NodeTopDocs eagerTopDocs = topDoc.getNodeTopDocs();
return (E) cache.get(eagerTopDocs.keys[scoreIndex]);
}
@Override
public final void remove() {
//TODO implement it?
throw new UnsupportedOperationException("This Iterator is read only");
}
@Override
public final boolean hasNext() {
int nextIndex = currentIndex + 1;
if (firstResult + nextIndex >= resultSize || nextIndex >= maxResults) {
return false;
}
return true;
}
}