Add int8 / byte-vector similarity support

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/disk/AbstractGraphIndexWriter.java

@@ -21,6 +21,7 @@
 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.FusedFeature;
+import io.github.jbellis.jvector.graph.disk.feature.InlineByteVectors;
 import io.github.jbellis.jvector.graph.disk.feature.InlineVectors;
 import io.github.jbellis.jvector.graph.disk.feature.NVQ;
 import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature;
@@ -386,6 +387,8 @@ public K build() throws IOException {
             int dimension;
             if (features.containsKey(FeatureId.INLINE_VECTORS)) {
                 dimension = ((InlineVectors) features.get(FeatureId.INLINE_VECTORS)).dimension();
+            } else if (features.containsKey(FeatureId.INLINE_BYTE_VECTORS)) {
+                dimension = ((InlineByteVectors) features.get(FeatureId.INLINE_BYTE_VECTORS)).dimension();
             } else if (features.containsKey(FeatureId.NVQ_VECTORS)) {
                 dimension = ((NVQ) features.get(FeatureId.NVQ_VECTORS)).dimension();
             } else if (features.containsKey(FeatureId.SEPARATED_VECTORS)) {

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

@@ -27,6 +27,7 @@
 import io.github.jbellis.jvector.graph.disk.feature.FeatureSource;
 import io.github.jbellis.jvector.graph.disk.feature.FusedPQ;
 import io.github.jbellis.jvector.graph.disk.feature.FusedFeature;
+import io.github.jbellis.jvector.graph.disk.feature.InlineByteVectors;
 import io.github.jbellis.jvector.graph.disk.feature.InlineVectors;
 import io.github.jbellis.jvector.graph.disk.feature.NVQ;
 import io.github.jbellis.jvector.graph.disk.feature.SeparatedFeature;
@@ -36,8 +37,10 @@
 import java.util.ArrayList;
 import io.github.jbellis.jvector.util.Bits;
 import io.github.jbellis.jvector.util.RamUsageEstimator;
+import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.ByteSequence;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
 
@@ -580,6 +583,41 @@ public void getVectorInto(int node, VectorFloat<?> vector, int offset) {
             }
         }
 
+        /**
+         * Returns the signed int8 vector stored for {@code node} via the
+         * {@link FeatureId#INLINE_BYTE_VECTORS} feature.
+         *
+         * @throws UnsupportedOperationException if the graph was not written with
+         *         {@link InlineByteVectors}
+         */
+        public ByteSequence<?> getByteVector(int node) {
+            if (!features.containsKey(FeatureId.INLINE_BYTE_VECTORS)) {
+                throw new UnsupportedOperationException("No inline byte vectors in this graph");
+            }
+            try {
+                long diskOffset = offsetFor(node, FeatureId.INLINE_BYTE_VECTORS);
+                reader.seek(diskOffset);
+                return vectorTypeSupport.readByteSequence(reader, dimension);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }
+
+        /**
+         * Returns a {@link ScoreFunction.ExactScoreFunction} that scores candidates by reading
+         * their int8 vectors from disk and comparing them byte×byte against {@code queryBytes}.
+         *
+         * @throws UnsupportedOperationException if the graph was not written with
+         *         {@link InlineByteVectors}
+         */
+        public ScoreFunction.ExactScoreFunction byteVectorRerankerFor(ByteSequence<?> queryBytes,
+                                                                       ByteVectorSimilarityFunction bvsf) {
+            if (!features.containsKey(FeatureId.INLINE_BYTE_VECTORS)) {
+                throw new UnsupportedOperationException("No inline byte vectors in this graph");
+            }
+            return node -> bvsf.compare(queryBytes, getByteVector(node));
+        }
+
         public NodesIterator getNeighborsIterator(int level, int node) {
             try {
                 int[] stored;

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java

@@ -33,7 +33,8 @@ public enum FeatureId {
     FUSED_PQ(FusedPQ::load),
     NVQ_VECTORS(NVQ::load),
     SEPARATED_VECTORS(SeparatedVectors::load),
-    SEPARATED_NVQ(SeparatedNVQ::load);
+    SEPARATED_NVQ(SeparatedNVQ::load),
+    INLINE_BYTE_VECTORS(InlineByteVectors::load);
 
     private final BiFunction<CommonHeader, RandomAccessReader, Feature> loader;