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 breaks the public IVF-PQ path in two
ways. `write_index` now accepts and writes `m=3` with a two-byte code, but
`IVFPQIndexReader::open_with_header` still rejects that file with `4-bit PQ
requires even m`. Before serialization, both `distance::scan_4bit_simd` and
`ivfpq::scan_codes_4bit_transposed` also still use `m / 2` and iterate only
complete nibble pairs. I reproduced the latter 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 generic IVF-PQ does not. Could we either retain the
even-`m` requirement for IVF-PQ 4-bit indexes or update its writer, reader, and
both scan layouts consistently?
--
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]
