package org.apache.lucene.facet.sortedset;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import org.apache.lucene.facet.FacetsConfig;
import org.apache.lucene.facet.sortedset.SortedSetDocValuesReaderState.OrdRange;
import org.apache.lucene.index.DocValues;
import org.apache.lucene.index.DocValuesType;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.MultiDocValues.MultiSortedSetDocValues;
import org.apache.lucene.index.MultiDocValues;
import org.apache.lucene.index.OrdinalMap;
import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.util.Accountable;
import org.apache.lucene.util.Accountables;
import org.apache.lucene.util.BytesRef;
public class DefaultSortedSetDocValuesReaderState extends SortedSetDocValuesReaderState {
private final String field;
private final int valueCount;
public final IndexReader reader;
private final Map<String,OrdinalMap> cachedOrdMaps = new HashMap<>();
private final Map<String,OrdRange> prefixToOrdRange = new HashMap<>();
public DefaultSortedSetDocValuesReaderState(IndexReader reader) throws IOException {
this(reader, FacetsConfig.DEFAULT_INDEX_FIELD_NAME);
}
public DefaultSortedSetDocValuesReaderState(IndexReader reader, String field) throws IOException {
this.field = field;
this.reader = reader;
SortedSetDocValues dv = getDocValues();
if (dv == null) {
throw new IllegalArgumentException("field \"" + field + "\" was not indexed with SortedSetDocValues");
}
if (dv.getValueCount() > Integer.MAX_VALUE) {
throw new IllegalArgumentException("can only handle valueCount < Integer.MAX_VALUE; got " + dv.getValueCount());
}
valueCount = (int) dv.getValueCount();
String lastDim = null;
int startOrd = -1;
for(int ord=0;ord<valueCount;ord++) {
final BytesRef term = dv.lookupOrd(ord);
String[] components = FacetsConfig.stringToPath(term.utf8ToString());
if (components.length != 2) {
throw new IllegalArgumentException("this class can only handle 2 level hierarchy (dim/value); got: " + Arrays.toString(components) + " " + term.utf8ToString());
}
if (!components[0].equals(lastDim)) {
if (lastDim != null) {
prefixToOrdRange.put(lastDim, new OrdRange(startOrd, ord-1));
}
startOrd = ord;
lastDim = components[0];
}
}
if (lastDim != null) {
prefixToOrdRange.put(lastDim, new OrdRange(startOrd, valueCount-1));
}
}
@Override
public long ramBytesUsed() {
synchronized (cachedOrdMaps) {
long bytes = 0;
for (OrdinalMap map : cachedOrdMaps.values()) {
bytes += map.ramBytesUsed();
}
return bytes;
}
}
@Override
public Collection<Accountable> getChildResources() {
synchronized (cachedOrdMaps) {
return Accountables.namedAccountables("DefaultSortedSetDocValuesReaderState", cachedOrdMaps);
}
}
@Override
public String toString() {
return "DefaultSortedSetDocValuesReaderState(field=" + field + " reader=" + reader + ")";
}
@Override
public SortedSetDocValues getDocValues() throws IOException {
OrdinalMap map = null;
synchronized (cachedOrdMaps) {
map = cachedOrdMaps.get(field);
if (map == null) {
SortedSetDocValues dv = MultiDocValues.getSortedSetValues(reader, field);
if (dv instanceof MultiDocValues.MultiSortedSetDocValues) {
map = ((MultiDocValues.MultiSortedSetDocValues)dv).mapping;
IndexReader.CacheHelper cacheHelper = reader.getReaderCacheHelper();
if (cacheHelper != null && map.owner == cacheHelper.getKey()) {
cachedOrdMaps.put(field, map);
}
}
return dv;
}
}
assert map != null;
int size = reader.leaves().size();
final SortedSetDocValues[] values = new SortedSetDocValues[size];
final int[] starts = new int[size+1];
long cost = 0;
for (int i = 0; i < size; i++) {
LeafReaderContext context = reader.leaves().get(i);
final LeafReader reader = context.reader();
final FieldInfo fieldInfo = reader.getFieldInfos().fieldInfo(field);
if (fieldInfo != null && fieldInfo.getDocValuesType() != DocValuesType.SORTED_SET) {
return null;
}
SortedSetDocValues v = reader.getSortedSetDocValues(field);
if (v == null) {
v = DocValues.emptySortedSet();
}
values[i] = v;
starts[i] = context.docBase;
cost += v.cost();
}
starts[size] = reader.maxDoc();
return new MultiSortedSetDocValues(values, starts, map, cost);
}
@Override
public Map<String,OrdRange> getPrefixToOrdRange() {
return prefixToOrdRange;
}
@Override
public OrdRange getOrdRange(String dim) {
return prefixToOrdRange.get(dim);
}
@Override
public String getField() {
return field;
}
@Override
public IndexReader getReader() {
return reader;
}
@Override
public int getSize() {
return valueCount;
}
}