Skip to content

Commit 46c0d27

Browse files
kaivalnpkaivalnp
committed
iter
1 parent 354e6e6 commit 46c0d27

9 files changed

Lines changed: 705 additions & 536 deletions

‎lucene/core/src/java/org/apache/lucene/codecs/dedup/DedupFlatVectorsReader.java‎

Lines changed: 202 additions & 176 deletions
Large diffs are not rendered by default.

‎lucene/core/src/java/org/apache/lucene/codecs/dedup/DedupFlatVectorsWriter.java‎

Lines changed: 371 additions & 250 deletions
Large diffs are not rendered by default.

‎lucene/core/src/java/org/apache/lucene/codecs/lucene95/HasIndexSlice.java‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,26 @@
2020

2121
/**
2222
* Implementors can return the IndexInput from which their values are read. For use by vector
23-
* quantizers.
23+
* quantizers and off-heap vector scorers.
2424
*/
2525
public interface HasIndexSlice {
2626

2727
/** Returns an IndexInput from which to read this instance's values, or null if not available. */
2828
IndexInput getSlice();
29+
30+
/**
31+
* Returns the byte offset within the backing {@link #getSlice() slice} for the vector at the
32+
* given ordinal. The default implementation assumes vectors are stored contiguously: {@code ord *
33+
* vectorByteSize}.
34+
*
35+
* <p>Formats that use indirection (e.g., de-duplicating formats with an ordinal mapping) override
36+
* this to return the correct offset.
37+
*
38+
* @param ord the vector ordinal
39+
* @param vectorByteSize the byte size of each vector
40+
* @return the byte offset for the vector
41+
*/
42+
default long ordToOffset(int ord, int vectorByteSize) {
43+
return (long) ord * vectorByteSize;
44+
}
2945
}

‎lucene/core/src/java/org/apache/lucene/index/KnnVectorValues.java‎

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -68,21 +68,6 @@ public int getVectorByteLength() {
6868
return dimension() * getEncoding().byteSize;
6969
}
7070

71-
/**
72-
* Returns the byte offset within the backing storage for the vector at the given ordinal. The
73-
* default implementation assumes vectors are stored contiguously: {@code ord *
74-
* getVectorByteLength()}.
75-
*
76-
* <p>Formats that use indirection (e.g., de-duplicating formats with an ordinal mapping) should
77-
* override this to return the correct offset.
78-
*
79-
* @param ord the vector ordinal
80-
* @return the byte offset for the vector
81-
*/
82-
public long ordToOffset(int ord) {
83-
return (long) ord * getVectorByteLength();
84-
}
85-
8671
/** The vector encoding of these values. */
8772
public abstract VectorEncoding getEncoding();
8873

‎lucene/core/src/java25/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentByteVectorScorer.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.io.IOException;
2020
import java.lang.foreign.MemorySegment;
2121
import java.util.Optional;
22+
import org.apache.lucene.codecs.lucene95.HasIndexSlice;
2223
import org.apache.lucene.index.ByteVectorValues;
2324
import org.apache.lucene.index.KnnVectorValues;
2425
import org.apache.lucene.index.VectorSimilarityFunction;
@@ -34,6 +35,7 @@ abstract sealed class Lucene99MemorySegmentByteVectorScorer
3435
final int vectorByteSize;
3536
final MemorySegmentAccessInput input;
3637
final KnnVectorValues values;
38+
final HasIndexSlice indexSlice;
3739
final byte[] query;
3840
byte[] scratch;
3941

@@ -63,13 +65,14 @@ public static Optional<Lucene99MemorySegmentByteVectorScorer> create(
6365
super(values);
6466
this.input = input;
6567
this.values = values;
68+
this.indexSlice = (HasIndexSlice) values;
6669
this.vectorByteSize = values.getVectorByteLength();
6770
this.query = queryVector;
6871
}
6972

7073
final MemorySegment getSegment(int ord) throws IOException {
7174
checkOrdinal(ord);
72-
long byteOffset = values.ordToOffset(ord);
75+
long byteOffset = indexSlice.ordToOffset(ord, vectorByteSize);
7376
MemorySegment seg = input.segmentSliceOrNull(byteOffset, vectorByteSize);
7477
if (seg == null) {
7578
if (scratch == null) {

‎lucene/core/src/java25/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentByteVectorScorerSupplier.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.io.IOException;
2020
import java.lang.foreign.MemorySegment;
2121
import java.util.Optional;
22+
import org.apache.lucene.codecs.lucene95.HasIndexSlice;
2223
import org.apache.lucene.index.ByteVectorValues;
2324
import org.apache.lucene.index.KnnVectorValues;
2425
import org.apache.lucene.index.VectorSimilarityFunction;
@@ -36,6 +37,7 @@ public abstract sealed class Lucene99MemorySegmentByteVectorScorerSupplier
3637
final int maxOrd;
3738
final MemorySegmentAccessInput input;
3839
final KnnVectorValues values; // to support ordToDoc/getAcceptOrds
40+
final HasIndexSlice indexSlice; // to support ordToOffset
3941
byte[] scratch1, scratch2;
4042

4143
/**
@@ -62,6 +64,7 @@ static Optional<RandomVectorScorerSupplier> create(
6264
MemorySegmentAccessInput input, KnnVectorValues values) {
6365
this.input = input;
6466
this.values = values;
67+
this.indexSlice = (HasIndexSlice) values;
6568
this.vectorByteSize = values.getVectorByteLength();
6669
this.maxOrd = values.size();
6770
}
@@ -79,7 +82,7 @@ final void checkOrdinal(int ord) {
7982
}
8083

8184
final MemorySegment getFirstSegment(int ord) throws IOException {
82-
long byteOffset = values.ordToOffset(ord);
85+
long byteOffset = indexSlice.ordToOffset(ord, vectorByteSize);
8386
MemorySegment seg = input.segmentSliceOrNull(byteOffset, vectorByteSize);
8487
if (seg == null) {
8588
if (scratch1 == null) {
@@ -92,7 +95,7 @@ final MemorySegment getFirstSegment(int ord) throws IOException {
9295
}
9396

9497
final MemorySegment getSecondSegment(int ord) throws IOException {
95-
long byteOffset = values.ordToOffset(ord);
98+
long byteOffset = indexSlice.ordToOffset(ord, vectorByteSize);
9699
MemorySegment seg = input.segmentSliceOrNull(byteOffset, vectorByteSize);
97100
if (seg == null) {
98101
if (scratch2 == null) {

‎lucene/core/src/java25/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorer.java‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.io.IOException;
2020
import java.lang.foreign.MemorySegment;
2121
import java.util.Optional;
22+
import org.apache.lucene.codecs.lucene95.HasIndexSlice;
2223
import org.apache.lucene.index.FloatVectorValues;
2324
import org.apache.lucene.index.VectorSimilarityFunction;
2425
import org.apache.lucene.store.FilterIndexInput;
@@ -31,6 +32,7 @@ abstract sealed class Lucene99MemorySegmentFloatVectorScorer
3132
extends RandomVectorScorer.AbstractRandomVectorScorer {
3233

3334
final FloatVectorValues values;
35+
final HasIndexSlice indexSlice;
3436
final int vectorByteSize;
3537
final MemorySegment seg;
3638
final float[] query;
@@ -62,6 +64,7 @@ public static Optional<Lucene99MemorySegmentFloatVectorScorer> create(
6264
MemorySegment seg, FloatVectorValues values, float[] query) {
6365
super(values);
6466
this.values = values;
67+
this.indexSlice = (HasIndexSlice) values;
6568
this.seg = seg;
6669
this.vectorByteSize = values.getVectorByteLength();
6770
this.query = query;
@@ -85,10 +88,10 @@ public float bulkScore(int[] nodes, float[] scores, int numNodes) throws IOExcep
8588
final int limit = numNodes & ~3;
8689
float maxScore = Float.NEGATIVE_INFINITY;
8790
for (; i < limit; i += 4) {
88-
long offset1 = values.ordToOffset(nodes[i]);
89-
long offset2 = values.ordToOffset(nodes[i + 1]);
90-
long offset3 = values.ordToOffset(nodes[i + 2]);
91-
long offset4 = values.ordToOffset(nodes[i + 3]);
91+
long offset1 = indexSlice.ordToOffset(nodes[i], vectorByteSize);
92+
long offset2 = indexSlice.ordToOffset(nodes[i + 1], vectorByteSize);
93+
long offset3 = indexSlice.ordToOffset(nodes[i + 2], vectorByteSize);
94+
long offset4 = indexSlice.ordToOffset(nodes[i + 3], vectorByteSize);
9295
vectorOp(seg, scratchScores, offset1, offset2, offset3, offset4, query.length);
9396
scores[i + 0] = normalizeRawScore(scratchScores[0]);
9497
maxScore = Math.max(maxScore, scores[i + 0]);
@@ -102,9 +105,9 @@ public float bulkScore(int[] nodes, float[] scores, int numNodes) throws IOExcep
102105
// Handle remaining 1–3 nodes in bulk (if any)
103106
int remaining = numNodes - i;
104107
if (remaining > 0) {
105-
long addr1 = values.ordToOffset(nodes[i]);
106-
long addr2 = (remaining > 1) ? values.ordToOffset(nodes[i + 1]) : addr1;
107-
long addr3 = (remaining > 2) ? values.ordToOffset(nodes[i + 2]) : addr1;
108+
long addr1 = indexSlice.ordToOffset(nodes[i], vectorByteSize);
109+
long addr2 = (remaining > 1) ? indexSlice.ordToOffset(nodes[i + 1], vectorByteSize) : addr1;
110+
long addr3 = (remaining > 2) ? indexSlice.ordToOffset(nodes[i + 2], vectorByteSize) : addr1;
108111
vectorOp(seg, scratchScores, addr1, addr2, addr3, addr3, query.length);
109112
scores[i] = normalizeRawScore(scratchScores[0]);
110113
maxScore = Math.max(maxScore, scores[i]);

‎lucene/core/src/java25/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorerSupplier.java‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.io.IOException;
2020
import java.lang.foreign.MemorySegment;
2121
import java.util.Optional;
22+
import org.apache.lucene.codecs.lucene95.HasIndexSlice;
2223
import org.apache.lucene.index.FloatVectorValues;
2324
import org.apache.lucene.index.VectorSimilarityFunction;
2425
import org.apache.lucene.store.FilterIndexInput;
@@ -36,6 +37,7 @@ public abstract sealed class Lucene99MemorySegmentFloatVectorScorerSupplier
3637
final int dims;
3738
final MemorySegment seg;
3839
final FloatVectorValues values; // to support ordToDoc/getAcceptOrds
40+
final HasIndexSlice indexSlice; // to support ordToOffset
3941

4042
/**
4143
* Return an optional whose value, if present, is the scorer supplier. Otherwise, an empty
@@ -62,6 +64,7 @@ static Optional<RandomVectorScorerSupplier> create(
6264
Lucene99MemorySegmentFloatVectorScorerSupplier(MemorySegment seg, FloatVectorValues values) {
6365
this.seg = seg;
6466
this.values = values;
67+
this.indexSlice = (HasIndexSlice) values;
6568
this.vectorByteSize = values.getVectorByteLength();
6669
this.maxOrd = values.size();
6770
this.dims = values.dimension();
@@ -285,23 +288,23 @@ abstract void vectorOp(
285288
@Override
286289
public float score(int node) {
287290
checkOrdinal(node);
288-
long queryAddr = values.ordToOffset(queryOrd);
289-
long addr = values.ordToOffset(node);
291+
long queryAddr = indexSlice.ordToOffset(queryOrd, vectorByteSize);
292+
long addr = indexSlice.ordToOffset(node, vectorByteSize);
290293
var raw = vectorOp(seg, queryAddr, addr, dims);
291294
return normalizeRawScore(raw);
292295
}
293296

294297
@Override
295298
public float bulkScore(int[] nodes, float[] scores, int numNodes) {
296299
int i = 0;
297-
long queryAddr = values.ordToOffset(queryOrd);
300+
long queryAddr = indexSlice.ordToOffset(queryOrd, vectorByteSize);
298301
float maxScore = Float.NEGATIVE_INFINITY;
299302
final int limit = numNodes & ~3;
300303
for (; i < limit; i += 4) {
301-
long offset1 = values.ordToOffset(nodes[i]);
302-
long offset2 = values.ordToOffset(nodes[i + 1]);
303-
long offset3 = values.ordToOffset(nodes[i + 2]);
304-
long offset4 = values.ordToOffset(nodes[i + 3]);
304+
long offset1 = indexSlice.ordToOffset(nodes[i], vectorByteSize);
305+
long offset2 = indexSlice.ordToOffset(nodes[i + 1], vectorByteSize);
306+
long offset3 = indexSlice.ordToOffset(nodes[i + 2], vectorByteSize);
307+
long offset4 = indexSlice.ordToOffset(nodes[i + 3], vectorByteSize);
305308
vectorOp(seg, scratchScores, queryAddr, offset1, offset2, offset3, offset4, dims);
306309
scores[i + 0] = normalizeRawScore(scratchScores[0]);
307310
maxScore = Math.max(maxScore, scores[i + 0]);
@@ -315,9 +318,9 @@ public float bulkScore(int[] nodes, float[] scores, int numNodes) {
315318
// Handle remaining 1–3 nodes in bulk (if any)
316319
int remaining = numNodes - i;
317320
if (remaining > 0) {
318-
long addr1 = values.ordToOffset(nodes[i]);
319-
long addr2 = (remaining > 1) ? values.ordToOffset(nodes[i + 1]) : addr1;
320-
long addr3 = (remaining > 2) ? values.ordToOffset(nodes[i + 2]) : addr1;
321+
long addr1 = indexSlice.ordToOffset(nodes[i], vectorByteSize);
322+
long addr2 = (remaining > 1) ? indexSlice.ordToOffset(nodes[i + 1], vectorByteSize) : addr1;
323+
long addr3 = (remaining > 2) ? indexSlice.ordToOffset(nodes[i + 2], vectorByteSize) : addr1;
321324
vectorOp(seg, scratchScores, queryAddr, addr1, addr2, addr3, addr1, dims);
322325
scores[i] = normalizeRawScore(scratchScores[0]);
323326
maxScore = Math.max(maxScore, scores[i]);

0 commit comments

Comments
 (0)