Add int8 / byte-vector similarity support

.github/workflows/unit-tests.yaml

@@ -45,12 +45,38 @@ jobs:
       - name: Set up GCC
         run: |
           sudo apt install -y gcc
-      - name: Install Meson and Ninja
+      - name: Install Meson, Ninja, and GTest
         run: |
-          sudo apt update && sudo apt install -y meson ninja-build
+          sudo apt update && sudo apt install -y meson ninja-build pkg-config libgtest-dev
       - uses: actions/checkout@v4
       - name: Initialize Git Submodules
         run: git submodule update --init
+
+      - name: Build test_simd_kernels (native C++)
+        working-directory: jvector-native/src/main/native
+        run: |
+          meson setup build --wipe
+          ninja -C build test_simd_kernels
+
+      - name: Run test_simd_kernels — no ISA cap (auto-detect)
+        if: matrix.max_isa == 'avx512f'
+        working-directory: jvector-native/src/main/native
+        run: ./build/test_simd_kernels
+
+      - name: Run test_simd_kernels — capped at avx2
+        if: matrix.max_isa == 'avx2'
+        working-directory: jvector-native/src/main/native
+        env:
+          JVECTOR_MAX_ISA: avx2
+        run: ./build/test_simd_kernels
+
+      - name: Run test_simd_kernels — capped at sse42
+        if: matrix.max_isa == 'sse42'
+        working-directory: jvector-native/src/main/native
+        env:
+          JVECTOR_MAX_ISA: sse42
+        run: ./build/test_simd_kernels
+
       - name: Set up JDK ${{ matrix.jdk }}
         uses: actions/setup-java@v3
         with:

jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java

@@ -24,9 +24,12 @@
 import io.github.jbellis.jvector.graph.diversity.VamanaDiversityProvider;
 import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;
 import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider;
 import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
 import io.github.jbellis.jvector.util.*;
+import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.types.ByteSequence;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -74,6 +77,10 @@ public class GraphIndexBuilder implements Closeable, Accountable {
 
     private final BuildScoreProvider scoreProvider;
 
+    // set only when built from a byte-vector constructor; used by addGraphNode(int, ByteSequence<?>)
+    private RandomAccessByteVectorValues byteVectorValues;
+    private ByteVectorSimilarityFunction byteVectorSimilarityFunction;
+
     private final ForkJoinPool simdExecutor;
     private final ForkJoinPool parallelExecutor;
 
@@ -97,6 +104,39 @@ public class GraphIndexBuilder implements Closeable, Accountable {
      *                         an HNSW graph will be created, which is usually not what you want.
      * @param addHierarchy     whether we want to add an HNSW-style hierarchy on top of the Vamana index.
      */
+    /**
+     * Convenience constructor for building a byte-vector (int8) graph.
+     * See {@link #GraphIndexBuilder(RandomAccessVectorValues, VectorSimilarityFunction, int, int, float, float, boolean)}
+     * for the float equivalent.
+     *
+     * @param vectorValues      the int8 vectors whose relations are represented by the graph
+     * @param similarityFunction the similarity metric to use during construction
+     * @param M                 the maximum number of connections a node can have
+     * @param beamWidth         the size of the beam search to use when finding nearest neighbors
+     * @param neighborOverflow  the ratio of extra neighbors to allow temporarily when inserting a node
+     * @param alpha             how aggressive pruning diverse neighbors should be
+     * @param addHierarchy      whether to add an HNSW-style hierarchy on top of the Vamana index
+     */
+    public GraphIndexBuilder(RandomAccessByteVectorValues vectorValues,
+                             ByteVectorSimilarityFunction similarityFunction,
+                             int M,
+                             int beamWidth,
+                             float neighborOverflow,
+                             float alpha,
+                             boolean addHierarchy)
+    {
+        this(BuildScoreProvider.byteVectorScoreProvider(vectorValues, similarityFunction),
+                vectorValues.dimension(),
+                M,
+                beamWidth,
+                neighborOverflow,
+                alpha,
+                addHierarchy,
+                true);
+        this.byteVectorValues = vectorValues;
+        this.byteVectorSimilarityFunction = similarityFunction;
+    }
+
     public GraphIndexBuilder(RandomAccessVectorValues vectorValues,
                              VectorSimilarityFunction similarityFunction,
                              int M,
@@ -446,6 +486,25 @@ public ImmutableGraphIndex build(RandomAccessVectorValues ravv) {
         cleanup();
         return graph;
     }
+
+    /**
+     * Builds the graph from a {@link RandomAccessByteVectorValues}.
+     * Each node is scored via the {@link BuildScoreProvider} supplied at construction time,
+     * so all comparisons remain byte×byte with no float round-trip.
+     */
+    public ImmutableGraphIndex build(RandomAccessByteVectorValues ravv) {
+        int size = ravv.size();
+
+        simdExecutor.submit(() -> {
+            IntStream.range(0, size).parallel().forEach(node -> {
+                var ssp = scoreProvider.searchProviderFor(node);
+                addGraphNode(node, ssp);
+            });
+        }).join();
+
+        cleanup();
+        return graph;
+    }
     /**
      * Validates that the current entry node has been completely added.
      */
@@ -590,6 +649,26 @@ public long addGraphNode(int node, VectorFloat<?> vector) {
         return addGraphNode(node, ssp);
     }
 
+    /**
+     * Inserts a node with the given int8 byte vector into the graph.
+     *
+     * @param node   the node ID to add
+     * @param vector the byte vector to add
+     * @return an estimate of the number of extra bytes used by the graph after adding the given node
+     * @throws UnsupportedOperationException if this builder was not constructed with a byte-vector score provider
+     */
+    public long addGraphNode(int node, ByteSequence<?> vector) {
+        if (byteVectorValues == null) {
+            throw new UnsupportedOperationException(
+                    "addGraphNode(int, ByteSequence<?>) requires a byte-vector GraphIndexBuilder; " +
+                    "use the GraphIndexBuilder(RandomAccessByteVectorValues, ...) constructor");
+        }
+        var bvsf = byteVectorSimilarityFunction;
+        var ravv = byteVectorValues;
+        var sf = (ScoreFunction.ExactScoreFunction) node2 -> bvsf.compare(vector, ravv.getVector(node2));
+        return addGraphNode(node, new DefaultSearchScoreProvider(sf));
+    }
+
     /**
      * Inserts a node with the given vector value to the graph.
      *

jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java

@@ -0,0 +1,70 @@
+/*
+ * 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.graph;
+
+import io.github.jbellis.jvector.vector.types.ByteSequence;
+
+import java.util.List;
+
+/**
+ * A List-backed implementation of the {@link RandomAccessByteVectorValues} interface.
+ * <p>
+ * It is acceptable to provide this class to a GraphBuilder, and then continue
+ * to add vectors to the backing List as you add to the graph.
+ * <p>
+ * This will be as threadsafe as the provided List.
+ */
+public class ListRandomAccessByteVectorValues implements RandomAccessByteVectorValues {
+    private final List<ByteSequence<?>> vectors;
+    private final int dimension;
+
+    /**
+     * Construct a new instance of {@link ListRandomAccessByteVectorValues}.
+     *
+     * @param vectors   a (potentially mutable) list of byte vectors.
+     * @param dimension the dimension of the vectors.
+     */
+    public ListRandomAccessByteVectorValues(List<ByteSequence<?>> vectors, int dimension) {
+        this.vectors = vectors;
+        this.dimension = dimension;
+    }
+
+    @Override
+    public int size() {
+        return vectors.size();
+    }
+
+    @Override
+    public int dimension() {
+        return dimension;
+    }
+
+    @Override
+    public ByteSequence<?> getVector(int nodeId) {
+        return vectors.get(nodeId);
+    }
+
+    @Override
+    public boolean isValueShared() {
+        return false;
+    }
+
+    @Override
+    public ListRandomAccessByteVectorValues copy() {
+        return this;
+    }
+}

jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java

@@ -0,0 +1,74 @@
+/*
+ * 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.graph;
+
+import io.github.jbellis.jvector.util.ExplicitThreadLocal;
+import io.github.jbellis.jvector.vector.types.ByteSequence;
+
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+
+/**
+ * Provides random access to byte (int8) vectors by dense ordinal.
+ * <p>
+ * This is the byte-vector parallel to {@link RandomAccessVectorValues}.
+ * It is used by graph-based index builders and searchers that operate natively
+ * on int8 vectors without a float32 round-trip.
+ */
+public interface RandomAccessByteVectorValues {
+    Logger LOG = Logger.getLogger(RandomAccessByteVectorValues.class.getName());
+
+    /** Return the number of vector values. */
+    int size();
+
+    /** Return the dimension of the returned vector values. */
+    int dimension();
+
+    /**
+     * Return the byte vector indexed at the given ordinal.
+     *
+     * @param nodeId a valid ordinal, &ge; 0 and &lt; {@link #size()}.
+     */
+    ByteSequence<?> getVector(int nodeId);
+
+    /**
+     * @return true iff the vector returned by {@link #getVector} is shared across calls.
+     * A shared vector is only valid until the next call to {@link #getVector} overwrites it.
+     */
+    boolean isValueShared();
+
+    /**
+     * Creates a new copy of this {@link RandomAccessByteVectorValues}.
+     * Un-shared implementations may simply return {@code this}.
+     */
+    RandomAccessByteVectorValues copy();
+
+    /**
+     * Returns a supplier of thread-local copies of the RABVV.
+     */
+    default Supplier<RandomAccessByteVectorValues> threadLocalSupplier() {
+        if (!isValueShared()) {
+            return () -> this;
+        }
+
+        if (this instanceof AutoCloseable) {
+            LOG.warning("RABVV is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks");
+        }
+        var tl = ExplicitThreadLocal.withInitial(this::copy);
+        return tl::get;
+    }
+}

jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java

@@ -16,8 +16,10 @@
 
 package io.github.jbellis.jvector.graph.similarity;
 
+import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues;
 import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
 import io.github.jbellis.jvector.graph.RemappedRandomAccessVectorValues;
+import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction;
 import io.github.jbellis.jvector.quantization.BQVectors;
 import io.github.jbellis.jvector.quantization.PQVectors;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
@@ -211,6 +213,63 @@ public VectorFloat<?> approximateCentroid() {
         };
     }
 
+    /**
+     * Returns a BSP that performs exact score comparisons using the given
+     * {@link RandomAccessByteVectorValues} and {@link ByteVectorSimilarityFunction}.
+     * All scoring is byte×byte with no float32 round-trip.
+     */
+    static BuildScoreProvider byteVectorScoreProvider(RandomAccessByteVectorValues ravv, ByteVectorSimilarityFunction bvsf) {
+        var vectors     = ravv.threadLocalSupplier();
+        var vectorsCopy = ravv.threadLocalSupplier();
+
+        return new BuildScoreProvider() {
+            @Override
+            public boolean isExact() {
+                return true;
+            }
+
+            @Override
+            public VectorFloat<?> approximateCentroid() {
+                var vv = vectors.get();
+                var centroid = vts.createFloatVector(vv.dimension());
+                for (int i = 0; i < vv.size(); i++) {
+                    var v = vv.getVector(i);
+                    for (int d = 0; d < vv.dimension(); d++) {
+                        centroid.set(d, centroid.get(d) + v.get(d));
+                    }
+                }
+                VectorUtil.scale(centroid, 1.0f / vv.size());
+                return centroid;
+            }
+
+            @Override
+            public SearchScoreProvider searchProviderFor(VectorFloat<?> vector) {
+                throw new UnsupportedOperationException(
+                        "byteVectorScoreProvider does not support float query vectors; use searchProviderFor(int node)");
+            }
+
+            @Override
+            public SearchScoreProvider searchProviderFor(int node1) {
+                var v  = vectors.get().getVector(node1);
+                var vc = vectorsCopy.get();
+                var sf = (ScoreFunction.ExactScoreFunction) node2 -> bvsf.compare(v, vc.getVector(node2));
+                return new DefaultSearchScoreProvider(sf);
+            }
+
+            @Override
+            public SearchScoreProvider diversityProviderFor(int node1) {
+                return searchProviderFor(node1);
+            }
+
+            @Override
+            public ScoreFunction diversityScoreFunctionFor(int node1) {
+                var v  = vectors.get().getVector(node1);
+                var vc = vectorsCopy.get();
+                return (ScoreFunction.ExactScoreFunction) node2 -> bvsf.compare(v, vc.getVector(node2));
+            }
+        };
+    }
+
     static BuildScoreProvider bqBuildScoreProvider(BQVectors bqv) {
         return new BuildScoreProvider() {
             @Override