Add new JMH benchmarks for FusedPQ and graph index queries

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/FusedPQQueryBenchmark.java

@@ -0,0 +1,204 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.bench;
+
+import io.github.jbellis.jvector.graph.GraphIndexBuilder;
+import io.github.jbellis.jvector.graph.GraphSearcher;
+import io.github.jbellis.jvector.graph.ImmutableGraphIndex;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.SearchResult;
+import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;
+import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider;
+import io.github.jbellis.jvector.quantization.PQVectors;
+import io.github.jbellis.jvector.quantization.ProductQuantization;
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter;
+import io.github.jbellis.jvector.graph.disk.OrdinalMapper;
+import io.github.jbellis.jvector.graph.disk.feature.Feature;
+import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
+import io.github.jbellis.jvector.graph.disk.feature.FusedPQ;
+import io.github.jbellis.jvector.graph.disk.feature.InlineVectors;
+import io.github.jbellis.jvector.disk.ReaderSupplierFactory;
+import java.util.*;
+import org.openjdk.jmh.annotations.*;
+import org.openjdk.jmh.infra.Blackhole;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.*;
+import java.util.concurrent.TimeUnit;
+import java.util.function.IntFunction;
+import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED;
+
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@State(Scope.Benchmark)
+@Fork(value = 1, jvmArgsAppend = {"--add-modules=jdk.incubator.vector", "--enable-preview", "-Djvector.experimental.enable_native_vectorization=true"})
+@Warmup(iterations = 3)
+@Measurement(iterations = 5)
+@Threads(1)
+public class FusedPQQueryBenchmark {
+    private static final VectorTypeSupport VECTOR_TYPE_SUPPORT = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+    private OnDiskGraphIndex index;
+    private ArrayList<VectorFloat<?>> queryVectors;
+    private Path indexPath;
+    private Path tempDir;
+
+    @Param({"1536"})
+    int dimension;
+
+    @Param({"96"})
+    int pqM;
+
+    @Param({"100000"})
+    int numBaseVectors;
+
+    @Param({"100"})
+    int numQueryVectors;
+
+    @Param({"10"})
+    int topK;
+
+    @Param({"100"})
+    int efSearch;
+
+    @Setup(Level.Trial)
+    public void setup() throws IOException {
+        System.out.println("Setting up FusedPQ index...");
+
+        // 1. Create base vectors
+        var baseVectors = new ArrayList<VectorFloat<?>>(numBaseVectors);
+        for (int i = 0; i < numBaseVectors; i++) {
+            baseVectors.add(createRandomVector(dimension));
+        }
+        RandomAccessVectorValues floatVectors = new ListRandomAccessVectorValues(baseVectors, dimension);
+
+        // 2. Create query vectors
+        queryVectors = new ArrayList<>(numQueryVectors);
+        for (int i = 0; i < numQueryVectors; i++) {
+            queryVectors.add(createRandomVector(dimension));
+        }
+
+        // 3. Compute PQ compression
+        System.out.println("Computing PQ compression...");
+        boolean centerData = false; // false for DOT_PRODUCT/COSINE
+        var pq = ProductQuantization.compute(floatVectors, pqM, 256, centerData, UNWEIGHTED);
+        var pqVectors = (PQVectors) pq.encodeAll(floatVectors);
+        System.out.printf("PQ: %d subspaces, 256 clusters%n", pqM);
+
+        // 4. Build graph with PQ-compressed vectors
+        System.out.println("Building graph...");
+        int M = 16;
+        int efConstruction = 100;
+        float neighborOverflow = 1.2f;
+        float alpha = 1.2f;
+        boolean addHierarchy = true;
+        boolean refineFinalGraph = true;
+
+        var bsp = BuildScoreProvider.pqBuildScoreProvider(VectorSimilarityFunction.DOT_PRODUCT, pqVectors);
+        var builder = new GraphIndexBuilder(bsp, dimension, M, efConstruction,
+                neighborOverflow, alpha, addHierarchy, refineFinalGraph);
+        var graph = builder.build(floatVectors);
+        System.out.printf("Graph built: %d nodes%n", graph.size(0));
+
+        // 5. Write FusedPQ index to disk
+        System.out.println("Writing FusedPQ index to disk...");
+        tempDir = Files.createTempDirectory("fusedpq-bench");
+        indexPath = tempDir.resolve("fusedpq-index");
+
+        var fusedPQFeature = new FusedPQ(graph.maxDegree(), pq);
+        var inlineVectors = new InlineVectors(dimension);
+
+        try (var writer = new OnDiskGraphIndexWriter.Builder(graph, indexPath)
+                .with(fusedPQFeature)
+                .with(inlineVectors)
+                .withMapper(new OrdinalMapper.IdentityMapper(floatVectors.size() - 1))
+                .build()) {
+
+            var view = graph.getView();
+            Map<FeatureId, IntFunction<Feature.State>> suppliers = new EnumMap<>(FeatureId.class);
+            suppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqVectors, ordinal));
+            suppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(floatVectors.getVector(ordinal)));
+
+            writer.write(suppliers);
+            view.close();
+        }
+
+        builder.close();
+        System.out.printf("Index written: %.2f MB%n", Files.size(indexPath) / 1024.0 / 1024.0);
+
+        // 6. Load the index
+        System.out.println("Loading index...");
+        index = OnDiskGraphIndex.load(ReaderSupplierFactory.open(indexPath));
+        System.out.println("Setup complete!");
+    }
+
+    @TearDown(Level.Trial)
+    public void tearDown() throws IOException {
+        if (index != null) {
+            index.close();
+        }
+        if (indexPath != null && Files.exists(indexPath)) {
+            Files.deleteIfExists(indexPath);
+        }
+        if (tempDir != null && Files.exists(tempDir)) {
+            Files.deleteIfExists(tempDir);
+        }
+        if (queryVectors != null) {
+            queryVectors.clear();
+        }
+    }
+
+     @Benchmark
+    public void queryFusedPQ(Blackhole blackhole) throws IOException {
+        // Perform queries on all query vectors
+        for (VectorFloat<?> queryVector : queryVectors) {
+            try (var view = index.getView()) {
+                var scoringView = (ImmutableGraphIndex.ScoringView) view;
+
+                // Get score functions - FusedPQ for approximate, then rerank
+                var asf = scoringView.approximateScoreFunctionFor(queryVector, VectorSimilarityFunction.DOT_PRODUCT);
+                var reranker = scoringView.rerankerFor(queryVector, VectorSimilarityFunction.DOT_PRODUCT);
+                var ssp = new io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider(asf, reranker);
+
+                // Search
+                var searcher = new GraphSearcher(index);
+                SearchResult result = searcher.search(ssp, topK, efSearch, 1.0f, 0.0f, io.github.jbellis.jvector.util.Bits.ALL);
+
+                blackhole.consume(result);
+            }
+        }
+    }
+
+    private VectorFloat<?> createRandomVector(int dimension) {
+        VectorFloat<?> vector = VECTOR_TYPE_SUPPORT.createFloatVector(dimension);
+        for (int i = 0; i < dimension; i++) {
+            vector.set(i, (float) Math.random());
+        }
+        return vector;
+    }
+}

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/QueryTimeBenchmark.java

@@ -0,0 +1,152 @@
+
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.github.jbellis.jvector.bench;
+
+import io.github.jbellis.jvector.graph.GraphIndexBuilder;
+import io.github.jbellis.jvector.graph.GraphSearcher;
+import io.github.jbellis.jvector.graph.ImmutableGraphIndex;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.SearchResult;
+import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;
+import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider;
+import io.github.jbellis.jvector.quantization.PQVectors;
+import io.github.jbellis.jvector.quantization.ProductQuantization;
+import io.github.jbellis.jvector.util.Bits;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+import org.openjdk.jmh.annotations.*;
+import org.openjdk.jmh.infra.Blackhole;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Benchmarks per-query search latency on a pre-built in-memory index with random vectors.
+ * Index construction happens once per trial in @Setup; only the search is measured.
+ */
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@State(Scope.Thread)
+@Fork(value = 1, jvmArgsAppend = {"--add-modules=jdk.incubator.vector", "--enable-preview", "-Djvector.experimental.enable_native_vectorization=false"})
+@Warmup(iterations = 3)
+@Measurement(iterations = 5)
+@Threads(1)
+public class QueryTimeBenchmark {
+    private static final Logger log = LoggerFactory.getLogger(QueryTimeBenchmark.class);
+    private static final VectorTypeSupport VECTOR_TYPE_SUPPORT = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+    @Param({"768", "1536"})
+    private int originalDimension;
+
+    @Param({"100000"})
+    private int numBaseVectors;
+
+    @Param({"0", "16"})
+    private int numberOfPQSubspaces;
+
+    @Param({"10"})
+    private int topK;
+
+    private RandomAccessVectorValues ravv;
+    private ImmutableGraphIndex graphIndex;
+    private PQVectors pqVectors;
+
+    /** Query vectors rotated through on each invocation to avoid caching effects. */
+    private VectorFloat<?>[] queryVectors;
+    private int queryIndex;
+
+    private static final int NUM_QUERY_VECTORS = 1000;
+    private static final int M = 32;
+    private static final int BEAM_WIDTH = 100;
+
+    @Setup(Level.Trial)
+    public void setup() throws IOException {
+        // Build base vectors
+        var baseVectors = new ArrayList<VectorFloat<?>>(numBaseVectors);
+        for (int i = 0; i < numBaseVectors; i++) {
+            baseVectors.add(createRandomVector(originalDimension));
+        }
+        ravv = new ListRandomAccessVectorValues(baseVectors, originalDimension);
+
+        // Build index once — not measured
+        final BuildScoreProvider buildScoreProvider;
+        if (numberOfPQSubspaces > 0) {
+            log.info("Building with PQ ({} subspaces), dim={}", numberOfPQSubspaces, originalDimension);
+            ProductQuantization pq = ProductQuantization.compute(ravv, numberOfPQSubspaces, 256, true);
+            pqVectors = (PQVectors) pq.encodeAll(ravv);
+            buildScoreProvider = BuildScoreProvider.pqBuildScoreProvider(VectorSimilarityFunction.EUCLIDEAN, pqVectors);
+        } else {
+            log.info("Building with exact scorer, dim={}", originalDimension);
+            pqVectors = null;
+            buildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(ravv, VectorSimilarityFunction.EUCLIDEAN);
+        }
+
+        try (var builder = new GraphIndexBuilder(buildScoreProvider, ravv.dimension(), M, BEAM_WIDTH, 1.2f, 1.2f, true)) {
+            graphIndex = builder.build(ravv);
+        }
+
+        // Pre-generate query vectors so vector creation is not part of the measurement
+        queryVectors = new VectorFloat<?>[NUM_QUERY_VECTORS];
+        for (int i = 0; i < NUM_QUERY_VECTORS; i++) {
+            queryVectors[i] = createRandomVector(originalDimension);
+        }
+        queryIndex = 0;
+    }
+
+    @TearDown(Level.Trial)
+    public void tearDown() {
+        // graphIndex is AutoCloseable only if wrapped; nothing to do for ImmutableGraphIndex
+    }
+
+    /**
+     * Measures the time to execute a single query against the pre-built index.
+     * A pool of pre-generated query vectors is cycled through
+     */
+    @Benchmark
+    public void queryBenchmark(Blackhole blackhole) throws IOException {
+        VectorFloat<?> queryVector = queryVectors[queryIndex];
+        queryIndex = (queryIndex + 1) % NUM_QUERY_VECTORS;
+
+        try (GraphSearcher searcher = new GraphSearcher(graphIndex)) {
+            final SearchResult result;
+            if (pqVectors != null) {
+                var asf = pqVectors.precomputedScoreFunctionFor(queryVector, VectorSimilarityFunction.EUCLIDEAN);
+                var reranker = ravv.rerankerFor(queryVector, VectorSimilarityFunction.EUCLIDEAN);
+                var ssp = new DefaultSearchScoreProvider(asf, reranker);
+                result = searcher.search(ssp, topK, topK * 2, 0.0f, 0.0f, Bits.ALL);
+            } else {
+                var ssp = DefaultSearchScoreProvider.exact(queryVector, VectorSimilarityFunction.EUCLIDEAN, ravv);
+                result = searcher.search(ssp, topK, Bits.ALL);
+            }
+            blackhole.consume(result);
+        }
+    }
+
+    private VectorFloat<?> createRandomVector(int dimension) {
+        VectorFloat<?> vector = VECTOR_TYPE_SUPPORT.createFloatVector(dimension);
+        for (int i = 0; i < dimension; i++) {
+            vector.set(i, (float) Math.random());
+        }
+        return vector;
+    }
+}