CavalliumDBEngine/src/main/java/it/cavallium/dbengine/lucene/LLSlotDocCodec.java
2022-10-02 03:09:50 +02:00

210 lines
6.3 KiB
Java

package it.cavallium.dbengine.lucene;
import io.netty5.buffer.Buffer;
import it.cavallium.dbengine.database.DiscardingCloseable;
import it.cavallium.dbengine.database.SafeCloseable;
import it.cavallium.dbengine.database.disk.LLTempHugePqEnv;
import it.cavallium.dbengine.utils.SimpleResource;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.FieldComparator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LeafFieldComparator;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.rocksdb.AbstractComparator;
import org.rocksdb.ComparatorOptions;
public class LLSlotDocCodec extends SimpleResource
implements HugePqCodec<LLSlotDoc>, FieldValueHitQueue, DiscardingCloseable {
private final SortField[] fields;
protected final FieldComparator<?>[] comparators;
protected final int[] reverseMul;
private final ComparatorOptions comparatorOptions;
private final AbstractComparator comparator;
public LLSlotDocCodec(LLTempHugePqEnv env, int numHits, SortField[] fields) {
// When we get here, fields.length is guaranteed to be > 0, therefore no
// need to check it again.
// All these are required by this class's API - need to return arrays.
// Therefore even in the case of a single comparator, create an array
// anyway.
this.fields = fields;
int numComparators = fields.length;
comparators = new FieldComparator<?>[numComparators];
reverseMul = new int[numComparators];
for (int i = 0; i < numComparators; ++i) {
SortField field = fields[i];
reverseMul[i] = field.getReverse() ? -1 : 1;
comparators[i] = HugePqComparator.getComparator(env, field, numHits, i == 0);
}
comparatorOptions = new ComparatorOptions().setMaxReusedBufferSize(0);
comparator = new AbstractComparator(comparatorOptions) {
@Override
public String name() {
return "slot-doc-codec-comparator";
}
@Override
public int compare(ByteBuffer hitA, ByteBuffer hitB) {
assert hitA != hitB;
hitA.position(hitA.position() + Float.BYTES);
hitB.position(hitB.position() + Float.BYTES);
var docA = readDoc(hitA);
var docB = readDoc(hitB);
if (docA == docB) {
return 0;
}
hitA.position(hitA.position() + Integer.BYTES);
hitB.position(hitB.position() + Integer.BYTES);
var slotA = readSlot(hitA);
var slotB = readSlot(hitB);
assert slotA != slotB : "Slot " + slotA + " is equal to slot " + slotB;
int numComparators = comparators.length;
for (int i = 0; i < numComparators; ++i) {
final int c = reverseMul[i] * comparators[i].compare(slotA, slotB);
if (c != 0) {
// Short circuit
return -c;
}
}
// avoid random sort order that could lead to duplicates (bug #31241):
return Integer.compare(docB, docA);
}
};
}
@Override
public Buffer serialize(Function<Integer, Buffer> allocator, LLSlotDoc data) {
var buf = allocator.apply(Float.BYTES + Integer.BYTES + Integer.BYTES + Integer.BYTES);
buf.writerOffset(Float.BYTES + Integer.BYTES + Integer.BYTES + Integer.BYTES);
setScore(buf, data.score());
setDoc(buf, data.doc());
setShardIndex(buf, data.shardIndex());
setSlot(buf, data.slot());
return buf;
}
@Override
public LLSlotDoc deserialize(Buffer buf) {
return new LLSlotDoc(getDoc(buf), getScore(buf), getShardIndex(buf), getSlot(buf));
}
@Override
public AbstractComparator getComparator() {
return comparator;
}
private static float getScore(Buffer hit) {
return hit.getFloat(0);
}
private static int getDoc(Buffer hit) {
return hit.getInt(Float.BYTES);
}
private static int readDoc(ByteBuffer hit) {
return hit.getInt();
}
private static int getShardIndex(Buffer hit) {
return hit.getInt(Float.BYTES + Integer.BYTES);
}
private static int getSlot(Buffer hit) {
return hit.getInt(Float.BYTES + Integer.BYTES + Integer.BYTES);
}
private static int readSlot(ByteBuffer hit) {
return hit.getInt();
}
private static void setScore(Buffer hit, float score) {
hit.setFloat(0, score);
}
private static void setDoc(Buffer hit, int doc) {
hit.setInt(Float.BYTES, doc);
}
private static void setShardIndex(Buffer hit, int shardIndex) {
hit.setInt(Float.BYTES + Integer.BYTES, shardIndex);
}
private static void setSlot(Buffer hit, int slot) {
hit.setInt(Float.BYTES + Integer.BYTES + Integer.BYTES, slot);
}
@Override
public FieldComparator<?>[] getComparators() {
return comparators;
}
@Override
public int[] getReverseMul() {
return reverseMul;
}
@Override
public LeafFieldComparator[] getComparators(LeafReaderContext context) throws IOException {
LeafFieldComparator[] comparators = new LeafFieldComparator[this.comparators.length];
for (int i = 0; i < comparators.length; ++i) {
comparators[i] = this.comparators[i].getLeafComparator(context);
}
return comparators;
}
/**
* Given a queue Entry, creates a corresponding FieldDoc that contains the values used to sort the
* given document. These values are not the raw values out of the index, but the internal
* representation of them. This is so the given search hit can be collated by a MultiSearcher with
* other search hits.
*
* @param entry The Entry used to create a FieldDoc
* @return The newly created FieldDoc
* @see IndexSearcher#search(Query,int, Sort)
*/
@Override
public LLFieldDoc fillFields(final LLSlotDoc entry) {
final int n = comparators.length;
final List<Object> fields = new ArrayList<>(n);
for (FieldComparator<?> comparator : comparators) {
fields.add(comparator.value(entry.slot()));
}
// if (maxscore > 1.0f) doc.score /= maxscore; // normalize scores
return new LLFieldDoc(entry.doc(), entry.score(), entry.shardIndex(), fields);
}
/** Returns the SortFields being used by this hit queue. */
@Override
public SortField[] getFields() {
return fields;
}
@Override
protected void onClose() {
for (FieldComparator<?> comparator : this.comparators) {
if (comparator instanceof SafeCloseable closeable) {
closeable.close();
}
}
comparator.close();
comparatorOptions.close();
}
@Override
public LLSlotDoc clone(LLSlotDoc obj) {
return new LLSlotDoc(obj.doc(), obj.score(), obj.shardIndex(), obj.slot());
}
}