Add int8 / byte-vector similarity support

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

@@ -19,6 +19,7 @@
 import io.github.jbellis.jvector.disk.ReaderSupplier;
 import io.github.jbellis.jvector.disk.ReaderSupplierFactory;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.FloatDataSet;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSetInfo;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
 import io.github.jbellis.jvector.example.reporting.GitInfo;
@@ -251,7 +252,7 @@ private static void writeCompletedCount(int count) {
     private List<VectorFloat<?>> queryVectors;
     private List<VectorFloat<?>> baseVectors;
     private List<? extends List<Integer>> groundTruth;
-    private DataSet ds;
+    private FloatDataSet ds;
     private VectorSimilarityFunction similarityFunction;
 
     private final List<OnDiskGraphIndex> graphs = new ArrayList<>();
@@ -393,7 +394,7 @@ public void setup() throws Exception {
             boolean needsRecallData = measureRecall && workloadMode != WorkloadMode.PARTITION;
 
             if (needsBaseVectors) {
-                ds = DataSets.loadDataSet(datasetNames)
+                ds = (FloatDataSet) DataSets.loadDataSet(datasetNames)
                         .orElseThrow(() -> new RuntimeException("Dataset not found: " + datasetNames))
                         .getDataSet();
 
@@ -432,7 +433,7 @@ public void setup() throws Exception {
                 dimension = -1;
 
                 if (needsRecallData) {
-                    ds = DataSets.loadDataSet(datasetNames)
+                    ds = (FloatDataSet) DataSets.loadDataSet(datasetNames)
                             .orElseThrow(() -> new RuntimeException("Dataset not found: " + datasetNames))
                             .getDataSet();
                     queryVectors = ds.getQueryVectors();
@@ -595,7 +596,7 @@ private void verifyPartitionsExist(Path partitionsDir, int numPartitions) {
         }
     }
 
-    private void buildPartitions(DataSet ds, List<VectorFloat<?>> baseVectors) throws Exception {
+    private void buildPartitions(DataSet<?> ds, List<VectorFloat<?>> baseVectors) throws Exception {
 
         var partitionedData = DataSetPartitioner.partition(baseVectors, numPartitions, splitDistribution);
         vectorsPerSourceCount = partitionedData.sizes;

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

@@ -23,6 +23,7 @@
 import io.github.jbellis.jvector.graph.SearchResult.NodeScore;
 import io.github.jbellis.jvector.graph.diversity.VamanaDiversityProvider;
 import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;
+import io.github.jbellis.jvector.graph.VectorValues;
 import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
 import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
 import io.github.jbellis.jvector.management.CompressionType;
@@ -32,7 +33,9 @@
 import io.github.jbellis.jvector.quantization.PQVectors;
 import io.github.jbellis.jvector.quantization.ProductQuantization;
 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;
@@ -114,6 +117,11 @@ private static BuildScoreProvider getBuildScoreProvider(RandomAccessVectorValues
         }
     }
 
+    private static BuildScoreProvider getBuildScoreProvider(RandomAccessByteVectorValues vectorValues, ByteVectorSimilarityFunction similarityFunction) {
+        // Byte vectors do not support PQ/BQ build-time compression; always use the direct provider.
+        return BuildScoreProvider.byteVectorScoreProvider(vectorValues, similarityFunction);
+    }
+
     /**
      * Reads all the vectors from vector values, builds a graph connecting them by their dense
      * ordinals, using the given hyperparameter settings, and returns the resulting graph.
@@ -407,6 +415,31 @@ public static Builder builder(RandomAccessVectorValues vectorValues, VectorSimil
         return builder(vectorValues, similarityFunction, List.of(M));
     }
 
+    /**
+     * Entry point for the fluent builder for int8 byte vectors, resolving the {@link BuildScoreProvider}
+     * from raw byte vectors via {@link BuildScoreProvider#byteVectorScoreProvider}.
+     *
+     * @param vectorValues       the int8 byte vectors whose relations are represented by the graph
+     * @param similarityFunction the byte-vector similarity function to score vectors with
+     * @param maxDegrees         the maximum number of connections a node can have in each layer; if fewer entries
+     *                           are specified than the number of layers, the last entry is used for all remaining layers.
+     */
+    public static Builder builder(RandomAccessByteVectorValues vectorValues, ByteVectorSimilarityFunction similarityFunction, List<Integer> maxDegrees) {
+        return new Builder(getBuildScoreProvider(vectorValues, similarityFunction), vectorValues.dimension(), maxDegrees);
+    }
+
+    /**
+     * Entry point for the fluent builder for int8 byte vectors, for the common case of a single
+     * (non-hierarchical) max degree.
+     *
+     * @param vectorValues       the int8 byte vectors whose relations are represented by the graph
+     * @param similarityFunction the byte-vector similarity function to score vectors with
+     * @param M                  the maximum number of connections a node can have
+     */
+    public static Builder builder(RandomAccessByteVectorValues vectorValues, ByteVectorSimilarityFunction similarityFunction, int M) {
+        return builder(vectorValues, similarityFunction, List.of(M));
+    }
+
     /**
      * Entry point for the fluent builder, building from an existing {@link MutableGraphIndex} (e.g. one just
      * loaded from disk) rather than constructing a fresh {@link OnHeapGraphIndex}. {@code addHierarchy} is not
@@ -695,13 +728,17 @@ public static GraphIndexBuilder rescore(GraphIndexBuilder other, BuildScoreProvi
         return newBuilder;
     }
 
-    public ImmutableGraphIndex build(RandomAccessVectorValues ravv) {
-        var vv = ravv.threadLocalSupplier();
+    /**
+     * Builds the graph from any {@link VectorValues} source — works for both
+     * {@link RandomAccessVectorValues} (float32) and {@link RandomAccessByteVectorValues} (int8).
+     * The score provider supplied at construction time determines how vectors are compared.
+     */
+    public ImmutableGraphIndex build(VectorValues<?> ravv) {
         int size = ravv.size();
 
         simdExecutor.submit(() -> {
             IntStream.range(0, size).parallel().forEach(node -> {
-                addGraphNode(node, vv.get().getVector(node));
+                addGraphNode(node, scoreProvider.searchProviderFor(node));
             });
         }).join();
 
@@ -852,6 +889,19 @@ 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) {
+        var ssp = scoreProvider.searchProviderFor(vector);
+        return addGraphNode(node, ssp);
+    }
+
     /**
      * Inserts a node with the given vector value to the graph.
      *
@@ -1337,4 +1387,4 @@ public static ImmutableGraphIndex buildAndMergeNewNodes(RandomAccessReader in,
             return builder.getGraph();
         }
     }
-}
+}

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,37 @@
+/*
+ * 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;
+
+/**
+ * Provides random access to byte (int8) vectors by dense ordinal.
+ * <p>
+ * This is the byte-vector parallel to {@link RandomAccessVectorValues}.
+ * Both extend the common super-interface {@link VectorValues}.
+ * It is used by graph-based index builders and searchers that operate natively
+ * on int8 vectors without a float32 round-trip.
+ */
+public interface RandomAccessByteVectorValues extends VectorValues<ByteSequence<?>> {
+
+    /**
+     * Creates a new copy of this {@link RandomAccessByteVectorValues}.
+     * Un-shared implementations may simply return {@code this}.
+     */
+    @Override
+    RandomAccessByteVectorValues copy();
+}

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

@@ -30,14 +30,15 @@
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 
 import java.util.function.Supplier;
-import java.util.logging.Logger;
 
 /**
- * Provides random access to vectors by dense ordinal. This interface is used by graph-based
+ * Provides random access to float32 vectors by dense ordinal. This interface is used by graph-based
  * implementations of KNN search.
+ * <p>
+ * For int8 vectors see {@link RandomAccessByteVectorValues}.
+ * Both extend the common super-interface {@link VectorValues}.
  */
-public interface RandomAccessVectorValues {
-    Logger LOG = Logger.getLogger(RandomAccessVectorValues.class.getName());
+public interface RandomAccessVectorValues extends VectorValues<VectorFloat<?>> {
 
     /**
      * Return the number of vector values.
@@ -46,23 +47,9 @@ public interface RandomAccessVectorValues {
      * (1) implementing a threadsafe, un-shared RAVV, where `copy` returns `this`, or
      * (2) implementing a fixed-size RAVV.
      */
+    @Override
     int size();
 
-    /** Return the dimension of the returned vector values */
-    int dimension();
-
-    /**
-     * Return the vector value indexed at the given ordinal.
-     *
-     * <p>For performance, implementations are free to re-use the same object across invocations.
-     * That is, you will get back the same VectorFloat&lt;?&gt;
-     * reference (for instance) for every requested ordinal. If you want to use those values across
-     * calls, you should make a copy.
-     *
-     * @param nodeId a valid ordinal, &ge; 0 and &lt; {@link #size()}.
-     */
-    VectorFloat<?> getVector(int nodeId);
-
     @Deprecated
     default VectorFloat<?> vectorValue(int targetOrd) {
         return getVector(targetOrd);
@@ -78,36 +65,16 @@ default void getVectorInto(int node, VectorFloat<?> destinationVector, int offse
         destinationVector.copyFrom(getVector(node), 0, offset, dimension());
     }
 
-    /**
-     * @return true iff the vector returned by `getVector` is shared.  A shared vector will
-     * only be valid until the next call to getVector overwrites it.
-     */
-    boolean isValueShared();
-
     /**
      * Creates a new copy of this {@link RandomAccessVectorValues}. This is helpful when you need to
      * access different values at once, to avoid overwriting the underlying float vector returned by
      * a shared {@link RandomAccessVectorValues#getVector}.
      * <p>
      * Un-shared implementations may simply return `this`.
      */
+    @Override
     RandomAccessVectorValues copy();
 
-    /**
-     * Returns a supplier of thread-local copies of the RAVV.
-     */
-    default Supplier<RandomAccessVectorValues> threadLocalSupplier() {
-        if (!isValueShared()) {
-            return () -> this;
-        }
-
-        if (this instanceof AutoCloseable) {
-            LOG.warning("RAVV is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks");
-        }
-        var tl = ExplicitThreadLocal.withInitial(this::copy);
-        return tl::get;
-    }
-
     /**
      * Convenience method to create an ExactScoreFunction for reranking.  The resulting function is NOT thread-safe.
      */