package org.apache.lucene.index;
import java.io.IOException;
import org.apache.lucene.codecs.DocValuesConsumer;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.SortField;
import org.apache.lucene.util.Counter;
import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.packed.PackedInts;
import org.apache.lucene.util.packed.PackedLongValues;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
class NumericDocValuesWriter extends DocValuesWriter {
private PackedLongValues.Builder pending;
private PackedLongValues finalValues;
private final Counter iwBytesUsed;
private long bytesUsed;
private DocsWithFieldSet docsWithField;
private final FieldInfo fieldInfo;
private int lastDocID = -1;
public NumericDocValuesWriter(FieldInfo fieldInfo, Counter iwBytesUsed) {
pending = PackedLongValues.deltaPackedBuilder(PackedInts.COMPACT);
docsWithField = new DocsWithFieldSet();
bytesUsed = pending.ramBytesUsed() + docsWithField.ramBytesUsed();
this.fieldInfo = fieldInfo;
this.iwBytesUsed = iwBytesUsed;
iwBytesUsed.addAndGet(bytesUsed);
}
public void addValue(int docID, long value) {
if (docID <= lastDocID) {
throw new IllegalArgumentException("DocValuesField \"" + fieldInfo.name + "\" appears more than once in this document (only one value is allowed per field)");
}
pending.add(value);
docsWithField.add(docID);
updateBytesUsed();
lastDocID = docID;
}
private void updateBytesUsed() {
final long newBytesUsed = pending.ramBytesUsed() + docsWithField.ramBytesUsed();
iwBytesUsed.addAndGet(newBytesUsed - bytesUsed);
bytesUsed = newBytesUsed;
}
@Override
public void finish(int maxDoc) {
}
@Override
Sorter.DocComparator getDocComparator(int maxDoc, SortField sortField) throws IOException {
assert finalValues == null;
finalValues = pending.build();
final BufferedNumericDocValues docValues =
new BufferedNumericDocValues(finalValues, docsWithField.iterator());
return Sorter.getDocComparator(maxDoc, sortField, () -> null, () -> docValues);
}
@Override
DocIdSetIterator getDocIdSet() {
return docsWithField.iterator();
}
static SortingLeafReader.CachedNumericDVs sortDocValues(int maxDoc, Sorter.DocMap sortMap, NumericDocValues oldDocValues) throws IOException {
FixedBitSet docsWithField = new FixedBitSet(maxDoc);
long[] values = new long[maxDoc];
while (true) {
int docID = oldDocValues.nextDoc();
if (docID == NO_MORE_DOCS) {
break;
}
int newDocID = sortMap.oldToNew(docID);
docsWithField.set(newDocID);
values[newDocID] = oldDocValues.longValue();
}
return new SortingLeafReader.CachedNumericDVs(values, docsWithField);
}
@Override
public void flush(SegmentWriteState state, Sorter.DocMap sortMap, DocValuesConsumer dvConsumer) throws IOException {
final PackedLongValues values;
if (finalValues == null) {
values = pending.build();
} else {
values = finalValues;
}
final SortingLeafReader.CachedNumericDVs sorted;
if (sortMap != null) {
NumericDocValues oldValues = new BufferedNumericDocValues(values, docsWithField.iterator());
sorted = sortDocValues(state.segmentInfo.maxDoc(), sortMap, oldValues);
} else {
sorted = null;
}
dvConsumer.addNumericField(fieldInfo,
new EmptyDocValuesProducer() {
@Override
public NumericDocValues getNumeric(FieldInfo fieldInfo) {
if (fieldInfo != NumericDocValuesWriter.this.fieldInfo) {
throw new IllegalArgumentException("wrong fieldInfo");
}
if (sorted == null) {
return new BufferedNumericDocValues(values, docsWithField.iterator());
} else {
return new SortingLeafReader.SortingNumericDocValues(sorted);
}
}
});
}
private static class BufferedNumericDocValues extends NumericDocValues {
final PackedLongValues.Iterator iter;
final DocIdSetIterator docsWithField;
private long value;
BufferedNumericDocValues(PackedLongValues values, DocIdSetIterator docsWithFields) {
this.iter = values.iterator();
this.docsWithField = docsWithFields;
}
@Override
public int docID() {
return docsWithField.docID();
}
@Override
public int nextDoc() throws IOException {
int docID = docsWithField.nextDoc();
if (docID != NO_MORE_DOCS) {
value = iter.next();
}
return docID;
}
@Override
public int advance(int target) {
throw new UnsupportedOperationException();
}
@Override
public boolean advanceExact(int target) throws IOException {
throw new UnsupportedOperationException();
}
@Override
public long cost() {
return docsWithField.cost();
}
@Override
public long longValue() {
return value;
}
}
}