package org.apache.lucene.index;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import org.apache.lucene.codecs.DocValuesProducer;
import org.apache.lucene.codecs.FieldsProducer;
import org.apache.lucene.codecs.NormsProducer;
import org.apache.lucene.codecs.PointsReader;
import org.apache.lucene.codecs.StoredFieldsReader;
import org.apache.lucene.codecs.TermVectorsReader;
import org.apache.lucene.search.Sort;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.InfoStream;
import org.apache.lucene.util.packed.PackedInts;
import org.apache.lucene.util.packed.PackedLongValues;
import static org.apache.lucene.index.IndexWriter.isCongruentSort;
public class MergeState {
public final DocMap[] docMaps;
final DocMap[] leafDocMaps;
public final SegmentInfo segmentInfo;
public FieldInfos mergeFieldInfos;
public final StoredFieldsReader[] storedFieldsReaders;
public final TermVectorsReader[] termVectorsReaders;
public final NormsProducer[] normsProducers;
public final DocValuesProducer[] docValuesProducers;
public final FieldInfos[] fieldInfos;
public final Bits[] liveDocs;
public final FieldsProducer[] fieldsProducers;
public final PointsReader[] pointsReaders;
public final int[] maxDocs;
public final InfoStream infoStream;
public boolean needsIndexSort;
MergeState(List<CodecReader> originalReaders, SegmentInfo segmentInfo, InfoStream infoStream) throws IOException {
this.infoStream = infoStream;
final Sort indexSort = segmentInfo.getIndexSort();
int numReaders = originalReaders.size();
leafDocMaps = new DocMap[numReaders];
List<CodecReader> readers = maybeSortReaders(originalReaders, segmentInfo);
maxDocs = new int[numReaders];
fieldsProducers = new FieldsProducer[numReaders];
normsProducers = new NormsProducer[numReaders];
storedFieldsReaders = new StoredFieldsReader[numReaders];
termVectorsReaders = new TermVectorsReader[numReaders];
docValuesProducers = new DocValuesProducer[numReaders];
pointsReaders = new PointsReader[numReaders];
fieldInfos = new FieldInfos[numReaders];
liveDocs = new Bits[numReaders];
int numDocs = 0;
for(int i=0;i<numReaders;i++) {
final CodecReader reader = readers.get(i);
maxDocs[i] = reader.maxDoc();
liveDocs[i] = reader.getLiveDocs();
fieldInfos[i] = reader.getFieldInfos();
normsProducers[i] = reader.getNormsReader();
if (normsProducers[i] != null) {
normsProducers[i] = normsProducers[i].getMergeInstance();
}
docValuesProducers[i] = reader.getDocValuesReader();
if (docValuesProducers[i] != null) {
docValuesProducers[i] = docValuesProducers[i].getMergeInstance();
}
storedFieldsReaders[i] = reader.getFieldsReader();
if (storedFieldsReaders[i] != null) {
storedFieldsReaders[i] = storedFieldsReaders[i].getMergeInstance();
}
termVectorsReaders[i] = reader.getTermVectorsReader();
if (termVectorsReaders[i] != null) {
termVectorsReaders[i] = termVectorsReaders[i].getMergeInstance();
}
fieldsProducers[i] = reader.getPostingsReader().getMergeInstance();
pointsReaders[i] = reader.getPointsReader();
if (pointsReaders[i] != null) {
pointsReaders[i] = pointsReaders[i].getMergeInstance();
}
numDocs += reader.numDocs();
}
segmentInfo.setMaxDoc(numDocs);
this.segmentInfo = segmentInfo;
this.docMaps = buildDocMaps(readers, indexSort);
}
private DocMap[] buildDeletionDocMaps(List<CodecReader> readers) {
int totalDocs = 0;
int numReaders = readers.size();
DocMap[] docMaps = new DocMap[numReaders];
for (int i = 0; i < numReaders; i++) {
LeafReader reader = readers.get(i);
Bits liveDocs = reader.getLiveDocs();
final PackedLongValues delDocMap;
if (liveDocs != null) {
delDocMap = removeDeletes(reader.maxDoc(), liveDocs);
} else {
delDocMap = null;
}
final int docBase = totalDocs;
docMaps[i] = new DocMap() {
@Override
public int get(int docID) {
if (liveDocs == null) {
return docBase + docID;
} else if (liveDocs.get(docID)) {
return docBase + (int) delDocMap.get(docID);
} else {
return -1;
}
}
};
totalDocs += reader.numDocs();
}
return docMaps;
}
private DocMap[] buildDocMaps(List<CodecReader> readers, Sort indexSort) throws IOException {
if (indexSort == null) {
return buildDeletionDocMaps(readers);
} else {
long t0 = System.nanoTime();
DocMap[] result = MultiSorter.sort(indexSort, readers);
if (result == null) {
return buildDeletionDocMaps(readers);
} else {
needsIndexSort = true;
}
long t1 = System.nanoTime();
if (infoStream.isEnabled("SM")) {
infoStream.message("SM", String.format(Locale.ROOT, "%.2f msec to build merge sorted DocMaps", (t1-t0)/1000000.0));
}
return result;
}
}
private List<CodecReader> maybeSortReaders(List<CodecReader> originalReaders, SegmentInfo segmentInfo) throws IOException {
for(int i=0;i<originalReaders.size();i++) {
leafDocMaps[i] = new DocMap() {
@Override
public int get(int docID) {
return docID;
}
};
}
Sort indexSort = segmentInfo.getIndexSort();
if (indexSort == null) {
return originalReaders;
}
List<CodecReader> readers = new ArrayList<>(originalReaders.size());
for (CodecReader leaf : originalReaders) {
Sort segmentSort = leaf.getMetaData().getSort();
if (segmentSort == null || isCongruentSort(indexSort, segmentSort) == false) {
throw new IllegalArgumentException("index sort mismatch: merged segment has sort=" + indexSort +
" but to-be-merged segment has sort=" + (segmentSort == null ? "null" : segmentSort));
}
readers.add(leaf);
}
return readers;
}
public static abstract class DocMap {
public DocMap() {
}
public abstract int get(int docID);
}
static PackedLongValues removeDeletes(final int maxDoc, final Bits liveDocs) {
final PackedLongValues.Builder docMapBuilder = PackedLongValues.monotonicBuilder(PackedInts.COMPACT);
int del = 0;
for (int i = 0; i < maxDoc; ++i) {
docMapBuilder.add(i - del);
if (liveDocs.get(i) == false) {
++del;
}
}
return docMapBuilder.build();
}
}