8
8
9
9
package org .elasticsearch .benchmark .vector ;
10
10
11
+ import org .apache .lucene .codecs .lucene99 .OffHeapQuantizedByteVectorValues ;
11
12
import org .apache .lucene .index .VectorSimilarityFunction ;
12
13
import org .apache .lucene .store .Directory ;
13
14
import org .apache .lucene .store .IOContext ;
14
15
import org .apache .lucene .store .IndexInput ;
15
16
import org .apache .lucene .store .IndexOutput ;
16
17
import org .apache .lucene .store .MMapDirectory ;
18
+ import org .apache .lucene .util .hnsw .RandomVectorScorer ;
19
+ import org .apache .lucene .util .quantization .RandomAccessQuantizedByteVectorValues ;
20
+ import org .apache .lucene .util .quantization .ScalarQuantizedRandomVectorScorer ;
17
21
import org .apache .lucene .util .quantization .ScalarQuantizedVectorSimilarity ;
18
22
import org .elasticsearch .common .logging .LogConfigurator ;
19
23
import org .elasticsearch .core .IOUtils ;
20
- import org .elasticsearch .vec .VectorScorer ;
21
24
import org .elasticsearch .vec .VectorScorerFactory ;
22
25
import org .openjdk .jmh .annotations .Benchmark ;
23
26
import org .openjdk .jmh .annotations .BenchmarkMode ;
@@ -71,10 +74,10 @@ public class VectorScorerBenchmark {
71
74
float vec2Offset ;
72
75
float scoreCorrectionConstant ;
73
76
74
- ScalarQuantizedVectorSimilarity luceneDotScorer ;
75
- ScalarQuantizedVectorSimilarity luceneSqrScorer ;
76
- VectorScorer nativeDotScorer ;
77
- VectorScorer nativeSqrScorer ;
77
+ RandomVectorScorer luceneDotScorer ;
78
+ RandomVectorScorer luceneSqrScorer ;
79
+ RandomVectorScorer nativeDotScorer ;
80
+ RandomVectorScorer nativeSqrScorer ;
78
81
79
82
@ Setup
80
83
public void setup () throws IOException {
@@ -107,14 +110,22 @@ public void setup() throws IOException {
107
110
out .writeInt (Float .floatToIntBits (vec2Offset ));
108
111
}
109
112
in = dir .openInput ("vector.data" , IOContext .DEFAULT );
113
+ var values = vectorValues (dims , 2 , in );
110
114
111
- luceneDotScorer = ScalarQuantizedVectorSimilarity .fromVectorSimilarity (
112
- VectorSimilarityFunction .DOT_PRODUCT ,
113
- scoreCorrectionConstant
115
+ luceneDotScorer = new ScalarQuantizedRandomVectorScorer (
116
+ ScalarQuantizedVectorSimilarity .fromVectorSimilarity (VectorSimilarityFunction .DOT_PRODUCT , scoreCorrectionConstant ),
117
+ values .copy (),
118
+ vec1 ,
119
+ vec1Offset
114
120
);
115
- luceneSqrScorer = ScalarQuantizedVectorSimilarity .fromVectorSimilarity (VectorSimilarityFunction .EUCLIDEAN , scoreCorrectionConstant );
116
- nativeDotScorer = factory .getInt7ScalarQuantizedVectorScorer (dims , size , scoreCorrectionConstant , DOT_PRODUCT , in ).get ();
117
- nativeSqrScorer = factory .getInt7ScalarQuantizedVectorScorer (dims , size , scoreCorrectionConstant , EUCLIDEAN , in ).get ();
121
+ luceneSqrScorer = new ScalarQuantizedRandomVectorScorer (
122
+ ScalarQuantizedVectorSimilarity .fromVectorSimilarity (VectorSimilarityFunction .EUCLIDEAN , scoreCorrectionConstant ),
123
+ values .copy (),
124
+ vec1 ,
125
+ vec1Offset
126
+ );
127
+ nativeDotScorer = factory .getInt7ScalarQuantizedVectorScorer (DOT_PRODUCT , in , values , scoreCorrectionConstant ).get ().scorer (0 );
128
+ nativeSqrScorer = factory .getInt7ScalarQuantizedVectorScorer (EUCLIDEAN , in , values , scoreCorrectionConstant ).get ().scorer (0 );
118
129
119
130
// sanity
120
131
var f1 = dotProductLucene ();
@@ -144,13 +155,13 @@ public void teardown() throws IOException {
144
155
}
145
156
146
157
@ Benchmark
147
- public float dotProductLucene () {
148
- return luceneDotScorer .score (vec1 , vec1Offset , vec2 , vec2Offset );
158
+ public float dotProductLucene () throws IOException {
159
+ return luceneDotScorer .score (1 );
149
160
}
150
161
151
162
@ Benchmark
152
163
public float dotProductNative () throws IOException {
153
- return nativeDotScorer .score (0 , 1 );
164
+ return nativeDotScorer .score (1 );
154
165
}
155
166
156
167
@ Benchmark
@@ -166,13 +177,13 @@ public float dotProductScalar() {
166
177
// -- square distance
167
178
168
179
@ Benchmark
169
- public float squareDistanceLucene () {
170
- return luceneSqrScorer .score (vec1 , vec1Offset , vec2 , vec2Offset );
180
+ public float squareDistanceLucene () throws IOException {
181
+ return luceneSqrScorer .score (1 );
171
182
}
172
183
173
184
@ Benchmark
174
185
public float squareDistanceNative () throws IOException {
175
- return nativeSqrScorer .score (0 , 1 );
186
+ return nativeSqrScorer .score (1 );
176
187
}
177
188
178
189
@ Benchmark
@@ -186,6 +197,10 @@ public float squareDistanceScalar() {
186
197
return 1 / (1f + adjustedDistance );
187
198
}
188
199
200
+ RandomAccessQuantizedByteVectorValues vectorValues (int dims , int size , IndexInput in ) throws IOException {
201
+ return new OffHeapQuantizedByteVectorValues .DenseOffHeapVectorValues (dims , size , in .slice ("values" , 0 , in .length ()));
202
+ }
203
+
189
204
// Unsigned int7 byte vectors have values in the range of 0 to 127 (inclusive).
190
205
static final byte MIN_INT7_VALUE = 0 ;
191
206
static final byte MAX_INT7_VALUE = 127 ;
0 commit comments