package com.browseengine.bobo.facets.data;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntList;
import java.io.IOException;
import java.io.Serializable;
import org.apache.lucene.index.AtomicReader;
import org.apache.lucene.index.DocsEnum;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.util.BytesRef;
import com.browseengine.bobo.api.BoboSegmentReader;
import com.browseengine.bobo.facets.FacetHandler;
import com.browseengine.bobo.facets.FacetHandler.TermCountSize;
import com.browseengine.bobo.sort.DocComparator;
import com.browseengine.bobo.sort.DocComparatorSource;
import com.browseengine.bobo.util.BigByteArray;
import com.browseengine.bobo.util.BigIntArray;
import com.browseengine.bobo.util.BigSegmentedArray;
import com.browseengine.bobo.util.BigShortArray;
public class FacetDataCache<T> implements Serializable {
/**
*
*/
private static final long serialVersionUID = 1L;
public BigSegmentedArray orderArray;
public TermValueList<T> valArray;
public int[] freqs;
public int[] minIDs;
public int[] maxIDs;
public FacetDataCache(BigSegmentedArray orderArray, TermValueList<T> valArray, int[] freqs,
int[] minIDs, int[] maxIDs, TermCountSize termCountSize) {
this.orderArray = orderArray;
this.valArray = valArray;
this.freqs = freqs;
this.minIDs = minIDs;
this.maxIDs = maxIDs;
}
public FacetDataCache() {
this.orderArray = null;
this.valArray = null;
this.maxIDs = null;
this.minIDs = null;
this.freqs = null;
}
public int getNumItems(int docid) {
int valIdx = orderArray.get(docid);
return valIdx <= 0 ? 0 : 1;
}
private final static BigSegmentedArray newInstance(int termCount, int maxDoc) {
// we use < instead of <= to take into consideration "missing" value (zero element in the
// dictionary)
if (termCount < Byte.MAX_VALUE) {
return new BigByteArray(maxDoc);
} else if (termCount < Short.MAX_VALUE) {
return new BigShortArray(maxDoc);
} else return new BigIntArray(maxDoc);
}
protected int getDictValueCount(AtomicReader reader, String field) throws IOException {
int ret = 0;
Terms terms = reader.terms(field);
if (terms == null) {
return ret;
}
return (int) terms.size();
}
protected int getNegativeValueCount(AtomicReader reader, String field) throws IOException {
int ret = 0;
Terms terms = reader.terms(field);
if (terms == null) {
return ret;
}
TermsEnum termsEnum = terms.iterator(null);
BytesRef text;
while ((text = termsEnum.next()) != null) {
if (!text.utf8ToString().startsWith("-")) {
break;
}
ret++;
}
return ret;
}
public void load(String fieldName, AtomicReader reader, TermListFactory<T> listFactory)
throws IOException {
String field = fieldName.intern();
int maxDoc = reader.maxDoc();
int dictValueCount = getDictValueCount(reader, fieldName);
BigSegmentedArray order = newInstance(dictValueCount, maxDoc);
this.orderArray = order;
IntArrayList minIDList = new IntArrayList();
IntArrayList maxIDList = new IntArrayList();
IntArrayList freqList = new IntArrayList();
int length = maxDoc + 1;
@SuppressWarnings("unchecked")
TermValueList<T> list = listFactory == null ? (TermValueList<T>) new TermStringList()
: listFactory.createTermList();
int negativeValueCount = getNegativeValueCount(reader, field);
int t = 1; // valid term id starts from 1
list.add(null);
minIDList.add(-1);
maxIDList.add(-1);
freqList.add(0);
int totalFreq = 0;
Terms terms = reader.terms(field);
if (terms != null) {
TermsEnum termsEnum = terms.iterator(null);
BytesRef text;
while ((text = termsEnum.next()) != null) {
// store term text
// we expect that there is at most one term per document
if (t >= length) throw new RuntimeException("there are more terms than "
+ "documents in field \"" + field + "\", but it's impossible to sort on "
+ "tokenized fields");
String strText = text.utf8ToString();
list.add(strText);
Term term = new Term(field, strText);
DocsEnum docsEnum = reader.termDocsEnum(term);
int minID = -1;
int maxID = -1;
int docID;
int df = 0;
int valId = (t - 1 < negativeValueCount) ? (negativeValueCount - t + 1) : t;
while ((docID = docsEnum.nextDoc()) != DocsEnum.NO_MORE_DOCS) {
df++;
order.add(docID, valId);
minID = docID;
while (docsEnum.nextDoc() != DocsEnum.NO_MORE_DOCS) {
docID = docsEnum.docID();
df++;
order.add(docID, valId);
}
maxID = docID;
}
freqList.add(df);
totalFreq += df;
minIDList.add(minID);
maxIDList.add(maxID);
t++;
}
}
list.seal();
// Process minIDList and maxIDList for negative number
for (int i = 1; i < negativeValueCount/2 + 1; ++i) {
int top = i;
int tail = negativeValueCount - i + 1;
int topValue = minIDList.getInt(top);
int tailValue = minIDList.getInt(tail);
minIDList.set(top, tailValue);
minIDList.set(tail, topValue);
topValue = maxIDList.getInt(top);
tailValue = maxIDList.getInt(tail);
maxIDList.set(top, tailValue);
maxIDList.set(tail, topValue);
}
this.valArray = list;
this.freqs = freqList.toIntArray();
this.minIDs = minIDList.toIntArray();
this.maxIDs = maxIDList.toIntArray();
int doc = 0;
while (doc < maxDoc && order.get(doc) != 0) {
++doc;
}
if (doc < maxDoc) {
this.minIDs[0] = doc;
// Try to get the max
doc = maxDoc - 1;
while (doc >= 0 && order.get(doc) != 0) {
--doc;
}
this.maxIDs[0] = doc;
}
this.freqs[0] = reader.numDocs() - totalFreq;
}
private static int[] convertString(FacetDataCache<?> dataCache, String[] vals) {
IntList list = new IntArrayList(vals.length);
for (int i = 0; i < vals.length; ++i) {
int index = dataCache.valArray.indexOf(vals[i]);
if (index >= 0) {
list.add(index);
}
}
return list.toIntArray();
}
/**
* Same as convert(FacetDataCache dataCache,String[] vals) except that the
* values are supplied in raw form so that we can take advantage of the type
* information to find index faster.
*
* @param <T>
* @param dataCache
* @param vals
* @return the array of order indices of the values.
*/
public static <T> int[] convert(FacetDataCache<T> dataCache, T[] vals) {
if (vals != null && (vals instanceof String[])) return convertString(dataCache, (String[]) vals);
IntList list = new IntArrayList(vals.length);
for (int i = 0; i < vals.length; ++i) {
int index = dataCache.valArray.indexOfWithType(vals[i]);
if (index >= 0) {
list.add(index);
}
}
return list.toIntArray();
}
public static class FacetDocComparatorSource extends DocComparatorSource {
private final FacetHandler<FacetDataCache<?>> _facetHandler;
public FacetDocComparatorSource(FacetHandler<FacetDataCache<?>> facetHandler) {
_facetHandler = facetHandler;
}
@Override
public DocComparator getComparator(AtomicReader reader, int docbase) throws IOException {
if (!(reader instanceof BoboSegmentReader)) throw new IllegalStateException(
"reader not instance of " + BoboSegmentReader.class);
BoboSegmentReader boboReader = (BoboSegmentReader) reader;
final FacetDataCache<?> dataCache = _facetHandler.getFacetData(boboReader);
final BigSegmentedArray orderArray = dataCache.orderArray;
return new DocComparator() {
@Override
public Comparable<?> value(ScoreDoc doc) {
int index = orderArray.get(doc.doc);
return dataCache.valArray.getComparableValue(index);
}
@Override
public int compare(ScoreDoc doc1, ScoreDoc doc2) {
return orderArray.get(doc1.doc) - orderArray.get(doc2.doc);
}
};
}
}
}