Make disk-search PQ pivot generation explicit - #1299
Conversation
|
@microsoft-github-policy-service agree company="Microsoft" |
There was a problem hiding this comment.
Pull request overview
This PR refactors disk-search PQ pivot generation so that pivot creation is an explicit step in the disk build flow, instead of being implicitly triggered inside PQGeneration::new(). The intent is to preserve behavior for fresh builds while making the flow clearer and ensuring PQGeneration::new() only loads existing pivots.
Changes:
- Added an explicit
PQGeneration::generate_pivots(...)step that generates pivots only when missing;PQGeneration::new()now validates and loads (and errors if pivots are absent). - Updated the disk build pipeline to call pivot generation before constructing the
QuantDataGenerator. - Refactored tests to cover explicit pivot generation and the new “missing pivots” error behavior.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| diskann-disk/src/storage/quant/pq/pq_generation.rs | Splits pivot generation from compressor construction; updates tests accordingly. |
| diskann-disk/src/build/builder/build.rs | Calls explicit pivot generation before creating the quantized data generator. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1299 +/- ##
==========================================
- Coverage 91.66% 90.68% -0.98%
==========================================
Files 514 516 +2
Lines 98732 99252 +520
==========================================
- Hits 90505 90010 -495
- Misses 8227 9242 +1015
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
diskann-disk/src/storage/quant/pq/pq_generation.rs:138
PQGeneration::new()loads the pivot metadata but does not verify that the pivot file's dimension matches the build context (context.dim). If a stale/incorrect pivots file exists for a different dimension, this can lead to confusing downstream failures or incorrect compression behavior. Consider validatingfull_dim == context.dimimmediately after reading pivot metadata and returning a clear error if it mismatches.
let (_, full_dim) = context
.pq_storage
.read_existing_pivot_metadata(context.storage_provider)?;
diskann-disk/src/storage/quant/pq/pq_generation.rs:133
- The new missing-pivots error message uses an inconsistent "ERROR:" prefix and prints the entire
PQStoragedebug struct with an extra trailing period. For consistency with nearby messages (e.g.,"Error: ...") and clearer logs, use a normal "Error:" prefix and avoid the trailing punctuation.
return Err(ANNError::log_pq_error(format_args!(
"ERROR: PQ k-means pivot file not found: {:?}.",
context.pq_storage
)));
|
Copilot resolve the merge conflicts in this pull request |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (1)
diskann-disk/src/storage/quant/pq/pq_generation.rs:147
ANNError::log_pq_errordoes not exist (andANNErrorisn’t imported), so this missing-pivots branch won’t compile. Construct the tagged error viadiskann_error!(ErrorKind::PQError, ...).
return Err(ANNError::log_pq_error(format_args!(
"ERROR: PQ k-means pivot file not found: {:?}.",
context.pq_storage
)));
}
The merge conflicts have been resolved locally. Here's a summary of what was done: Merge conflict:
Resolution: Kept our structural approach (
The resolved code builds successfully ( Push status: The fix commit ( |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
diskann-disk/src/storage/quant/pq/pq_generation.rs:131
- There are unresolved git merge conflict markers (<<<<<<< / ======= / >>>>>>>) in
PQGeneration::new, which will prevent this file from compiling. Resolve the conflict and keep a single validation call (the newvalidate_context).
fn new(context: &Self::CompressorContext) -> diskann::ANNResult<Self> {
Self::validate_context(context)?;
if !context
.pq_storage
diskann-disk/src/storage/quant/pq/pq_generation.rs:147
ANNError::log_pq_erroris not defined anywhere in the repo, andANNErrorisn’t in scope in this module. This will fail to compile. Use the existingdiskann_error!macro (tagged withErrorKind::PQError) to construct the error instead.
let (_, full_dim) = context
.pq_storage
.read_existing_pivot_metadata(context.storage_provider)?;
//Load the pivots
let num_chunks = context.num_chunks;
let (mut full_pivot_data, centroid, chunk_offsets) =
context.pq_storage.load_existing_pivot_data(
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
diskann-disk/src/storage/quant/pq/pq_generation.rs:137
ANNError::log_pq_erroris not defined anywhere in the workspace (andANNErrorisn’t in scope here), so this will fail to compile. Use the existingdiskann::ANNError::message(...)constructor (already used elsewhere for missing pivot errors) instead.
if !context
.pq_storage
.pivot_data_exist(context.storage_provider)
{
return Err(ANNError::log_pq_error(format_args!(
"ERROR: PQ k-means pivot file not found: {:?}.",
context.pq_storage
)));
diskann-disk/src/storage/quant/pq/pq_generation.rs:81
generate_pivotsreturns early on any existing pivots file, which means it will not regenerate when the file exists but its recorded metadata (dim / num_centers) doesn’t match the current context.generate_pq_pivotsalready has this metadata check; mirror it here before returning so stale pivots get regenerated rather than silently reused.
if context
.pq_storage
.pivot_data_exist(context.storage_provider)
{
return Ok(());
49547c0 to
e5aadc3
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
diskann-disk/src/storage/quant/pq/pq_generation.rs:66
- The validation error message is too generic; it doesn’t include the actual values for
num_chunksanddim, which makes failures harder to diagnose (especially now that context validation is shared by bothnew()andgenerate_pivots()). Consider including the values in the message.
if context.num_chunks > context.dim {
return Err(diskann_error!(
ErrorKind::PQError,
"Error: number of chunks more than dimension."
));
diskann-disk/src/storage/quant/pq/pq_generation.rs:138
- This new error path uses an inconsistent prefix ("ERROR:") and a trailing period. Consider using the same "Error:" prefix used elsewhere in the crate, and format the message to clearly indicate what value is being printed.
return Err(diskann_error!(
ErrorKind::PQError,
"ERROR: PQ k-means pivot file not found: {:?}.",
context.pq_storage
));
e5aadc3 to
6ddf44e
Compare
6ddf44e to
af9769d
Compare
|
Suggesting a different direction on this one The problem this PR is going after is real: PQGeneration::new() is nominally a constructor, but it silently runs k-means and writes pivots to disk. A method called new shouldn't do that. That said, hoisting the training to the call site trades one problem for another — it introduces an ordering contract that the compiler doesn't check. Callers of QuantDataGenerator must now call generate_pivots before new. That line returns (), produces no value anything downstream consumes, and still compiles if you delete it — it only fails at runtime. It also breaks up what used to be a clean new(...) → generate_data(...) flow in build.rs. The deeper point: new shouldn't do this work, but not because the caller should — because it should be deferred to generate_data. new builds an object from config; generate_data is what actually does the work. Concretely: Have QuantDataGenerator hold &Q::CompressorContext instead of a Q, so new is pure config and doesn't need to return Result This also fixes a real defect along the way: the max_block_size == 0 and empty-dataset checks live at the top of generate_data, but pivots have already been trained and loaded back in QuantDataGenerator::new by then. So a bad parameter costs you a full wasted k-means run before it reports. Deferring construction puts validation first and drops that cost. Two more points, independent of which approach you pick: The num_chunks > dim check and the pivot_data_exist call are now duplicated across generate_pivots and new — introduced by the refactor. Under the approach above, the copy in prepare can just go away. The old test_create_and_load_pivots_file had a step that called generate_pq_pivots independently and compared the result byte-for-byte against what the compressor wrote — that was verifying codebook contents are correct. The new test only asserts the file exists and the dimensions match, with nothing replacing that check. It's unrelated to this restructuring, so I'd restore it. |
af9769d to
2d9f1bb
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.
Suppressed comments (5)
diskann-disk/src/storage/quant/pq/pq_generation.rs:242
- The test name
explicit_generation_creates_pivots_fileis misleading: the test creates pivots by callingPQGeneration::prepare(&context), which implicitly generates missing pivots. If pivot generation is meant to be explicit, the test should either callPQGeneration::generate_pivots(...)directly or be renamed to reflect the actual behavior.
#[rstest]
fn explicit_generation_creates_pivots_file() {
let storage_provider = VirtualStorageProvider::new_memory();
diskann-disk/src/storage/quant/pq/pq_generation.rs:125
PQGeneration::prepare()still callsSelf::generate_pivots(context)?, which means pivot generation remains an implicit side effect of creating a compressor. This contradicts the PR description (“generation step is no longer hidden inside PQGeneration::new(); new only loads existing pivots”) and keeps the fallback logic from #1016 inside the compressor initialization path.
If the intent is to make the build flow explicit, consider moving pivot generation to the build orchestration (e.g. DiskIndexBuilder::generate_compressed_data) and making prepare() only load pivots (returning an error when the pivots file is missing).
This issue also appears on line 240 of the same file.
fn prepare(context: &Self::CompressorContext) -> diskann::ANNResult<Self> {
Self::generate_pivots(context)?;
diskann-disk/src/build/builder/build.rs:168
- The PR goal is to make disk-search PQ pivot generation explicit in the disk build flow, but the build currently only calls
QuantDataGenerator::generate_data(), which internally callsQ::prepare(...)(andPQGeneration::prepare()still triggers pivot generation).
To make the generation step explicit at the orchestration layer, call PQGeneration::generate_pivots(&quantizer_context)? here before constructing/running the generator (even if prepare() is later changed to only load).
let quantizer_context = PQGenerationContext {
pq_storage: self.pq_storage.clone(),
num_chunks: num_chunks.get(),
max_kmeans_reps: NUM_KMEANS_REPS_PQ,
num_centers: NUM_PQ_CENTROIDS,
seed: self.index_configuration.random_seed,
p_val: MAX_PQ_TRAINING_SET_SIZE / (num_points as f64),
storage_provider,
pool,
dim: self.index_configuration.dim,
metric: self.index_configuration.dist_metric,
};
let generator = QuantDataGenerator::<
Data::VectorDataType,
PQGeneration<Data::VectorDataType, StorageProvider>,
>::new(
self.index_writer.get_dataset_file(),
self.pq_storage.get_compressed_data_path().into(),
&quantizer_context,
);
generator.generate_data(
diskann-disk/src/storage/quant/compressor.rs:34
QuantCompressoris re-exported fromstorage::quant(a public module), so renaming the required method fromnewtoprepareis a breaking API change for downstream crates. If this is intended to be internal-only, consider making the traitpub(crate)(and adjusting re-exports). If it’s intended as public API, consider keepingnewas a default method (delegating toprepare) or otherwise documenting the breaking change / bumping the crate’s semver appropriately.
/// # Methods
/// - `prepare`: Performs any setup needed before compression and returns a compressor.
/// - `compress`: Compresses a batch of vectors into the output buffer.
/// - `compressed_bytes`: Returns the size in bytes of each compressed vector
pub trait QuantCompressor<T>: Sized + Sync
where
T: VectorRepr,
{
type CompressorContext;
fn prepare(context: &Self::CompressorContext) -> ANNResult<Self>;
fn compress(&self, vector: MatrixView<f32>, output: MutMatrixView<u8>) -> ANNResult<()>;
fn compressed_bytes(&self) -> usize;
diskann-disk/src/storage/quant/generator.rs:55
QuantDataGeneratoris publicly re-exported (storage::quant::QuantDataGenerator), and this change makes it borrow aCompressorContextvia a new lifetime parameter and makesnew()infallible. That’s a breaking API change for external callers (signature and type name change) and also changes when initialization failures surface (now ingenerate_data()viaQ::prepare).
If the intent is to keep this API stable, consider keeping the old constructor shape (e.g., keep new(...) -> ANNResult<Self> and/or add an additional constructor like with_context_ref(...)) to avoid forcing downstream lifetime plumbing.
pub struct QuantDataGenerator<'a, T, Q>
where
T: Copy + VectorRepr,
Q: QuantCompressor<T>,
{
quantizer_context: &'a Q::CompressorContext,
pub data_path: String,
pub compressed_data_path: String,
phantom: PhantomData<T>,
}
impl<'a, T, Q> QuantDataGenerator<'a, T, Q>
where
T: Copy + VectorRepr,
Q: QuantCompressor<T>,
{
pub fn new(
data_path: String,
compressed_data_path: String,
quantizer_context: &'a Q::CompressorContext,
) -> Self {
Self {
data_path,
compressed_data_path,
quantizer_context,
phantom: PhantomData,
}
}
Reference Issues/PRs
Refs #1016.
What does this implement/fix? Briefly explain your changes.
This refactors disk-search PQ pivot generation so it is explicit in the disk build flow. Fresh builds still generate
<prefix>_pq_pivots.bin, but the generation step is no longer hidden insidePQGeneration::new().PQGeneration::new()now only loads existing pivots.No intended behavior change.
Any other comments?
Testing:
cargo test -p diskann-diskcargo fmt --all --checkcargo clippy --workspace --all-targets -- -D warnings