leaves12138 commented on code in PR #62:
URL:
https://github.com/apache/paimon-vector-index/pull/62#discussion_r3651330052
##########
core/src/pq.rs:
##########
@@ -105,63 +182,85 @@ impl ProductQuantizer {
let m = self.m;
let d = self.d;
- let dsub = self.dsub;
let ksub = self.ksub;
+ let chunk_offsets = &self.chunk_offsets;
// Train all M sub-quantizers in parallel
let sub_results: Vec<Vec<f32>> = (0..m)
.into_par_iter()
.map(|sub| {
- let offset = sub * dsub;
+ let start = chunk_offsets[sub];
+ let stop = chunk_offsets[sub + 1];
+ let chunk_dim = stop - start;
- let mut sub_data = vec![0.0f32; n * dsub];
+ let mut sub_data = vec![0.0f32; n * chunk_dim];
for i in 0..n {
- sub_data[i * dsub..(i + 1) * dsub]
- .copy_from_slice(&data[i * d + offset..i * d + offset
+ dsub]);
+ sub_data[i * chunk_dim..(i + 1) * chunk_dim]
+ .copy_from_slice(&data[i * d + start..i * d + stop]);
}
let init: Option<Vec<f32>> = prev_centroids.as_ref().map(|pc| {
- let src = sub * ksub * dsub;
- pc[src..src + ksub * dsub].to_vec()
+ let src = start * ksub;
+ pc[src..src + ksub * chunk_dim].to_vec()
});
- kmeans::kmeans_train_with_init(km_config, &sub_data, n, dsub,
ksub, init.as_deref())
+ kmeans::kmeans_train_with_init(
+ km_config,
+ &sub_data,
+ n,
+ chunk_dim,
+ ksub,
+ init.as_deref(),
+ )
})
.collect();
- self.centroids = vec![0.0f32; m * ksub * dsub];
+ self.centroids = vec![0.0f32; d * ksub];
for (sub, sub_centroids) in sub_results.into_iter().enumerate() {
- let dst_offset = sub * ksub * dsub;
- self.centroids[dst_offset..dst_offset + ksub *
dsub].copy_from_slice(&sub_centroids);
+ let chunk_dim = self.chunk_dim(sub);
+ let dst_offset = self.centroid_chunk_base(sub);
+ self.centroids[dst_offset..dst_offset + ksub * chunk_dim]
+ .copy_from_slice(&sub_centroids);
}
self.rebuild_norms_cache();
}
/// Rebuild the centroid norms cache. Called after training or loading
centroids.
pub fn rebuild_norms_cache(&mut self) {
- self.centroid_norms_cache = vec![0.0f32; self.m * self.ksub];
+ self.try_rebuild_norms_cache()
+ .expect("PQ centroid norms allocation failed");
+ }
+
+ pub fn try_rebuild_norms_cache(&mut self) -> Result<(),
std::collections::TryReserveError> {
+ let mut norms = Vec::new();
+ norms.try_reserve_exact(self.m * self.ksub)?;
+ norms.resize(self.m * self.ksub, 0.0f32);
for sub in 0..self.m {
- let c_base = sub * self.ksub * self.dsub;
+ let chunk_dim = self.chunk_dim(sub);
+ let c_base = self.centroid_chunk_base(sub);
for j in 0..self.ksub {
- let c_off = c_base + j * self.dsub;
- self.centroid_norms_cache[sub * self.ksub + j] =
- fvec_norm_l2sqr(&self.centroids[c_off..c_off + self.dsub]);
+ let c_off = c_base + j * chunk_dim;
+ norms[sub * self.ksub + j] =
+ fvec_norm_l2sqr(&self.centroids[c_off..c_off + chunk_dim]);
}
}
+ self.centroid_norms_cache = norms;
+ Ok(())
}
/// Bytes per encoded vector.
pub fn code_size(&self) -> usize {
if self.nbits == 4 {
- self.m / 2
+ self.m.div_ceil(2)
Review Comment:
Allowing odd `m` for 4-bit PQ here makes IVF-PQ search silently ignore the
final subquantizer. Both `distance::scan_4bit_simd` and
`ivfpq::scan_codes_4bit_transposed` still use `m / 2` and iterate only complete
nibble pairs. I reproduced this with `m=3`:
`ProductQuantizer::distance_from_table` returns 100 and 1 for two codes whose
only difference is the third subquantizer, while `scan_codes_4bit` returns 0
for both and changes the ranking. DiskANN has dedicated packed-4-bit scanners
that handle odd `m`, but the generic IVF-PQ paths do not. Could we either
retain the even-`m` requirement for IVF-PQ 4-bit indexes or update both
row-major and transposed scanners to process the final low nibble?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
