Skip to content

Commit

Permalink
fix math in MutablePQV allocation
Browse files Browse the repository at this point in the history
  • Loading branch information
jbellis committed Dec 23, 2024
1 parent 0092e8c commit e0f670c
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -41,15 +41,18 @@ public MutablePQVectors(ProductQuantization pq, int maximumVectorCount) {
long totalSize = (long) maximumVectorCount * compressedDimension;
this.vectorsPerChunk = totalSize <= MAX_CHUNK_SIZE ? maximumVectorCount : MAX_CHUNK_SIZE / compressedDimension;

int numChunks = maximumVectorCount / vectorsPerChunk;
ByteSequence<?>[] chunks = new ByteSequence<?>[numChunks];
int chunkSize = vectorsPerChunk * compressedDimension;
for (int i = 0; i < numChunks - 1; i++)
chunks[i] = vectorTypeSupport.createByteSequence(chunkSize);
int fullSizeChunks = maximumVectorCount / vectorsPerChunk;
int totalChunks = maximumVectorCount % vectorsPerChunk == 0 ? fullSizeChunks : fullSizeChunks + 1;
ByteSequence<?>[] chunks = new ByteSequence<?>[totalChunks];
int chunkBytes = vectorsPerChunk * compressedDimension;
for (int i = 0; i < fullSizeChunks; i++)
chunks[i] = vectorTypeSupport.createByteSequence(chunkBytes);

// Last chunk might be smaller
int remainingVectors = maximumVectorCount - (vectorsPerChunk * (numChunks - 1));
chunks[numChunks - 1] = vectorTypeSupport.createByteSequence(remainingVectors * compressedDimension);
if (totalChunks > fullSizeChunks) {
int remainingVectors = maximumVectorCount % vectorsPerChunk;
chunks[fullSizeChunks] = vectorTypeSupport.createByteSequence(remainingVectors * compressedDimension);
}

this.compressedDataChunks = chunks;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -250,4 +250,14 @@ public void testSaveVersion0() throws Exception {
var contents2 = Files.readAllBytes(fileOut.toPath());
assertArrayEquals(contents1, contents2);
}

@Test
public void testMutablePQVectors() {
// test that MPVQ gets the math right in an allocation edge case
var R = getRandom();
VectorFloat<?>[] vectors = generate(2 * DEFAULT_CLUSTERS, 2, 1_000);
var ravv = new ListRandomAccessVectorValues(List.of(vectors), vectors[0].length());
var pq = ProductQuantization.compute(ravv, 1, DEFAULT_CLUSTERS, false);
var pqv = new MutablePQVectors(pq, Integer.MAX_VALUE);
}
}

0 comments on commit e0f670c

Please sign in to comment.