Skip to content

Add a new vectorize ES91OSQVectorsScorer in case there is no MMAP #133847

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 2, 2025
Merged

Add a new vectorize ES91OSQVectorsScorer in case there is no MMAP #133847

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e5179b1
Add a new vectorize ES91OSQVectorsScorer in case there is no MMAP
iverase Aug 29, 2025
764acec
Merge branch 'main' into OnHeapES91OSQVectorsScorer
iverase Aug 29, 2025
6fcd266
Merge branch 'main' into OnHeapES91OSQVectorsScorer
benwtrent Aug 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 88 additions & 37 deletions benchmarks/src/main/java/org/elasticsearch/benchmark/vector/OSQScorerBenchmark.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,10 @@
import org.apache.lucene.store.IndexInput;
import org.apache.lucene.store.IndexOutput;
import org.apache.lucene.store.MMapDirectory;
import org.apache.lucene.util.VectorUtil;
import org.apache.lucene.store.NIOFSDirectory;
import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
import org.elasticsearch.common.logging.LogConfigurator;
import org.elasticsearch.core.IOUtils;
import org.elasticsearch.simdvec.ES91OSQVectorsScorer;
import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider;
import org.openjdk.jmh.annotations.Benchmark;
Expand All @@ -29,6 +30,7 @@
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.TearDown;
import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.Blackhole;

Expand Down Expand Up @@ -66,9 +68,14 @@ public class OSQScorerBenchmark {
float centroidDp;

byte[] scratch;
ES91OSQVectorsScorer scorer;
ES91OSQVectorsScorer scorerMmap;
ES91OSQVectorsScorer scorerNfios;

IndexInput in;
Directory dirMmap;
IndexInput inMmap;

Directory dirNiofs;
IndexInput inNiofs;

float[] scratchScores;
float[] corrections;
Expand All @@ -84,18 +91,24 @@ public void setup() throws IOException {
random.nextBytes(binaryVector);
}

Directory dir = new MMapDirectory(Files.createTempDirectory("vectorData"));
IndexOutput out = dir.createOutput("vectors", IOContext.DEFAULT);
dirMmap = new MMapDirectory(Files.createTempDirectory("vectorDataMmap"));
dirNiofs = new NIOFSDirectory(Files.createTempDirectory("vectorDataNFIOS"));
IndexOutput outMmap = dirMmap.createOutput("vectors", IOContext.DEFAULT);
IndexOutput outNfios = dirNiofs.createOutput("vectors", IOContext.DEFAULT);
byte[] correctionBytes = new byte[14 * ES91OSQVectorsScorer.BULK_SIZE];
for (int i = 0; i < numVectors; i += ES91OSQVectorsScorer.BULK_SIZE) {
for (int j = 0; j < ES91OSQVectorsScorer.BULK_SIZE; j++) {
out.writeBytes(binaryVectors[i + j], 0, binaryVectors[i + j].length);
outMmap.writeBytes(binaryVectors[i + j], 0, binaryVectors[i + j].length);
outNfios.writeBytes(binaryVectors[i + j], 0, binaryVectors[i + j].length);
}
random.nextBytes(correctionBytes);
out.writeBytes(correctionBytes, 0, correctionBytes.length);
outMmap.writeBytes(correctionBytes, 0, correctionBytes.length);
outNfios.writeBytes(correctionBytes, 0, correctionBytes.length);
}
out.close();
in = dir.openInput("vectors", IOContext.DEFAULT);
outMmap.close();
outNfios.close();
inMmap = dirMmap.openInput("vectors", IOContext.DEFAULT);
inNiofs = dirNiofs.openInput("vectors", IOContext.DEFAULT);

binaryQueries = new byte[numVectors][4 * length];
for (byte[] binaryVector : binaryVectors) {
Expand All @@ -110,42 +123,40 @@ public void setup() throws IOException {
centroidDp = random.nextFloat();

scratch = new byte[length];
scorer = ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(in, dims);
scorerMmap = ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(inMmap, dims);
scorerNfios = ESVectorizationProvider.getInstance().newES91OSQVectorsScorer(inNiofs, dims);
scratchScores = new float[16];
corrections = new float[3];
}

@TearDown
public void teardown() throws IOException {
IOUtils.close(dirMmap, inMmap, dirNiofs, inNiofs);
}

@Benchmark
public void scoreFromMemorySegmentOnlyVectorMmapScalar(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVector(bh, inMmap, scorerMmap);
}

@Benchmark
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
public void scoreFromArray(Blackhole bh) throws IOException {
for (int j = 0; j < numQueries; j++) {
in.seek(0);
for (int i = 0; i < numVectors; i++) {
in.readBytes(scratch, 0, length);
float qDist = VectorUtil.int4BitDotProduct(binaryQueries[j], scratch);
in.readFloats(corrections, 0, corrections.length);
int addition = Short.toUnsignedInt(in.readShort());
float score = scorer.score(
result.lowerInterval(),
result.upperInterval(),
result.quantizedComponentSum(),
result.additionalCorrection(),
VectorSimilarityFunction.EUCLIDEAN,
centroidDp,
corrections[0],
corrections[1],
addition,
corrections[2],
qDist
);
bh.consume(score);
}
}
public void scoreFromMemorySegmentOnlyVectorMmapVect(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVector(bh, inMmap, scorerMmap);
}

@Benchmark
public void scoreFromMemorySegmentOnlyVectorNiofsScalar(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVector(bh, inNiofs, scorerNfios);
}

@Benchmark
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
public void scoreFromMemorySegmentOnlyVectorNiofsVect(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVector(bh, inNiofs, scorerNfios);
}

private void scoreFromMemorySegmentOnlyVector(Blackhole bh, IndexInput in, ES91OSQVectorsScorer scorer) throws IOException {
for (int j = 0; j < numQueries; j++) {
in.seek(0);
for (int i = 0; i < numVectors; i++) {
Expand All @@ -170,9 +181,29 @@ public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
}
}

@Benchmark
public void scoreFromMemorySegmentOnlyVectorBulkMmapScalar(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVectorBulk(bh, inMmap, scorerMmap);
}

@Benchmark
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
public void scoreFromMemorySegmentOnlyVectorBulkMmapVect(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVectorBulk(bh, inMmap, scorerMmap);
}

@Benchmark
public void scoreFromMemorySegmentOnlyVectorBulkNiofsScalar(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVectorBulk(bh, inNiofs, scorerNfios);
}

@Benchmark
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
public void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh) throws IOException {
public void scoreFromMemorySegmentOnlyVectorBulkNiofsVect(Blackhole bh) throws IOException {
scoreFromMemorySegmentOnlyVectorBulk(bh, inNiofs, scorerNfios);
}

private void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh, IndexInput in, ES91OSQVectorsScorer scorer) throws IOException {
for (int j = 0; j < numQueries; j++) {
in.seek(0);
for (int i = 0; i < numVectors; i += 16) {
Expand All @@ -199,9 +230,29 @@ public void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh) throws IOExceptio
}
}

@Benchmark
public void scoreFromMemorySegmentAllBulkMmapScalar(Blackhole bh) throws IOException {
scoreFromMemorySegmentAllBulk(bh, inMmap, scorerMmap);
}

@Benchmark
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
public void scoreFromMemorySegmentAllBulkMmapVect(Blackhole bh) throws IOException {
scoreFromMemorySegmentAllBulk(bh, inMmap, scorerMmap);
}

@Benchmark
public void scoreFromMemorySegmentAllBulkNiofsScalar(Blackhole bh) throws IOException {
scoreFromMemorySegmentAllBulk(bh, inNiofs, scorerNfios);
}

@Benchmark
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
public void scoreFromMemorySegmentAllBulk(Blackhole bh) throws IOException {
public void scoreFromMemorySegmentAllBulkNiofsVect(Blackhole bh) throws IOException {
scoreFromMemorySegmentAllBulk(bh, inNiofs, scorerNfios);
}

private void scoreFromMemorySegmentAllBulk(Blackhole bh, IndexInput in, ES91OSQVectorsScorer scorer) throws IOException {
for (int j = 0; j < numQueries; j++) {
in.seek(0);
for (int i = 0; i < numVectors; i += 16) {
Expand Down
2 changes: 1 addition & 1 deletion ...a/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentES91OSQVectorsScorer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN;
import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT;

/** Panamized scorer for quantized vectors stored as an {@link IndexInput}. */
/** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */
public final class MemorySegmentES91OSQVectorsScorer extends ES91OSQVectorsScorer {

private static final VectorSpecies<Integer> INT_SPECIES_128 = IntVector.SPECIES_128;
Expand Down
Loading