Skip to content

CNDB-16051: Include pq vectors in CompactionGraph bytesUsed calculation #2144

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michaeljmarshall merged 4 commits into from
Feb 10, 2026
Merged

CNDB-16051: Include pq vectors in CompactionGraph bytesUsed calculation #2144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
024e2a0
CNDB-16051: Include pq vectors in CompactionGraph bytesUsed calculation
michaeljmarshall Dec 1, 2025
405e511
Merge remote-tracking branch 'datastax/main' into cndb-16051
michaeljmarshall Feb 3, 2026
d7ea0d8
Remove assertion to keep code coverage without introducing a failing …
michaeljmarshall Feb 6, 2026
4840e99
Merge remote-tracking branch 'datastax/main' into cndb-16051
michaeljmarshall Feb 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/java/org/apache/cassandra/index/sai/disk/vector/CompactionGraph.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,6 +343,7 @@ public InsertionResult maybeAddVector(ByteBuffer term, int segmentRowId) throws
compressor = ((ProductQuantization) compressor).refine(new ListRandomAccessVectorValues(trainingVectors, dimension));
trainingVectors.clear(); // don't need these anymore so let GC reclaim if it wants to

long originalBytesUsed = compressedVectors.ramBytesUsed();
// re-encode the vectors added so far
int encodedVectorCount = compressedVectors.count();
compressedVectors = new MutablePQVectors((ProductQuantization) compressor);
Expand All @@ -358,6 +359,10 @@ public InsertionResult maybeAddVector(ByteBuffer term, int segmentRowId) throws
});
}).join();

// Update bytes to account for new encoding. This isn't expected to change, but just
// in case it does, we track it here.
bytesUsed += (compressedVectors.ramBytesUsed() - originalBytesUsed);

// Keep the existing edges but recompute their scores
builder = GraphIndexBuilder.rescore(builder, BuildScoreProvider.pqBuildScoreProvider(similarityFunction, (PQVectors) compressedVectors));
}
Expand All @@ -381,12 +386,15 @@ public InsertionResult maybeAddVector(ByteBuffer term, int segmentRowId) throws
for (int i = 0; i < dimension; i++)
vectorsByOrdinalBufferedWriter.writeFloat(vector.get(i));

// Track the bytes used as a result of this operation
long compressedVectorsBytesUsed = compressedVectors.ramBytesUsed();
// Fill in any holes in the pqVectors (setZero has the side effect of increasing the count)
while (compressedVectors.count() < ordinal)
compressedVectors.setZero(compressedVectors.count());
compressedVectors.encodeAndSet(ordinal, vector);

bytesUsed += postings.ramBytesUsed();
bytesUsed += (compressedVectors.ramBytesUsed() - compressedVectorsBytesUsed);
return new InsertionResult(bytesUsed, ordinal, vector);
}

Expand Down
15 changes: 15 additions & 0 deletions test/unit/org/apache/cassandra/index/sai/cql/VectorSiftSmallTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import org.apache.cassandra.index.sai.StorageAttachedIndex;
import org.apache.cassandra.index.sai.disk.v1.SegmentBuilder;
import org.apache.cassandra.index.sai.disk.v2.V2VectorIndexSearcher;
import org.apache.cassandra.index.sai.disk.vector.CompactionGraph;
import org.apache.cassandra.index.sai.disk.vector.NVQUtil;

import static org.junit.Assert.assertEquals;
Expand Down Expand Up @@ -221,6 +222,20 @@ public void testCompaction() throws Throwable
var recall = testRecall(topK, queryVectors, groundTruth);
assertTrue("Post-compaction recall is " + recall, recall > postCompactionRecall);
}

// Set force PQ training size to ensure we hit the refine code path and apply it to half the vectors.
// TODO this test fails as of this commit due to recall issues. Will investigate further.
CompactionGraph.PQ_TRAINING_SIZE = baseVectors.size() / 2;
Comment on lines +226 to +228
Copy link
Member Author

@michaeljmarshall michaeljmarshall Feb 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The problem here seems likely to be with GraphIndexBuilder.rescore in the CompactionGraph. The failure is unrelated to this PR. I am looking into possible fixes. Generally speaking, I think we can entirely remove the rescore logic, which duplicates work unnecessarily. We could instead accumulate N vectors, refine a PQ, encode the vectors, and then insert them into a graph builder properly built using these PQ values. I have a local version of this that passes, but is a bit hacky.

Copy link

@pkolaczk pkolaczk Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you explain how is this added test here related to the memory consumption tracking modification in the code under test? I can see we're likely invoking vector compression here, but we're not testing the memory consumption counters. Or am I missing sth?

Copy link
Member Author

@michaeljmarshall michaeljmarshall Feb 6, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By reducing the PQ_TRAINING_SIZE, we hit code that was previously uncovered and then we asserted on the recall for the graph built by that code and the test didn't pass. The failure is unrelated to my change.


// Compact again to take the CompactionGraph code path that calls the refine logic
compact();
for (int topK : List.of(1, 100))
{
var recall = testRecall(topK, queryVectors, groundTruth);
// This assertion will fail until we address the design the bug discussed
// in https://github.com/riptano/cndb/issues/16637.
// assertTrue("Post-compaction recall is " + recall, recall > postCompactionRecall);
}
}

// exercise the path where we use the PQ from the first segment (constructed on-heap)
Expand Down