Check values of clusterCount in PQ (#464)

Addresses issue #463 by checking the value of clusterCount in ProductQuantization. Since each PQ segment is encoded using a byte, using clusterCount = 256 is the optimal choice. With this PR, we now throw an exception if clusterCount > 256 and issue a warning if clusterCount < 256.

Author: marianotepper

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

@@ -36,6 +36,7 @@
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -53,6 +54,8 @@
 public class ProductQuantization implements VectorCompressor<ByteSequence<?>>, Accountable {
     private static final int MAGIC = 0x75EC4012; // JVECTOR, with some imagination
 
+    protected static final Logger LOG = Logger.getLogger(ProductQuantization.class.getName());
+
     private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
     static final int DEFAULT_CLUSTERS = 256; // number of clusters per subspace = one byte's worth
     static final int K_MEANS_ITERATIONS = 6;
@@ -114,6 +117,8 @@ public static ProductQuantization compute(RandomAccessVectorValues ravv,
                                               ForkJoinPool simdExecutor,
                                               ForkJoinPool parallelExecutor)
     {
+        checkClusterCount(clusterCount);
+
         var subvectorSizesAndOffsets = getSubvectorSizesAndOffsets(ravv.dimension(), M);
         var vectors = extractTrainingVectors(ravv, parallelExecutor);
 
@@ -190,6 +195,8 @@ public ProductQuantization refine(RandomAccessVectorValues ravv,
     }
 
     ProductQuantization(VectorFloat<?>[] codebooks, int clusterCount, int[][] subvectorSizesAndOffsets, VectorFloat<?> globalCentroid, float anisotropicThreshold) {
+        checkClusterCount(clusterCount);
+
         this.codebooks = codebooks;
         this.globalCentroid = globalCentroid;
         this.M = codebooks.length;
@@ -718,4 +725,13 @@ public String toString() {
                              anisotropicThreshold,
                              KMeansPlusPlusClusterer.computeParallelCostMultiplier(anisotropicThreshold, originalDimension));
     }
+
+    private static void checkClusterCount(int clusterCount) {
+        if (clusterCount > 256) {
+            throw new IllegalArgumentException("Too many PQ clusters: " + clusterCount + " > 256");
+        }
+        if (clusterCount < 256) {
+            LOG.warning("Using less than 256 PQ clusters will not reduce the memory footprint.");
+        }
+    }
 }