parentVal.setLongValue(random().nextInt(50));
documents.add(parentDoc);
writer.addDocuments(documents);
}
writer.forceMerge(1);
final DirectoryReader indexReader = writer.getReader();
writer.close();
final AtomicReader reader = getOnlySegmentReader(indexReader);
final Filter parentsFilter = new FixedBitSetCachingWrapperFilter(new QueryWrapperFilter(new TermQuery(new Term("parent", "true"))));
final FixedBitSet parentBits = (FixedBitSet) parentsFilter.getDocIdSet(reader.getContext(), null);
final NumericDocValues parentValues = reader.getNumericDocValues("parent_val");
final Sorter.DocComparator parentComparator = new Sorter.DocComparator() {
@Override
public int compare(int docID1, int docID2) {
assertTrue(parentBits.get(docID1));
assertTrue(parentBits.get(docID2));
final long v1 = parentValues.get(docID1);
final long v2 = parentValues.get(docID2);
return v1 < v2 ? -1 : v1 == v2 ? 0 : 1;
}
};
final NumericDocValues childValues = reader.getNumericDocValues("child_val");
final Sorter.DocComparator childComparator = new Sorter.DocComparator() {
@Override
public int compare(int docID1, int docID2) {
assertFalse(parentBits.get(docID1));
assertFalse(parentBits.get(docID2));
final long v1 = childValues.get(docID1);
final long v2 = childValues.get(docID2);
return v1 < v2 ? -1 : v1 == v2 ? 0 : 1;
}
};
final Sorter sorter = new BlockJoinSorter(parentsFilter) {
@Override
public String getID() {
return "Dummy";
}
@Override
protected DocComparator getParentComparator(AtomicReader r) {
assertEquals(reader, r);
return parentComparator;
}
@Override
protected DocComparator getChildComparator(AtomicReader r) {
assertEquals(reader, r);
return childComparator;
}
};
final Sorter.DocMap docMap = sorter.sort(reader);
assertEquals(reader.maxDoc(), docMap.size());
int[] children = new int[1];
int numChildren = 0;
int previousParent = -1;
for (int i = 0; i < docMap.size(); ++i) {
final int oldID = docMap.newToOld(i);
if (parentBits.get(oldID)) {
// check that we have the right children
for (int j = 0; j < numChildren; ++j) {
assertEquals(oldID, parentBits.nextSetBit(children[j]));
}
// check that children are sorted
for (int j = 1; j < numChildren; ++j) {
final int doc1 = children[j-1];
final int doc2 = children[j];
if (childValues.get(doc1) == childValues.get(doc2)) {
assertTrue(doc1 < doc2); // sort is stable
} else {
assertTrue(childValues.get(doc1) < childValues.get(doc2));
}
}
// check that parents are sorted
if (previousParent != -1) {
if (parentValues.get(previousParent) == parentValues.get(oldID)) {
assertTrue(previousParent < oldID);
} else {
assertTrue(parentValues.get(previousParent) < parentValues.get(oldID));
}
}
// reset
previousParent = oldID;
numChildren = 0;
} else {
children = ArrayUtil.grow(children, numChildren+1);
children[numChildren++] = oldID;
}
}
indexReader.close();
writer.w.getDirectory().close();
}