Merge pull request #143 from jbellis/cv-dimensions

add getOriginalSize and getCompressedSize to CompressedVectors interface

Author: jkni

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

@@ -98,6 +98,16 @@ public long[] get(int i) {
         return compressedVectors[i];
     }
 
+    @Override
+    public int getOriginalSize() {
+        return bq.getOriginalDimension() * Float.BYTES;
+    }
+
+    @Override
+    public int getCompressedSize() {
+        return compressedVectors[0].length * Long.BYTES;
+    }
+
     @Override
     public long ramBytesUsed() {
         return compressedVectors.length * RamUsageEstimator.sizeOf(compressedVectors[0]);

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

@@ -107,6 +107,10 @@ public void write(DataOutput out) throws IOException {
         Io.writeFloats(out, globalCentroid);
     }
 
+    public int getOriginalDimension() {
+        return globalCentroid.length;
+    }
+
     public static BinaryQuantization load(RandomAccessReader in) throws IOException {
         int length = in.readInt();
         var centroid = new float[length];

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

@@ -24,7 +24,18 @@
 import java.io.IOException;
 
 public interface CompressedVectors extends Accountable {
+    /** write the compressed vectors to the given DataOutput */
     void write(DataOutput out) throws IOException;
 
+    /**
+     * @return a ScoreFunction suitable for performing search against the compressed vectors,
+     * potentially without decompression them first
+     */
     NeighborSimilarity.ApproximateScoreFunction approximateScoreFunctionFor(float[] q, VectorSimilarityFunction similarityFunction);
+
+    /** @return the original size of the vectors, in bytes, before compression */
+    int getOriginalSize();
+
+    /** @return the compressed size of the vectors, in bytes */
+    int getCompressedSize();
 }

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

@@ -125,6 +125,16 @@ float[] reusablePartialMagnitudes() {
         return partialMagnitudes.get();
     }
 
+    @Override
+    public int getOriginalSize() {
+        return pq.originalDimension * Float.BYTES;
+    }
+
+    @Override
+    public int getCompressedSize() {
+        return pq.codebooks.length;
+    }
+
     @Override
     public long ramBytesUsed() {
         long codebooksSize = pq.memorySize();

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

@@ -48,8 +48,8 @@ public class ProductQuantization implements VectorCompressor<byte[]> {
     static final int MAX_PQ_TRAINING_SET_SIZE = 128000;
 
     final float[][][] codebooks;
-    final int M;
-    private final int originalDimension;
+    private final int M; // codebooks.length, redundantly reproduced for convenience
+    final int originalDimension;
     private final float[] globalCentroid;
     final int[][] subvectorSizesAndOffsets;
 

jvector-tests/src/test/java/io/github/jbellis/jvector/pq/TestCompressedVectors.java

@@ -46,6 +46,8 @@ public void testSaveLoadPQ() throws Exception {
         // Compress the vectors
         var compressed = pq.encodeAll(vectors);
         var cv = new PQVectors(pq, compressed);
+        assertEquals(2 * Float.BYTES, cv.getOriginalSize());
+        assertEquals(1, cv.getCompressedSize());
 
         // Write compressed vectors
         File cvFile = File.createTempFile("pqtest", ".cv");
@@ -68,6 +70,8 @@ public void testSaveLoadBQ() throws Exception {
         // Compress the vectors
         var compressed = bq.encodeAll(vectors);
         var cv = new BQVectors(bq, compressed);
+        assertEquals(64 * Float.BYTES, cv.getOriginalSize());
+        assertEquals(8, cv.getCompressedSize());
 
         // Write compressed vectors
         File cvFile = File.createTempFile("bqtest", ".cv");