Add binary quantization (#135)

* extract CompressedVectors interface, implemented by PQVectors

* BQ implementation

* introduce VectorComporessor interface so Bench can mix PQ and BQ in the same grid

* remove diskGrid/useDisk so we only do the uncompressed queries once

Author: jbellis

README.md

@@ -62,12 +62,16 @@ JVector implements [DiskANN](https://suhasjs.github.io/files/diskann_neurips19.p
 search, meaning that vectors can be compressed using product quantization so that searches
 can be performed using the compressed representation that is kept in memory.  You can enable
 this with the following steps:
-- Create a [`ProductQuantization`](./jvector-base/src/main/java/io/github/jbellis/jvector/pq/ProductQuantization.java) object with your vectors using `ProductQuantization.compute`.  This will take some time
-  to compute the codebooks.
-- Use `ProductQuantization::encode` or `encodeAll` to encode your vectors.
-- Create a [`CompressedVectors`](./jvector-base/src/main/java/io/github/jbellis/jvector/disk/CompressedVectors.java) object from the encoded vectors.
-- Create a [`NeighborSimilarity.ApproximateScoreFunction`](./jvector-base/src/main/java/io/github/jbellis/jvector/graph/NeighborSimilarity.java) for your query that uses the
-  `ProductQuantization` object and `CompressedVectors` to compute scores, and pass this
+- Create a [`VectorCompressor`](./jvector-base/src/main/java/io/github/jbellis/jvector/pq/VectorCompressor.java) object with your vectors using either `ProductQuantization.compute`
+- or `BinaryQuantization.compute`.  PQ is more flexible than BQ and is less lossy: even at the same compressed size, 
+  in the datasets tested by Bench, only the ada002 vectors in the wikipedia dataset
+  are large enough and/or overparameterized enough to benefit from BQ while achieving recall
+  competitive with PQ.  However, if you are dealing with very large vectors and/or your
+  recall requirement is not strict, you may still want to try BQ since it is MUCH faster to both compute and search with.
+- Use `VectorCompressor::encode` or `encodeAll` to encode your vectors, then call
+  `VectorCompressor::createCompressedVectors` to create a `CompressedVectors` object.
+- Call `CompressedVectors::approximateScoreFunctionFor` to create a [`NeighborSimilarity.ApproximateScoreFunction`](./jvector-base/src/main/java/io/github/jbellis/jvector/graph/NeighborSimilarity.java) for your query that uses the
+  compressed vectors to accelerate search, and pass this
   to the `GraphSearcher.search` method.
 
 ## Saving and loading indexes
@@ -77,7 +81,7 @@ this with the following steps:
   implementation of [`RandomAccessReader`](./jvector-base/src/main/java/io/github/jbellis/jvector/disk/RandomAccessReader.java) and the related `ReaderSupplier` to wrap your
   preferred i/o class for best performance. See `SimpleMappedReader` and `SimpleMappedReaderSupplier` for an example.
 - Building a graph does not technically require your RandomAccessVectorValues object
-  to live in memory, but it will perform much better if it does.  OnDiskGraphIndex,
+  to live in memory, but it will perform much better if it does.  `OnDiskGraphIndex`,
   by contrast, is designed to live on disk and use minimal memory otherwise.
 - You can optionally wrap `OnDiskGraphIndex` in a [`CachingGraphIndex`](./jvector-base/src/main/java/io/github/jbellis/jvector/disk/CachingGraphIndex.java) to keep the most commonly accessed
   nodes (the ones nearest to the graph entry point) in memory.

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

@@ -16,7 +16,6 @@
 
 package io.github.jbellis.jvector.disk;
 
-import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
 
@@ -26,12 +25,4 @@ public static void writeFloats(DataOutput out, float[] v) throws IOException {
             out.writeFloat(a);
         }
     }
-
-    public static float[] readFloats(DataInput in, int size) throws IOException {
-        var v = new float[size];
-        for (int i = 0; i < size; i++) {
-            v[i] = in.readFloat();
-        }
-        return v;
-    }
 }

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

@@ -31,6 +31,8 @@ public interface RandomAccessReader extends AutoCloseable {
 
     void readFully(float[] floats) throws IOException;
 
+    void readFully(long[] vector) throws IOException;
+
     void read(int[] ints, int offset, int count) throws IOException;
 
     void close() throws IOException;

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

@@ -83,6 +83,13 @@ public void readFully(byte[] b) {
         mbb.get(b);
     }
 
+    @Override
+    public void readFully(long[] vector) throws IOException {
+        for (int i = 0; i < vector.length; i++) {
+            vector[i] = mbb.getLong();
+        }
+    }
+
     @Override
     public int readInt() {
         return mbb.getInt();

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

@@ -171,8 +171,7 @@ SearchResult searchInternal(
     }
     while (candidates.size() > 0 && !resultsQueue.incomplete()) {
       // get the best candidate (closest or best scoring)
-      float topCandidateSimilarity = candidates.topScore();
-      if (topCandidateSimilarity < minAcceptedSimilarity) {
+      if (candidates.topScore() < minAcceptedSimilarity) {
         break;
       }
 

jvector-base/src/main/java/io/github/jbellis/jvector/pq/BQVectors.java

@@ -0,0 +1,120 @@
+/*
+ * 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.pq;
+
+import io.github.jbellis.jvector.disk.RandomAccessReader;
+import io.github.jbellis.jvector.graph.NeighborSimilarity;
+import io.github.jbellis.jvector.util.RamUsageEstimator;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorUtil;
+
+import java.io.DataOutput;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+
+public class BQVectors implements CompressedVectors {
+    private final BinaryQuantization bq;
+    private final long[][] compressedVectors;
+
+    public BQVectors(BinaryQuantization bq, long[][] compressedVectors) {
+        this.bq = bq;
+        this.compressedVectors = compressedVectors;
+    }
+
+    @Override
+    public void write(DataOutput out) throws IOException {
+        // BQ centering data
+        bq.write(out);
+
+        // compressed vectors
+        out.writeInt(compressedVectors.length);
+        if (compressedVectors.length <= 0) {
+            return;
+        }
+        out.writeInt(compressedVectors[0].length);
+        for (var v : compressedVectors) {
+            for (int i = 0; i < v.length; i++) {
+                out.writeLong(v[i]);
+            }
+        }
+    }
+
+    public static BQVectors load(RandomAccessReader in, int offset) throws IOException {
+        in.seek(offset);
+
+        // BQ
+        var bq = BinaryQuantization.load(in);
+
+        // check validity of compressed vectors header
+        int size = in.readInt();
+        if (size < 0) {
+            throw new IOException("Invalid compressed vector count " + size);
+        }
+        var compressedVectors = new long[size][];
+        if (size == 0) {
+            return new BQVectors(bq, compressedVectors);
+        }
+        int compressedLength = in.readInt();
+        if (compressedLength < 0) {
+            throw new IOException("Invalid compressed vector dimension " + compressedLength);
+        }
+
+        // read the compressed vectors
+        for (int i = 0; i < size; i++)
+        {
+            long[] vector = new long[compressedLength];
+            in.readFully(vector);
+            compressedVectors[i] = vector;
+        }
+
+        return new BQVectors(bq, compressedVectors);
+    }
+
+    @Override
+    public NeighborSimilarity.ApproximateScoreFunction approximateScoreFunctionFor(float[] q, VectorSimilarityFunction similarityFunction) {
+        var qBQ = bq.encode(q);
+        return node2 -> {
+            var vBQ = compressedVectors[node2];
+            return 1 - (float) VectorUtil.hammingDistance(qBQ, vBQ) / q.length;
+        };
+    }
+
+    public long[] get(int i) {
+        return compressedVectors[i];
+    }
+
+    @Override
+    public long ramBytesUsed() {
+        return compressedVectors.length * RamUsageEstimator.sizeOf(compressedVectors[0]);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        BQVectors bqVectors = (BQVectors) o;
+        return Objects.equals(bq, bqVectors.bq) && Arrays.deepEquals(compressedVectors, bqVectors.compressedVectors);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = Objects.hash(bq);
+        result = 31 * result + Arrays.deepHashCode(compressedVectors);
+        return result;
+    }
+}

jvector-base/src/main/java/io/github/jbellis/jvector/pq/BinaryQuantization.java

@@ -0,0 +1,134 @@
+/*
+ * 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.pq;
+
+import io.github.jbellis.jvector.disk.Io;
+import io.github.jbellis.jvector.disk.RandomAccessReader;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
+import io.github.jbellis.jvector.util.PoolingSupport;
+import io.github.jbellis.jvector.vector.VectorUtil;
+
+import java.io.DataOutput;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.lang.Math.min;
+
+/**
+ * Binary Quantization of float vectors: each float is compressed to a single bit,
+ * and similarity is computed with a simple Hamming distance.
+ */
+public class BinaryQuantization implements VectorCompressor<long[]> {
+    private final float[] globalCentroid;
+
+    public BinaryQuantization(float[] globalCentroid) {
+        this.globalCentroid = globalCentroid;
+    }
+
+    public static BinaryQuantization compute(RandomAccessVectorValues<float[]> ravv) {
+        // limit the number of vectors we train on
+        var P = min(1.0f, ProductQuantization.MAX_PQ_TRAINING_SET_SIZE / (float) ravv.size());
+        var ravvCopy = ravv.isValueShared() ? PoolingSupport.newThreadBased(ravv::copy) : PoolingSupport.newNoPooling(ravv);
+        var vectors = IntStream.range(0, ravv.size()).parallel()
+                .filter(i -> ThreadLocalRandom.current().nextFloat() < P)
+                .mapToObj(targetOrd -> {
+                    try (var pooledRavv = ravvCopy.get()) {
+                        var localRavv = pooledRavv.get();
+                        float[] v = localRavv.vectorValue(targetOrd);
+                        return localRavv.isValueShared() ? Arrays.copyOf(v, v.length) : v;
+                    }
+                })
+                .collect(Collectors.toList());
+
+        // compute the centroid of the training set
+        float[] globalCentroid = KMeansPlusPlusClusterer.centroidOf(vectors);
+        return new BinaryQuantization(globalCentroid);
+    }
+
+    @Override
+    public CompressedVectors createCompressedVectors(Object[] compressedVectors) {
+        return new BQVectors(this, (long[][]) compressedVectors);
+    }
+
+    @Override
+    public long[][] encodeAll(List<float[]> vectors) {
+        return PhysicalCoreExecutor.instance.submit(() -> vectors.stream().parallel().map(this::encode).toArray(long[][]::new));
+    }
+
+    /**
+     * Encodes the input vector
+     *
+     * @return one bit per original f32
+     */
+    @Override
+    public long[] encode(float[] v) {
+        var centered = VectorUtil.sub(v, globalCentroid);
+
+        int M = (int) Math.ceil(centered.length / 64.0);
+        long[] encoded = new long[M];
+        for (int i = 0; i < M; i++) {
+            long bits = 0;
+            for (int j = 0; j < 64; j++) {
+                int idx = i * 64 + j;
+                if (idx >= centered.length) {
+                    break;
+                }
+                if (centered[idx] > 0) {
+                    bits |= 1L << j;
+                }
+            }
+            encoded[i] = bits;
+        }
+        return encoded;
+    }
+
+    @Override
+    public void write(DataOutput out) throws IOException {
+        out.writeInt(globalCentroid.length);
+        Io.writeFloats(out, globalCentroid);
+    }
+
+    public static BinaryQuantization load(RandomAccessReader in) throws IOException {
+        int length = in.readInt();
+        var centroid = new float[length];
+        in.readFully(centroid);
+        return new BinaryQuantization(centroid);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        BinaryQuantization that = (BinaryQuantization) o;
+        return Arrays.equals(globalCentroid, that.globalCentroid);
+    }
+
+    @Override
+    public int hashCode() {
+        return Arrays.hashCode(globalCentroid);
+    }
+
+    @Override
+    public String toString() {
+        return "BinaryQuantization";
+    }
+}