This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
The following commit(s) were added to refs/heads/main by this push:
new 9464181 perf(pq): add transposed SIMD encoder with SGEMM fallback
(#90)
9464181 is described below
commit 9464181434ab5d8bf73301c8b2d2924d7db8db16
Author: jerry <[email protected]>
AuthorDate: Wed Sep 9 11:19:40 2026 +0800
perf(pq): add transposed SIMD encoder with SGEMM fallback (#90)
---
.github/workflows/ci.yml | 13 +
core/benches/ann_bench.rs | 1 +
core/benches/ivfpq_add_bench.rs | 17 +-
core/benches/ivfpq_batch_reuse_bench.rs | 8 +-
core/benches/ivfpq_filter_scan_bench.rs | 8 +-
core/benches/ivfpq_train_bench.rs | 2 +-
core/src/diskann.rs | 34 +-
core/src/diskann_io.rs | 40 +-
core/src/diskann_search.rs | 8 +-
core/src/index.rs | 103 +-
core/src/io.rs | 41 +-
core/src/ivfpq.rs | 209 ++--
core/src/pq.rs | 1024 +++++++++++++++++---
core/src/vamana.rs | 60 +-
core/tests/storage_format_fixtures.rs | 50 +-
docs/api.html | 1 +
docs/ivf-pq.html | 4 +-
docs/releases.html | 3 +-
.../index/vector/VectorIndexJavaApiTest.java | 1 +
.../vector/VectorIndexNativeValidationTest.java | 9 +-
python/tests/test_vindex.py | 12 +
21 files changed, 1309 insertions(+), 339 deletions(-)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index bc4533d..c206376 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -288,6 +288,19 @@ jobs:
- name: Build FFI library
run: cargo build --release -p paimon-vindex-ffi
+ - name: Test AArch64 PQ kernels
+ if: runner.os == 'macOS'
+ run: cargo test --release -p paimon-vindex-core pq::tests
+
+ - name: Test no-AVX2/FMA PQ fallback under Rosetta
+ if: runner.os == 'macOS'
+ env:
+ PAIMON_EXPECT_PQ_SGEMM_FALLBACK: '1'
+ run: |
+ rustup target add x86_64-apple-darwin
+ cargo test --release -p paimon-vindex-core --target
x86_64-apple-darwin config_from_options_pq_encoding
+ cargo test --release -p paimon-vindex-core --target
x86_64-apple-darwin test_encode_batch_8bit_sgemm_is_thread_and_split_invariant
+
- name: Build wheel
working-directory: python
run: |
diff --git a/core/benches/ann_bench.rs b/core/benches/ann_bench.rs
index a62566b..0d79609 100644
--- a/core/benches/ann_bench.rs
+++ b/core/benches/ann_bench.rs
@@ -667,6 +667,7 @@ fn index_specs(config: &Config) -> Vec<IndexSpec> {
metric: MetricType::L2,
use_opq: false,
use_approximate_coarse_assignment: true,
+ canonical_pq_encoding: false,
},
searches: vec![ivf_search],
},
diff --git a/core/benches/ivfpq_add_bench.rs b/core/benches/ivfpq_add_bench.rs
index 3893cb8..1cfffff 100644
--- a/core/benches/ivfpq_add_bench.rs
+++ b/core/benches/ivfpq_add_bench.rs
@@ -106,16 +106,10 @@ const CASES: [Case; 10] = [
},
];
-fn new_index(
- case: Case,
- quantizer_centroids: &[f32],
- centroids: &[f32],
- norms: &[f32],
-) -> IVFPQIndex {
+fn new_index(case: Case, quantizer_centroids: &[f32], centroids: &[f32]) ->
IVFPQIndex {
let mut index = IVFPQIndex::new(case.d, case.nlist, case.m,
MetricType::L2, false);
index.set_quantizer_centroids(quantizer_centroids.to_vec());
- index.pq.centroids = centroids.to_vec();
- index.pq.centroid_norms_cache = norms.to_vec();
+ index.pq.set_centroids(centroids.to_vec());
index
}
@@ -138,11 +132,6 @@ fn bench_ivfpq_add(c: &mut Criterion) {
let centroids = (0..case.m * 256 * dsub)
.map(|_| rng.gen_range(-1.0f32..1.0))
.collect::<Vec<_>>();
- let norms = centroids
- .chunks_exact(dsub)
- .map(|centroid| centroid.iter().map(|value| value * value).sum())
- .collect::<Vec<_>>();
-
group.throughput(Throughput::Elements(case.rows as u64));
group.bench_with_input(
BenchmarkId::new(
@@ -155,7 +144,7 @@ fn bench_ivfpq_add(c: &mut Criterion) {
&case,
|b, &case| {
b.iter_batched(
- || new_index(case, &quantizer_centroids, ¢roids,
&norms),
+ || new_index(case, &quantizer_centroids, ¢roids),
|mut index| index.add(black_box(&data), black_box(&ids),
case.rows),
BatchSize::LargeInput,
);
diff --git a/core/benches/ivfpq_batch_reuse_bench.rs
b/core/benches/ivfpq_batch_reuse_bench.rs
index 7a67db8..4713e42 100644
--- a/core/benches/ivfpq_batch_reuse_bench.rs
+++ b/core/benches/ivfpq_batch_reuse_bench.rs
@@ -103,9 +103,11 @@ fn main() {
.map(|_| rng.gen_range(-1.0f32..1.0))
.collect(),
);
- index.pq.centroids = (0..M * index.pq.ksub * index.pq.dsub)
- .map(|_| rng.gen_range(-1.0f32..1.0))
- .collect();
+ index.pq.set_centroids(
+ (0..M * index.pq.ksub() * index.pq.dsub())
+ .map(|_| rng.gen_range(-1.0f32..1.0))
+ .collect(),
+ );
for list_id in 0..NLIST {
let first_id = list_id * ROWS_PER_LIST;
index.ids[list_id] = (first_id..first_id + ROWS_PER_LIST)
diff --git a/core/benches/ivfpq_filter_scan_bench.rs
b/core/benches/ivfpq_filter_scan_bench.rs
index a886c51..9f8b48b 100644
--- a/core/benches/ivfpq_filter_scan_bench.rs
+++ b/core/benches/ivfpq_filter_scan_bench.rs
@@ -104,9 +104,11 @@ fn main() {
.map(|_| rng.gen_range(-1.0f32..1.0))
.collect(),
);
- index.pq.centroids = (0..M * index.pq.ksub * index.pq.dsub)
- .map(|_| rng.gen_range(-1.0f32..1.0))
- .collect();
+ index.pq.set_centroids(
+ (0..M * index.pq.ksub() * index.pq.dsub())
+ .map(|_| rng.gen_range(-1.0f32..1.0))
+ .collect(),
+ );
for list_id in 0..NLIST {
let first_id = list_id * ROWS_PER_LIST;
index.ids[list_id] = (first_id..first_id + ROWS_PER_LIST)
diff --git a/core/benches/ivfpq_train_bench.rs
b/core/benches/ivfpq_train_bench.rs
index 07ce98b..34a79e1 100644
--- a/core/benches/ivfpq_train_bench.rs
+++ b/core/benches/ivfpq_train_bench.rs
@@ -103,7 +103,7 @@ fn run_scenario(s: &Scenario) {
// Keep results observable so nothing is optimized away.
let checksum: f32 =
- centroids.iter().take(8).sum::<f32>() +
pq.centroids.iter().take(8).sum::<f32>();
+ centroids.iter().take(8).sum::<f32>() +
pq.centroids().iter().take(8).sum::<f32>();
println!(
"{:<11} {:>8} {:>5} {:>6} {:>5} {:>8.3} {:>9.3} {:>8.3} {:>14.3}",
diff --git a/core/src/diskann.rs b/core/src/diskann.rs
index 4b2cb95..30ae782 100644
--- a/core/src/diskann.rs
+++ b/core/src/diskann.rs
@@ -317,8 +317,8 @@ impl DiskAnnIndex {
fn training_plan(&self, n: usize) -> io::Result<PqTrainingPlan> {
pq_training_plan_with_sample_buffers(
self.d,
- self.pq.m,
- self.pq.ksub,
+ self.pq.m(),
+ self.pq.ksub(),
n,
self.build_params.memory_budget_bytes,
usize::from(self.metric == MetricType::Cosine) + 1,
@@ -338,14 +338,15 @@ impl DiskAnnIndex {
let row_ids = checked_bytes(n, size_of::<i64>(), "row IDs")?;
let row_id_encoding_scratch = row_id_encoding_scratch_bytes(n)?;
let pq_codes = checked_bytes(n, self.pq.code_size(), "PQ codes")?;
- let pq_codebook = checked_bytes(self.pq.centroids.len(),
size_of::<f32>(), "PQ codebook")?;
+ let pq_codebook =
+ checked_bytes(self.pq.centroids().len(), size_of::<f32>(), "PQ
codebook")?;
let pq_build_distances = if self.build_params.build_distance
== DiskAnnBuildDistance::ProductQuantized
{
self.pq
- .m
- .checked_mul(self.pq.ksub)
- .and_then(|value| value.checked_mul(self.pq.ksub))
+ .m()
+ .checked_mul(self.pq.ksub())
+ .and_then(|value| value.checked_mul(self.pq.ksub()))
.and_then(|value| value.checked_mul(size_of::<f32>()))
.ok_or_else(|| invalid_input("DiskANN PQ build-distance table
size overflows"))?
} else {
@@ -438,12 +439,17 @@ impl DiskAnnIndex {
}
pub(crate) fn validate_for_write(&self) -> io::Result<()> {
- validate_diskann_format_configuration(self.d, self.pq.m,
self.pq.nbits, self.build_params)?;
+ validate_diskann_format_configuration(
+ self.d,
+ self.pq.m(),
+ self.pq.nbits(),
+ self.build_params,
+ )?;
validate_diskann_training_budget(
self.d,
self.metric,
- self.pq.m,
- self.pq.nbits,
+ self.pq.m(),
+ self.pq.nbits(),
self.build_params.memory_budget_bytes,
)?;
if self.build_params.memory_budget_bytes == 0 {
@@ -489,22 +495,22 @@ impl DiskAnnIndex {
}
let expected_ksub = 1usize
- .checked_shl(self.pq.nbits as u32)
+ .checked_shl(self.pq.nbits() as u32)
.ok_or_else(|| invalid_input("DiskANN PQ centroid count overflows
usize"))?;
let expected_centroids = self
.d
.checked_mul(expected_ksub)
.ok_or_else(|| invalid_input("DiskANN PQ codebook shape overflows
usize"))?;
- if self.pq.d != self.d
- || self.pq.ksub != expected_ksub
- || self.pq.centroids.len() != expected_centroids
+ if self.pq.d() != self.d
+ || self.pq.ksub() != expected_ksub
+ || self.pq.centroids().len() != expected_centroids
|| !self.pq.has_valid_layout()
{
return Err(invalid_input("DiskANN PQ codebook shape is invalid"));
}
if let Some(offset) = self
.pq
- .centroids
+ .centroids()
.iter()
.position(|value| !value.is_finite())
{
diff --git a/core/src/diskann_io.rs b/core/src/diskann_io.rs
index e2106cf..b22a2e8 100644
--- a/core/src/diskann_io.rs
+++ b/core/src/diskann_io.rs
@@ -1433,10 +1433,8 @@ impl<R: SeekRead> DiskAnnIndexReader<R> {
)));
}
- let (mut pq, row_ids, pq_codes, adjacency_index) =
+ let (pq, row_ids, pq_codes, adjacency_index) =
read_resident_sections(&mut self.reader, &self.header)?;
- pq.try_rebuild_norms_cache()
- .map_err(|_| invalid_data("DiskANN PQ norms allocation failed"))?;
validate_pq_code_padding(&self.header, &pq_codes)?;
let adjacency_validation =
AdjacencyValidationCache::new(adjacency_page_count(&self.header)?)?;
@@ -2312,8 +2310,8 @@ pub fn write_diskann_index_with_stats(
index.d,
index.ids.len(),
prepared.graph.entry_node,
- index.pq.m,
- index.pq.nbits,
+ index.pq.m(),
+ index.pq.nbits(),
index.metric,
index.build_params,
row_ids_len,
@@ -2588,38 +2586,38 @@ fn write_pq_codebook(
put_u32(
&mut header,
8,
- u32::try_from(pq.d).map_err(|_| invalid_input("DiskANN PQ dimension
exceeds u32"))?,
+ u32::try_from(pq.d()).map_err(|_| invalid_input("DiskANN PQ dimension
exceeds u32"))?,
);
put_u32(
&mut header,
12,
- u32::try_from(pq.m).map_err(|_| invalid_input("DiskANN PQ m exceeds
u32"))?,
+ u32::try_from(pq.m()).map_err(|_| invalid_input("DiskANN PQ m exceeds
u32"))?,
);
put_u32(
&mut header,
16,
- u32::try_from(pq.nbits).map_err(|_| invalid_input("DiskANN PQ bits
exceeds u32"))?,
+ u32::try_from(pq.nbits()).map_err(|_| invalid_input("DiskANN PQ bits
exceeds u32"))?,
);
put_u32(
&mut header,
20,
- u32::try_from(pq.ksub).map_err(|_| invalid_input("DiskANN PQ ksub
exceeds u32"))?,
+ u32::try_from(pq.ksub()).map_err(|_| invalid_input("DiskANN PQ ksub
exceeds u32"))?,
);
put_u32(
&mut header,
24,
- u32::try_from(pq.chunk_offsets.len())
+ u32::try_from(pq.chunk_offsets().len())
.map_err(|_| invalid_input("DiskANN PQ chunk-offset count exceeds
u32"))?,
);
writer.write_bytes(&header)?;
- for &offset in &pq.chunk_offsets {
+ for &offset in pq.chunk_offsets() {
writer.write_bytes(
&u32::try_from(offset)
.map_err(|_| invalid_input("DiskANN PQ chunk offset exceeds
u32"))?
.to_le_bytes(),
)?;
}
- for &value in &pq.centroids {
+ for &value in pq.centroids() {
writer.write_bytes(&value.to_le_bytes())?;
}
Ok(())
@@ -3050,7 +3048,7 @@ fn decode_pq_codebook(bytes: &[u8], header:
&DiskAnnHeader) -> io::Result<Produc
.map_err(invalid_data)?;
let mut centroids = Vec::new();
centroids
- .try_reserve_exact(header.dimension as usize * pq.ksub)
+ .try_reserve_exact(header.dimension as usize * pq.ksub())
.map_err(|_| invalid_data("DiskANN PQ centroid allocation failed"))?;
for encoded in bytes[centroids_offset..].chunks_exact(4) {
let value =
@@ -3060,7 +3058,8 @@ fn decode_pq_codebook(bytes: &[u8], header:
&DiskAnnHeader) -> io::Result<Produc
}
centroids.push(value);
}
- pq.centroids = centroids;
+ pq.try_set_centroids(centroids)
+ .map_err(|_| invalid_data("DiskANN PQ norms allocation failed"))?;
if !pq.has_valid_layout() {
return Err(invalid_data("invalid DiskANN PQ codebook layout"));
}
@@ -4120,15 +4119,18 @@ mod tests {
..DiskAnnBuildParams::default()
},
);
- index.pq.centroids = (0..256).map(|code| code as f32).collect();
- index.pq.rebuild_norms_cache();
+ index
+ .pq
+ .set_centroids((0..256).map(|code| code as f32).collect());
index.ids = vec![7];
index.vectors = vec![0.0];
index
}
let mut invalid_codebook = one_vector_index();
- invalid_codebook.pq.centroids[0] = f32::NAN;
+ let mut centroids = invalid_codebook.pq.centroids().to_vec();
+ centroids[0] = f32::NAN;
+ invalid_codebook.pq.set_centroids(centroids);
let mut codebook_output = Vec::new();
assert!(
write_diskann_index(&invalid_codebook, &mut PosWriter::new(&mut
codebook_output))
@@ -4154,7 +4156,7 @@ mod tests {
assert!(f16_output.is_empty());
let mut invalid_pq_shape = one_vector_index();
- invalid_pq_shape.pq.chunk_offsets[1] = 0;
+ invalid_pq_shape.pq = ProductQuantizer::new(invalid_pq_shape.d + 1, 1);
let mut pq_shape_output = Vec::new();
assert!(
write_diskann_index(&invalid_pq_shape, &mut PosWriter::new(&mut
pq_shape_output))
@@ -4840,7 +4842,7 @@ mod tests {
);
assert_eq!(reader.row_id_count().unwrap(), indexed_count);
assert_eq!(reader.pq_codes().unwrap().len(), indexed_count * 2);
- assert_eq!(reader.pq().unwrap().centroids, index.pq.centroids);
+ assert_eq!(reader.pq().unwrap().centroids(), index.pq.centroids());
let limited_rounds = Arc::new(Mutex::new(Vec::new()));
let limited_recording = RoundRecordingReader {
diff --git a/core/src/diskann_search.rs b/core/src/diskann_search.rs
index e1a557b..959e8f0 100644
--- a/core/src/diskann_search.rs
+++ b/core/src/diskann_search.rs
@@ -1460,7 +1460,7 @@ impl<R: SeekRead>
crate::diskann_io::DiskAnnIndexReader<R> {
let query_count = queries.len() / dimension;
let pq_m = self.header.pq_m as usize;
let pq = self.pq()?;
- let pq_ksub = pq.ksub;
+ let pq_ksub = pq.ksub();
let pq_code_size = pq.code_size();
let pq_codes = self.pq_codes()?;
let metric = self.header.metric_type();
@@ -1881,7 +1881,7 @@ impl<R: SeekRead>
crate::diskann_io::DiskAnnIndexReader<R> {
self.header.max_degree as usize,
)?;
let pq = self.pq()?;
- let distance_table_len = pq.m * pq.ksub;
+ let distance_table_len = pq.m() * pq.ksub();
pq.compute_distance_table(
query,
self.header.metric_type(),
@@ -2024,7 +2024,7 @@ impl<R: SeekRead>
crate::diskann_io::DiskAnnIndexReader<R> {
let result = (|| {
scratch.begin_rerank();
let pq = self.pq()?;
- let distance_table_len = pq.m * pq.ksub;
+ let distance_table_len = pq.m() * pq.ksub();
pq.compute_distance_table(
query,
self.header.metric_type(),
@@ -5432,7 +5432,7 @@ mod tests {
reader.ensure_resident().unwrap();
assert_eq!(reader.header.pq_bits, 4);
- assert_eq!(reader.pq().unwrap().ksub, 16);
+ assert_eq!(reader.pq().unwrap().ksub(), 16);
assert_eq!(reader.pq_codes().unwrap().len(), indexed_count);
let (result_ids, distances) = reader.search(query, 5, 100).unwrap();
diff --git a/core/src/index.rs b/core/src/index.rs
index 773984f..9acdee5 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -155,6 +155,9 @@ pub enum VectorIndexConfig {
metric: MetricType,
use_opq: bool,
use_approximate_coarse_assignment: bool,
+ /// Use canonical expanded-form PQ encoding instead of the default
+ /// transposed direct-L2 encoder.
+ canonical_pq_encoding: bool,
},
IvfRq {
dimension: usize,
@@ -207,6 +210,7 @@ impl VectorIndexConfig {
metric,
use_opq,
use_approximate_coarse_assignment: true,
+ canonical_pq_encoding: false,
};
validate_config(&config)?;
Ok(config)
@@ -283,6 +287,8 @@ pub struct ResolvedVectorIndexConfig {
pub rq_bits: Option<usize>,
pub use_opq: bool,
pub use_approximate_coarse_assignment: bool,
+ /// Only meaningful for IVF-PQ.
+ pub canonical_pq_encoding: bool,
pub diskann_build: Option<DiskAnnBuildParams>,
}
@@ -310,6 +316,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
rq_bits: None,
use_opq: false,
use_approximate_coarse_assignment:
*use_approximate_coarse_assignment,
+ canonical_pq_encoding: false,
diskann_build: None,
},
VectorIndexConfig::IvfPq {
@@ -319,6 +326,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
metric,
use_opq,
use_approximate_coarse_assignment,
+ canonical_pq_encoding,
} => Self {
index_type: IndexType::IvfPq,
dimension: *dimension,
@@ -329,6 +337,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
rq_bits: None,
use_opq: *use_opq,
use_approximate_coarse_assignment:
*use_approximate_coarse_assignment,
+ canonical_pq_encoding: *canonical_pq_encoding,
diskann_build: None,
},
VectorIndexConfig::IvfRq {
@@ -347,6 +356,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
rq_bits: Some(*bits),
use_opq: false,
use_approximate_coarse_assignment:
*use_approximate_coarse_assignment,
+ canonical_pq_encoding: false,
diskann_build: None,
},
VectorIndexConfig::DiskAnn {
@@ -365,6 +375,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
rq_bits: None,
use_opq: false,
use_approximate_coarse_assignment: false,
+ canonical_pq_encoding: false,
diskann_build: Some(*build),
},
}
@@ -431,6 +442,7 @@ impl VectorIndexBuildPlan {
}
IndexType::DiskAnn => true,
};
+ let canonical_pq_encoding = parse_ivf_pq_encoding_option(&mut options,
index_type)?;
let config = match index_type {
IndexType::IvfFlat => VectorIndexConfig::IvfFlat {
@@ -460,6 +472,7 @@ impl VectorIndexBuildPlan {
None => target_recall.is_some_and(|recall| recall >= 0.9),
},
use_approximate_coarse_assignment,
+ canonical_pq_encoding,
},
IndexType::IvfRq => {
let explicit_bits = options
@@ -639,6 +652,28 @@ fn parse_ivf_coarse_assignment_option(options: &mut
ConfigOptions) -> io::Result
}
}
+fn parse_ivf_pq_encoding_option(
+ options: &mut ConfigOptions,
+ index_type: IndexType,
+) -> io::Result<bool> {
+ let value = options.optional("ivf.pq-encoding");
+ if index_type != IndexType::IvfPq {
+ return match value {
+ None => Ok(false),
+ Some(_) => Err(invalid_input(
+ "option 'ivf.pq-encoding' is only valid for IVF-PQ",
+ )),
+ };
+ }
+ match value.as_deref().map(str::trim) {
+ None | Some("auto") => Ok(false),
+ Some("canonical") => Ok(true),
+ Some(_) => Err(invalid_input(
+ "option 'ivf.pq-encoding' must be auto or canonical",
+ )),
+ }
+}
+
fn parse_deployment_profile_option(name: &str, value: &str) ->
io::Result<DeploymentProfile> {
match value.trim() {
"auto" => Ok(DeploymentProfile::Auto),
@@ -1339,9 +1374,11 @@ impl VectorIndexWriter {
metric,
use_opq,
use_approximate_coarse_assignment,
+ canonical_pq_encoding,
} => {
let mut index = IVFPQIndex::new(dimension, nlist, m, metric,
use_opq);
index.set_approximate_coarse_assignment(use_approximate_coarse_assignment);
+ index.set_canonical_pq_encoding(canonical_pq_encoding);
Self::IvfPq(index)
}
VectorIndexConfig::IvfRq {
@@ -1524,7 +1561,7 @@ impl<R: SeekRead> VectorIndexReader<R> {
metric: reader.metric,
total_vectors: reader.total_vectors,
pq_m: Some(reader.m),
- pq_bits: Some(reader.pq.nbits),
+ pq_bits: Some(reader.pq.nbits()),
rq_bits: None,
diskann: None,
},
@@ -2946,6 +2983,7 @@ mod tests {
metric: MetricType::L2,
use_opq: false,
use_approximate_coarse_assignment: true,
+ canonical_pq_encoding: false,
},
VectorIndexConfig::IvfRq {
dimension: 8,
@@ -3046,6 +3084,7 @@ mod tests {
metric: MetricType::L2,
use_opq: false,
use_approximate_coarse_assignment: true,
+ canonical_pq_encoding: false,
}) {
Ok(_) => panic!("invalid PQ config should be rejected"),
Err(err) => err,
@@ -4295,6 +4334,68 @@ mod tests {
assert!(error.to_string().contains("must be auto or exact"));
}
+ #[test]
+ fn config_from_options_pq_encoding() {
+ let supports_transposed = crate::pq::supports_transposed_pq_encoding();
+ if std::env::var("PAIMON_EXPECT_PQ_SGEMM_FALLBACK").as_deref() ==
Ok("1") {
+ assert!(!supports_transposed, "expected the SGEMM fallback
backend");
+ }
+
+ let base = [
+ ("index.type", "ivf_pq"),
+ ("dimension", "4"),
+ ("nlist", "1"),
+ ("pq.m", "1"),
+ ("metric", "l2"),
+ ];
+ let with = |extra: &[(&str, &str)]| {
+ let mut all = base.to_vec();
+ all.extend_from_slice(extra);
+ VectorIndexConfig::from_options(&options(&all))
+ };
+
+ let auto = with(&[]).unwrap();
+ assert!(!auto.resolved().canonical_pq_encoding);
+ assert!(
+ !with(&[("ivf.pq-encoding", "auto")])
+ .unwrap()
+ .resolved()
+ .canonical_pq_encoding
+ );
+ let canonical = with(&[("ivf.pq-encoding", "canonical")]).unwrap();
+ assert!(canonical.resolved().canonical_pq_encoding);
+
+ let encoded_code = |config| {
+ let VectorIndexWriter::IvfPq(mut index) =
+ VectorIndexWriter::from_config(config).unwrap()
+ else {
+ unreachable!()
+ };
+ index.set_quantizer_centroids(vec![0.0; 4]);
+ let mut centroids = vec![100_000_016.0; 4 * 256];
+ centroids[0..4].fill(100_000_008.0);
+ centroids[4..8].fill(100_000_000.0);
+ index.pq.set_centroids(centroids);
+ index.add(&[100_000_000.0; 4], &[0], 1);
+ index.codes[0][0]
+ };
+ assert_eq!(encoded_code(auto), u8::from(supports_transposed));
+ assert_eq!(encoded_code(canonical), 0);
+
+ let error = with(&[("ivf.pq-encoding", "sgemm")]).unwrap_err();
+ assert!(error.to_string().contains("must be auto or canonical"));
+
+ let error = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "ivf_flat"),
+ ("dimension", "8"),
+ ("nlist", "4"),
+ ("metric", "l2"),
+ ("ivf.pq-encoding", "canonical"),
+ ]))
+ .unwrap_err();
+ assert!(error.to_string().contains("only valid for IVF-PQ"));
+ }
+
#[test]
fn config_from_options_rejects_unknown_options() {
let err = VectorIndexConfig::from_options(&options(&[
diff --git a/core/src/io.rs b/core/src/io.rs
index 3a55823..3067be1 100644
--- a/core/src/io.rs
+++ b/core/src/io.rs
@@ -280,9 +280,9 @@ fn checked_list_bytes(count: usize, bytes_per_entry: usize)
-> io::Result<usize>
pub fn write_index(index: &IVFPQIndex, out: &mut dyn SeekWrite) ->
io::Result<()> {
let d = index.d;
let nlist = index.nlist;
- let m = index.pq.m;
- let ksub = index.pq.ksub;
- let dsub = index.pq.dsub;
+ let m = index.pq.m();
+ let ksub = index.pq.ksub();
+ let dsub = index.pq.dsub();
let code_size = index.pq.code_size();
if ksub == 16 && !m.is_multiple_of(2) {
return Err(io::Error::new(
@@ -364,7 +364,7 @@ pub fn write_index(index: &IVFPQIndex, out: &mut dyn
SeekWrite) -> io::Result<()
}
write_f32_slice(out, index.quantizer_centroids())?;
- write_f32_slice(out, &index.pq.centroids)?;
+ write_f32_slice(out, index.pq.centroids())?;
// Compute offsets for inverted lists
// Delta-varint format per list: [base_id: i64][id_bytes_len:
u32][id_bytes][codes]
@@ -640,16 +640,7 @@ impl<R: SeekRead> IVFPQIndexReader<R> {
total_vectors,
opq: None,
quantizer_centroids: Vec::new(),
- pq: ProductQuantizer {
- d,
- m,
- nbits: ksub.trailing_zeros() as usize,
- dsub,
- ksub,
- chunk_offsets: (0..=m).map(|chunk| chunk * dsub).collect(),
- centroids: Vec::new(),
- centroid_norms_cache: Vec::new(),
- },
+ pq: ProductQuantizer::with_nbits(d, m, ksub.trailing_zeros() as
usize),
list_offsets: Vec::new(),
list_counts: Vec::new(),
list_id_bytes_lens: Vec::new(),
@@ -743,17 +734,7 @@ impl<R: SeekRead> IVFPQIndexReader<R> {
let pq_centroids = bytes_to_f32_vec(&metadata[position..position +
pq_centroid_bytes])?;
position += pq_centroid_bytes;
- self.pq = ProductQuantizer {
- d,
- m,
- nbits: ksub.trailing_zeros() as usize,
- dsub,
- ksub,
- chunk_offsets: (0..=m).map(|chunk| chunk * dsub).collect(),
- centroids: pq_centroids,
- centroid_norms_cache: Vec::new(),
- };
- self.pq.rebuild_norms_cache();
+ self.pq.set_centroids(pq_centroids);
self.list_offsets = vec![0i64; nlist];
self.list_counts = vec![0i32; nlist];
@@ -1226,9 +1207,9 @@ fn compute_precomputed_table(
nlist: usize,
d: usize,
) -> Vec<f32> {
- let m = pq.m;
- let ksub = pq.ksub;
- let dsub = pq.dsub;
+ let m = pq.m();
+ let ksub = pq.ksub();
+ let dsub = pq.dsub();
let table_size = nlist * m * ksub;
let mut table = vec![0.0f32; table_size];
@@ -1247,7 +1228,7 @@ fn compute_precomputed_table(
let pq_off = pq_base + j * dsub;
let mut ip = 0.0f32;
for dd in 0..dsub {
- ip += sub_centroid[dd] * pq.centroids[pq_off + dd];
+ ip += sub_centroid[dd] * pq.centroids()[pq_off + dd];
}
list_table[sub * ksub + j] = pq_norms[sub * ksub + j] +
2.0 * ip;
}
@@ -1547,7 +1528,7 @@ mod tests {
let mut cursor = Cursor::new(&buf);
let mut reader = IVFPQIndexReader::open(&mut cursor).unwrap();
- assert_eq!(reader.pq.nbits, 4);
+ assert_eq!(reader.pq.nbits(), 4);
assert_eq!(reader.pq.code_size(), m / 2);
let (result_ids, result_dists) = reader.search(&data[0..d], 5,
4).unwrap();
diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs
index 15e530c..a7b1082 100644
--- a/core/src/ivfpq.rs
+++ b/core/src/ivfpq.rs
@@ -80,6 +80,7 @@ pub struct IVFPQIndex {
/// Block-layout packed codes for 4-bit FastScan. One per list.
fastscan_codes: Vec<Vec<u8>>,
coarse_assignment: CoarseAssignment,
+ canonical_pq_encoding: bool,
}
impl IVFPQIndex {
@@ -117,6 +118,7 @@ impl IVFPQIndex {
precomputed_table: Vec::new(),
fastscan_codes: Vec::new(),
coarse_assignment: CoarseAssignment::default(),
+ canonical_pq_encoding: false,
}
}
@@ -124,6 +126,16 @@ impl IVFPQIndex {
&self.quantizer_centroids
}
+ /// Uses the canonical per-vector PQ encoder instead of the default
+ /// automatic encoder. This build-time setting is not serialized.
+ pub(crate) fn set_canonical_pq_encoding(&mut self, canonical: bool) {
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot change PQ encoding after vectors have been added"
+ );
+ self.canonical_pq_encoding = canonical;
+ }
+
/// Enables automatic Vamana coarse assignment for large centroid matrices.
/// Disable it to keep vector assignment exact.
pub(crate) fn set_approximate_coarse_assignment(&mut self, enabled: bool) {
@@ -174,16 +186,7 @@ impl IVFPQIndex {
metric: trained.metric,
by_residual: trained.by_residual,
quantizer_centroids: trained.quantizer_centroids.clone(),
- pq: ProductQuantizer {
- d: trained.pq.d,
- m: trained.pq.m,
- nbits: trained.pq.nbits,
- dsub: trained.pq.dsub,
- ksub: trained.pq.ksub,
- chunk_offsets: trained.pq.chunk_offsets.clone(),
- centroids: trained.pq.centroids.clone(),
- centroid_norms_cache: trained.pq.centroid_norms_cache.clone(),
- },
+ pq: trained.pq.clone_without_transposed_cache(),
opq: trained.opq.as_ref().map(|o| OPQMatrix {
d: o.d,
m: o.m,
@@ -199,6 +202,7 @@ impl IVFPQIndex {
precomputed_table: Vec::new(),
fastscan_codes: Vec::new(),
coarse_assignment: CoarseAssignment::default(),
+ canonical_pq_encoding: trained.canonical_pq_encoding,
};
index.set_approximate_coarse_assignment(trained.coarse_assignment.approximate_enabled());
index
@@ -303,7 +307,11 @@ impl IVFPQIndex {
let code_size = self.pq.code_size();
let mut codes = vec![0u8; n * code_size];
- self.pq.encode_batch_blocked(&to_encode, n, &mut codes);
+ if self.canonical_pq_encoding {
+ self.pq.encode_batch(&to_encode, n, &mut codes);
+ } else {
+ self.pq.encode_batch_blocked(&to_encode, n, &mut codes);
+ }
for i in 0..n {
let list_id = assignments[i];
@@ -322,7 +330,7 @@ impl IVFPQIndex {
/// Build fastscan block codes for 4-bit search acceleration.
/// Call after all vectors are added. Lightweight — only reorganizes
existing codes.
pub fn build_search_structures(&mut self) {
- if self.pq.nbits == 4 {
+ if self.pq.nbits() == 4 {
let cs = self.pq.code_size();
self.fastscan_codes = self
.codes
@@ -345,8 +353,8 @@ impl IVFPQIndex {
/// Costs ~10ms to build and uses nlist * M * ksub * 4 bytes of memory.
pub fn build_precomputed_table(&mut self) {
let d = self.d;
- let m = self.pq.m;
- let ksub = self.pq.ksub;
+ let m = self.pq.m();
+ let ksub = self.pq.ksub();
let nlist = self.nlist;
if self.metric != MetricType::L2 || !self.by_residual {
@@ -362,14 +370,15 @@ impl IVFPQIndex {
.for_each(|(i, list_table)| {
let centroid = &self.quantizer_centroids[i * d..(i + 1) *
d];
for sub in 0..m {
- let sub_centroid = ¢roid[sub * self.pq.dsub..(sub
+ 1) * self.pq.dsub];
- let pq_base = sub * ksub * self.pq.dsub;
+ let sub_centroid =
+ ¢roid[sub * self.pq.dsub()..(sub + 1) *
self.pq.dsub()];
+ let pq_base = sub * ksub * self.pq.dsub();
for j in 0..ksub {
- let pq_off = pq_base + j * self.pq.dsub;
+ let pq_off = pq_base + j * self.pq.dsub();
let ip = fvec_inner_product(
sub_centroid,
- &self.pq.centroids[pq_off..pq_off +
self.pq.dsub],
+ &self.pq.centroids()[pq_off..pq_off +
self.pq.dsub()],
);
list_table[sub * ksub + j] = pq_norms[sub * ksub +
j] + 2.0 * ip;
}
@@ -413,8 +422,8 @@ impl IVFPQIndex {
result_labels: &mut [i64],
) {
let d = self.d;
- let m = self.pq.m;
- let ksub = self.pq.ksub;
+ let m = self.pq.m();
+ let ksub = self.pq.ksub();
let processed_queries = self.preprocess_queries(queries, nq);
@@ -428,7 +437,7 @@ impl IVFPQIndex {
);
let use_precomputed = !self.precomputed_table.is_empty();
- let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits ==
4;
+ let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits()
== 4;
let matching_rows_by_list = filter.map(|filter| {
let mut probed_lists = vec![false; self.nlist];
for probe_indices in &all_probe_indices {
@@ -524,7 +533,7 @@ impl IVFPQIndex {
heap.push(dis0 + dists[i],
self.ids[list_id][i]);
}
}
- } else if self.pq.nbits == 4 {
+ } else if self.pq.nbits() == 4 {
scan_codes_4bit(
&sim_table,
&self.codes[list_id],
@@ -621,8 +630,8 @@ impl IVFPQIndex {
result_labels: &mut [i64],
) {
let d = self.d;
- let m = self.pq.m;
- let ksub = self.pq.ksub;
+ let m = self.pq.m();
+ let ksub = self.pq.ksub();
let processed_queries = self.preprocess_queries(queries, nq);
let (all_probe_indices, all_coarse_dists) = kmeans::find_topk_batch(
@@ -635,7 +644,7 @@ impl IVFPQIndex {
);
let use_precomputed = !self.precomputed_table.is_empty();
- let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits ==
4;
+ let use_fastscan = !self.fastscan_codes.is_empty() && self.pq.nbits()
== 4;
let results: Vec<Vec<(f32, i64)>> = (0..nq)
.into_par_iter()
@@ -697,7 +706,7 @@ impl IVFPQIndex {
for i in 0..scan_count {
heap.push(dis0 + dists[i], self.ids[list_id][i]);
}
- } else if self.pq.nbits == 4 {
+ } else if self.pq.nbits() == 4 {
scan_codes_4bit(
&sim_table,
&self.codes[list_id],
@@ -786,24 +795,24 @@ impl IVFPQIndex {
self.by_residual, other.by_residual
)));
}
- if self.pq.d != other.pq.d
- || self.pq.m != other.pq.m
- || self.pq.nbits != other.pq.nbits
- || self.pq.dsub != other.pq.dsub
- || self.pq.ksub != other.pq.ksub
+ if self.pq.d() != other.pq.d()
+ || self.pq.m() != other.pq.m()
+ || self.pq.nbits() != other.pq.nbits()
+ || self.pq.dsub() != other.pq.dsub()
+ || self.pq.ksub() != other.pq.ksub()
{
return Err(invalid_merge_input(format!(
"PQ layout mismatch: self=(d={}, m={}, nbits={}, dsub={},
ksub={}), other=(d={}, m={}, nbits={}, dsub={}, ksub={})",
- self.pq.d,
- self.pq.m,
- self.pq.nbits,
- self.pq.dsub,
- self.pq.ksub,
- other.pq.d,
- other.pq.m,
- other.pq.nbits,
- other.pq.dsub,
- other.pq.ksub
+ self.pq.d(),
+ self.pq.m(),
+ self.pq.nbits(),
+ self.pq.dsub(),
+ self.pq.ksub(),
+ other.pq.d(),
+ other.pq.m(),
+ other.pq.nbits(),
+ other.pq.dsub(),
+ other.pq.ksub()
)));
}
if self.opq.is_some() != other.opq.is_some() {
@@ -823,7 +832,7 @@ impl IVFPQIndex {
if self.quantizer_centroids != other.quantizer_centroids {
return Err(invalid_merge_input("coarse centroids mismatch"));
}
- if self.pq.centroids != other.pq.centroids {
+ if self.pq.centroids() != other.pq.centroids() {
return Err(invalid_merge_input("PQ codebooks mismatch"));
}
@@ -1034,37 +1043,37 @@ fn fill_list_precomputed_table(
pq_norms: &[f32],
table: &mut Vec<f32>,
) {
- debug_assert_eq!(coarse_centroid.len(), pq.d);
- debug_assert_eq!(pq_norms.len(), pq.m * pq.ksub);
- table.resize(pq.m * pq.ksub, 0.0);
- for sub in 0..pq.m {
+ debug_assert_eq!(coarse_centroid.len(), pq.d());
+ debug_assert_eq!(pq_norms.len(), pq.m() * pq.ksub());
+ table.resize(pq.m() * pq.ksub(), 0.0);
+ for sub in 0..pq.m() {
let range = pq.chunk_range(sub);
let chunk_dim = range.len();
- let pq_base = range.start * pq.ksub;
- for code in 0..pq.ksub {
+ let pq_base = range.start * pq.ksub();
+ for code in 0..pq.ksub() {
let pq_offset = pq_base + code * chunk_dim;
let mut inner_product = 0.0f32;
for dimension in 0..chunk_dim {
- inner_product +=
- coarse_centroid[range.start + dimension] *
pq.centroids[pq_offset + dimension];
+ inner_product += coarse_centroid[range.start + dimension]
+ * pq.centroids()[pq_offset + dimension];
}
- let table_offset = sub * pq.ksub + code;
+ let table_offset = sub * pq.ksub() + code;
table[table_offset] = pq_norms[table_offset] + 2.0 * inner_product;
}
}
}
fn compute_stable_ephemeral_pq_norms(pq: &ProductQuantizer) -> Vec<f64> {
- let mut norms = vec![0.0f64; pq.m * pq.ksub];
- for sub in 0..pq.m {
+ let mut norms = vec![0.0f64; pq.m() * pq.ksub()];
+ for sub in 0..pq.m() {
let range = pq.chunk_range(sub);
let chunk_dim = range.len();
- let pq_base = range.start * pq.ksub;
- for code in 0..pq.ksub {
+ let pq_base = range.start * pq.ksub();
+ for code in 0..pq.ksub() {
let pq_offset = pq_base + code * chunk_dim;
- norms[sub * pq.ksub + code] = (0..chunk_dim)
+ norms[sub * pq.ksub() + code] = (0..chunk_dim)
.map(|dimension| {
- let value = f64::from(pq.centroids[pq_offset + dimension]);
+ let value = f64::from(pq.centroids()[pq_offset +
dimension]);
value * value
})
.sum();
@@ -1079,38 +1088,38 @@ fn fill_stable_ephemeral_list_table(
pq_norms: &[f64],
table: &mut Vec<f64>,
) {
- table.resize(pq.m * pq.ksub, 0.0);
- for sub in 0..pq.m {
+ table.resize(pq.m() * pq.ksub(), 0.0);
+ for sub in 0..pq.m() {
let range = pq.chunk_range(sub);
let chunk_dim = range.len();
- let pq_base = range.start * pq.ksub;
- for code in 0..pq.ksub {
+ let pq_base = range.start * pq.ksub();
+ for code in 0..pq.ksub() {
let pq_offset = pq_base + code * chunk_dim;
let mut inner_product = 0.0f64;
for dimension in 0..chunk_dim {
- let pq_value = f64::from(pq.centroids[pq_offset + dimension]);
+ let pq_value = f64::from(pq.centroids()[pq_offset +
dimension]);
inner_product += f64::from(coarse_centroid[range.start +
dimension]) * pq_value;
}
- let offset = sub * pq.ksub + code;
+ let offset = sub * pq.ksub() + code;
table[offset] = pq_norms[offset] + 2.0 * inner_product;
}
}
}
fn fill_stable_ephemeral_query_table(query: &[f32], pq: &ProductQuantizer,
table: &mut Vec<f64>) {
- table.resize(pq.m * pq.ksub, 0.0);
- for sub in 0..pq.m {
+ table.resize(pq.m() * pq.ksub(), 0.0);
+ for sub in 0..pq.m() {
let range = pq.chunk_range(sub);
let chunk_dim = range.len();
- let pq_base = range.start * pq.ksub;
- for code in 0..pq.ksub {
+ let pq_base = range.start * pq.ksub();
+ for code in 0..pq.ksub() {
let pq_offset = pq_base + code * chunk_dim;
let mut inner_product = 0.0f64;
for dimension in 0..chunk_dim {
inner_product += f64::from(query[range.start + dimension])
- * f64::from(pq.centroids[pq_offset + dimension]);
+ * f64::from(pq.centroids()[pq_offset + dimension]);
}
- table[sub * pq.ksub + code] = inner_product;
+ table[sub * pq.ksub() + code] = inner_product;
}
}
}
@@ -1123,16 +1132,16 @@ fn combine_stable_ephemeral_tables(
pq: &ProductQuantizer,
sim_table: &mut Vec<f32>,
) {
- sim_table.resize(pq.m * pq.ksub, 0.0);
- for sub in 0..pq.m {
+ sim_table.resize(pq.m() * pq.ksub(), 0.0);
+ for sub in 0..pq.m() {
let range = pq.chunk_range(sub);
let mut residual_norm = 0.0f64;
for dimension in range {
let residual = f64::from(query[dimension]) -
f64::from(coarse_centroid[dimension]);
residual_norm += residual * residual;
}
- let table_base = sub * pq.ksub;
- for code in 0..pq.ksub {
+ let table_base = sub * pq.ksub();
+ for code in 0..pq.ksub() {
let offset = table_base + code;
sim_table[offset] =
(residual_norm + list_table[offset] - 2.0 *
query_table[offset]).max(0.0) as f32;
@@ -1531,7 +1540,7 @@ pub fn search_with_reader_filter<R: SeekRead>(
let (_, _, dis0) = lists_to_read[batch_start];
let sim_table = by_residual
.then(|| reader_sim_table(reader, first_list, &q, &ip_table,
use_precomputed));
- let pq_nbits = reader.pq.nbits;
+ let pq_nbits = reader.pq.nbits();
let transposed_codes = reader.transposed_codes;
let mut scratch = ReaderScanScratch::default();
reader.for_each_streamed_list_chunk(first_list, |pq, ids, codes| {
@@ -1674,7 +1683,7 @@ fn scan_reader_list(
&entry.ids,
ctx.m,
ctx.ksub,
- ctx.pq.nbits,
+ ctx.pq.nbits(),
ctx.transposed_codes,
dis0,
matching_rows,
@@ -2274,7 +2283,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
&& by_residual
&& !reader.precomputed_table.is_empty()
&& reused_query_tables_fit_budget;
- let allow_ephemeral_precomputed = reader.pq.nbits == 8
+ let allow_ephemeral_precomputed = reader.pq.nbits() == 8
&& metric == MetricType::L2
&& by_residual
&& !use_precomputed
@@ -2299,7 +2308,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
};
// Non-residual tables depend only on the query and PQ codebook, so every
// probed list for that query can share one table.
- let reuse_non_residual_tables = reader.pq.nbits == 8
+ let reuse_non_residual_tables = reader.pq.nbits() == 8
&& !by_residual
&& probe_end - probe_start > 1
&& match reuse_mode {
@@ -2375,7 +2384,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
})
})
.collect::<Vec<_>>();
- let pq_nbits = reader.pq.nbits;
+ let pq_nbits = reader.pq.nbits();
let transposed_codes = reader.transposed_codes;
// The loop is sequential across queries. Reuse one chunk-sized
// distance buffer instead of retaining one per query.
@@ -2471,7 +2480,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
.get(list.list_id as u32)
.copied()
.unwrap_or_default(),
- reader.pq.nbits,
+ reader.pq.nbits(),
reader.transposed_codes,
);
}
@@ -2615,7 +2624,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
&loaded_lists[position].ids,
m,
ksub,
- reader.pq.nbits,
+ reader.pq.nbits(),
reader.transposed_codes,
dis0,
matching_rows_by_list[position].as_ref(),
@@ -2680,7 +2689,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
elapsed_since(total_started),
nq,
scanned_nprobe,
- reader.pq.nbits,
+ reader.pq.nbits(),
k,
unique_lists.len(),
filter.is_some(),
@@ -2700,7 +2709,7 @@ fn
search_batch_reader_filter_with_reuse_mode_and_observer<R: SeekRead>(
budget_bytes={reuse_max_bytes} tables_built={tables_built}",
tables_built > 0,
metric.as_str(),
- reader.pq.nbits,
+ reader.pq.nbits(),
unique_lists.len(),
filter.is_some(),
reuse_required_bytes,
@@ -3303,7 +3312,7 @@ mod tests {
let ids: Vec<i64> = (0..n as i64).collect();
let mut index = IVFPQIndex::with_nbits(d, nlist, m, 4, MetricType::L2,
false);
- assert_eq!(index.pq.ksub, 16);
+ assert_eq!(index.pq.ksub(), 16);
assert_eq!(index.pq.code_size(), 4);
index.train(&data, n);
@@ -3365,6 +3374,34 @@ mod tests {
assert!(dists_full[0] <= dists_limited[0] + 1e-6);
}
+ #[test]
+ fn canonical_pq_encoding_is_preserved_by_from_trained() {
+ let mut trainer = IVFPQIndex::new(4, 1, 1, MetricType::L2, false);
+ assert!(!trainer.canonical_pq_encoding);
+ trainer.set_quantizer_centroids(vec![0.0; 4]);
+ let mut centroids = vec![100_000_016.0; 4 * 256];
+ centroids[0..4].fill(100_000_008.0);
+ centroids[4..8].fill(100_000_000.0);
+ trainer.pq.set_centroids(centroids);
+ trainer.set_canonical_pq_encoding(true);
+
+ let n = 7;
+ let data = vec![100_000_000.0; n * 4];
+ let mut expected = vec![0; n];
+ trainer.pq.encode_batch(&data, n, &mut expected);
+ assert_eq!(expected, vec![0; n]);
+
+ let mut worker = IVFPQIndex::from_trained(&trainer);
+ assert!(worker.canonical_pq_encoding);
+ worker.add(&data, &(0..n as i64).collect::<Vec<_>>(), n);
+ assert_eq!(worker.codes[0], expected);
+
+ assert!(std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+ worker.set_canonical_pq_encoding(false);
+ }))
+ .is_err());
+ }
+
#[test]
fn test_from_trained_and_merge() {
let d = 16;
@@ -3435,7 +3472,9 @@ mod tests {
assert_invalid_merge(&base, &mismatched_centroids, "coarse centroids
mismatch");
let mut mismatched_codebooks = IVFPQIndex::from_trained(&trainer);
- mismatched_codebooks.pq.centroids[0] += 1.0;
+ let mut centroids = mismatched_codebooks.pq.centroids().to_vec();
+ centroids[0] += 1.0;
+ mismatched_codebooks.pq.set_centroids(centroids);
assert_invalid_merge(&base, &mismatched_codebooks, "PQ codebooks
mismatch");
let mismatched_opq = IVFPQIndex::new(d, nlist, m, MetricType::L2,
true);
@@ -3837,7 +3876,7 @@ mod tests {
&pq_norms,
&mut actual,
);
- let table_size = m * index.pq.ksub;
+ let table_size = m * index.pq.ksub();
assert_eq!(
actual,
index.precomputed_table[list_id * table_size..(list_id + 1) *
table_size]
diff --git a/core/src/pq.rs b/core/src/pq.rs
index edc6a3e..96fd06f 100644
--- a/core/src/pq.rs
+++ b/core/src/pq.rs
@@ -21,6 +21,7 @@ use crate::distance::{
};
use crate::kmeans::{self, KMeansConfig};
use rayon::prelude::*;
+use std::sync::OnceLock;
/// Product Quantizer aligned with Faiss's ProductQuantizer.
///
@@ -31,23 +32,63 @@ use rayon::prelude::*;
///
/// Centroids are chunk-major. Chunk `m` starts at
/// `chunk_offsets[m] * ksub`, and each of its `ksub` centroids contains
-/// `chunk_offsets[m + 1] - chunk_offsets[m]` contiguous components.
+/// `chunk_offsets[m + 1] - chunk_offsets[m]` contiguous components. Use
+/// [`Self::set_centroids`] to replace them so derived caches stay
synchronized.
+/// Layout is immutable after construction so derived caches stay valid.
+///
+/// ```compile_fail,E0616
+/// use paimon_vindex_core::pq::ProductQuantizer;
+/// let mut pq = ProductQuantizer::new(12, 2);
+/// pq.chunk_offsets = vec![0, 4, 12];
+/// ```
pub struct ProductQuantizer {
- pub d: usize,
- pub m: usize,
- pub nbits: usize,
+ d: usize,
+ m: usize,
+ nbits: usize,
/// Uniform chunk width for legacy formats, or the largest chunk width for
/// a balanced non-uniform layout.
- pub dsub: usize,
- pub ksub: usize,
- pub chunk_offsets: Vec<usize>,
- pub centroids: Vec<f32>,
+ dsub: usize,
+ ksub: usize,
+ chunk_offsets: Vec<usize>,
+ centroids: Vec<f32>,
+ centroids_are_finite: bool,
/// Pre-computed squared norms of each centroid: [M * ksub].
/// Avoids recomputing per query for L2 distance table.
- pub centroid_norms_cache: Vec<f32>,
+ centroid_norms_cache: Vec<f32>,
+ transposed_codebook_cache: OnceLock<(Vec<f32>, usize)>,
}
impl ProductQuantizer {
+ /// Vector dimension.
+ pub fn d(&self) -> usize {
+ self.d
+ }
+
+ /// Number of subquantizers.
+ pub fn m(&self) -> usize {
+ self.m
+ }
+
+ /// Bits per PQ code.
+ pub fn nbits(&self) -> usize {
+ self.nbits
+ }
+
+ /// Largest subvector dimension.
+ pub fn dsub(&self) -> usize {
+ self.dsub
+ }
+
+ /// Centroids per subquantizer.
+ pub fn ksub(&self) -> usize {
+ self.ksub
+ }
+
+ /// Contiguous subvector boundaries.
+ pub fn chunk_offsets(&self) -> &[usize] {
+ &self.chunk_offsets
+ }
+
pub fn new(d: usize, m: usize) -> Self {
Self::with_nbits(d, m, 8)
}
@@ -125,7 +166,9 @@ impl ProductQuantizer {
ksub,
chunk_offsets,
centroids: Vec::new(),
+ centroids_are_finite: true,
centroid_norms_cache: Vec::new(),
+ transposed_codebook_cache: OnceLock::new(),
}
}
@@ -155,6 +198,50 @@ impl ProductQuantizer {
)
}
+ /// Return the chunk-major PQ codebook.
+ pub fn centroids(&self) -> &[f32] {
+ &self.centroids
+ }
+
+ /// Replace the PQ codebook and refresh its derived norm and transpose
caches.
+ pub fn set_centroids(&mut self, centroids: Vec<f32>) {
+ self.try_set_centroids(centroids)
+ .expect("PQ centroid norms allocation failed");
+ }
+
+ pub(crate) fn try_set_centroids(
+ &mut self,
+ centroids: Vec<f32>,
+ ) -> Result<(), std::collections::TryReserveError> {
+ assert_eq!(
+ centroids.len(),
+ self.d * self.ksub,
+ "PQ centroids must hold d * ksub values"
+ );
+ let norms = self.try_compute_centroid_norms(¢roids)?;
+ self.centroids_are_finite = centroids.iter().all(|value|
value.is_finite());
+ self.centroids = centroids;
+ self.centroid_norms_cache = norms;
+ self.transposed_codebook_cache.take();
+ Ok(())
+ }
+
+ /// Copy trained state while leaving the writer-local transpose lazy.
+ pub(crate) fn clone_without_transposed_cache(&self) -> Self {
+ Self {
+ d: self.d,
+ m: self.m,
+ nbits: self.nbits,
+ dsub: self.dsub,
+ ksub: self.ksub,
+ chunk_offsets: self.chunk_offsets.clone(),
+ centroids: self.centroids.clone(),
+ centroids_are_finite: self.centroids_are_finite,
+ centroid_norms_cache: self.centroid_norms_cache.clone(),
+ transposed_codebook_cache: OnceLock::new(),
+ }
+ }
+
/// Train the codebooks from training data.
/// data: flat [n * d], n training vectors.
pub fn train(&mut self, data: &[f32], n: usize) {
@@ -243,23 +330,19 @@ impl ProductQuantizer {
.install(train_subquantizers)
};
- self.centroids = vec![0.0f32; d * ksub];
+ let mut centroids = vec![0.0f32; d * ksub];
for (sub, sub_centroids) in sub_results.into_iter().enumerate() {
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);
+ centroids[dst_offset..dst_offset + ksub *
chunk_dim].copy_from_slice(&sub_centroids);
}
- self.rebuild_norms_cache();
+ self.set_centroids(centroids);
}
- /// Rebuild the centroid norms cache. Called after training or loading
centroids.
- pub fn rebuild_norms_cache(&mut self) {
- self.try_rebuild_norms_cache()
- .expect("PQ centroid norms allocation failed");
- }
-
- pub fn try_rebuild_norms_cache(&mut self) -> Result<(),
std::collections::TryReserveError> {
+ fn try_compute_centroid_norms(
+ &self,
+ centroids: &[f32],
+ ) -> Result<Vec<f32>, std::collections::TryReserveError> {
let mut norms = Vec::new();
norms.try_reserve_exact(self.m * self.ksub)?;
norms.resize(self.m * self.ksub, 0.0f32);
@@ -268,12 +351,10 @@ impl ProductQuantizer {
let c_base = self.centroid_chunk_base(sub);
for j in 0..self.ksub {
let c_off = c_base + j * chunk_dim;
- norms[sub * self.ksub + j] =
- fvec_norm_l2sqr(&self.centroids[c_off..c_off + chunk_dim]);
+ norms[sub * self.ksub + j] =
fvec_norm_l2sqr(¢roids[c_off..c_off + chunk_dim]);
}
}
- self.centroid_norms_cache = norms;
- Ok(())
+ Ok(norms)
}
/// Bytes per encoded vector.
@@ -357,11 +438,10 @@ impl ProductQuantizer {
/// Encode multiple vectors in parallel.
///
- /// This is the byte-stable path: results are bit-identical to the
- /// per-vector [`Self::encode`], which golden storage fixtures rely on
- /// (DiskANN serializes these codes). IVF-PQ's add path uses
- /// [`Self::encode_batch_blocked`] instead, which batches the same distance
- /// calculation into larger SGEMM calls.
+ /// This is the canonical path: results are bit-identical to the
+ /// per-vector [`Self::encode`] on the same runtime backend, which golden
+ /// storage fixtures rely on (DiskANN serializes these codes). IVF-PQ's add
+ /// path uses [`Self::encode_batch_blocked`] instead.
pub fn encode_batch(&self, data: &[f32], n: usize, codes: &mut [u8]) {
let d = self.d;
let cs = self.code_size();
@@ -376,19 +456,41 @@ impl ProductQuantizer {
);
}
- /// Blocked batch encode for the IVF-PQ add path.
+ /// Automatic batch encode for the IVF-PQ add path.
+ ///
+ /// The backend depends only on shape and CPU features, never on `n`:
+ /// - 8-bit shapes with `dsub >= 4` use the transposed direct-L2 encoder
+ /// when this CPU has a fast fused kernel.
+ /// - The same finite shape uses blocked SGEMM otherwise, avoiding billions
+ /// of software `fmaf` calls on x86 CPUs without FMA.
+ /// - Other shapes, and non-finite codebooks on the SGEMM fallback, use the
+ /// canonical encoder.
+ ///
+ /// Each backend is deterministic across thread counts and batch splits,
+ /// but their floating-point contracts differ. The transposed path treats
+ /// NaN distances as losing, lets infinite distances lose, returns code 0
+ /// when no distance is below `f32::MAX`, and computes direct squared L2.
+ /// The SGEMM and canonical paths use the expanded form
+ /// `|q|² + |c|² - 2q·c`, including its existing NaN and large-offset
+ /// behavior. Canonical mode reproduces the per-vector encoder on the same
+ /// CPU and runtime backend; neither mode promises cross-CPU code identity.
pub(crate) fn encode_batch_blocked(&self, data: &[f32], n: usize, codes:
&mut [u8]) {
- if self.nbits == 8
- && n >= encode_sgemm_min_rows(rayon::current_num_threads())
- && (0..self.m).all(|sub| self.chunk_dim(sub) >= 4)
- && self.centroids.iter().all(|value| value.is_finite())
- {
- self.encode_batch_8bit_sgemm(data, n, codes);
+ if self.nbits == 8 && self.ksub == 256 && (0..self.m).all(|sub|
self.chunk_dim(sub) >= 4) {
+ if supports_transposed_pq_encoding() {
+ self.encode_batch_8bit_transposed(data, n, codes);
+ } else if self.centroids_are_finite {
+ self.encode_batch_8bit_sgemm(data, n, codes);
+ } else {
+ self.encode_batch(data, n, codes);
+ }
return;
}
self.encode_batch(data, n, codes);
}
+ /// Blocked SGEMM fallback for CPUs without a fast transposed kernel.
+ /// Unlike the former main path, this runs for every `n`; block sizes only
+ /// divide work and never select another encoder.
fn encode_batch_8bit_sgemm(&self, data: &[f32], n: usize, codes: &mut
[u8]) {
let d = self.d;
let m = self.m;
@@ -396,6 +498,8 @@ impl ProductQuantizer {
let cs = self.code_size();
debug_assert_eq!(cs, m);
debug_assert_eq!(ksub, 256);
+ debug_assert!((0..m).all(|sub| self.chunk_dim(sub) >= 4));
+ debug_assert!(self.centroids_are_finite);
let max_dsub = (0..m).map(|sub|
self.chunk_dim(sub)).max().unwrap_or(0);
let computed_norms = self
@@ -455,6 +559,73 @@ impl ProductQuantizer {
);
}
+ /// Transposed codebook: per sub a `[dsub][ksub]` block at a uniform
+ /// stride of `max_dsub * ksub`, so sub lookup stays O(1) for balanced
+ /// non-uniform chunks. Returns the table and the per-sub stride.
+ fn build_transposed_codebook(&self) -> (Vec<f32>, usize) {
+ let m = self.m;
+ let ksub = self.ksub;
+ let max_dsub = (0..m).map(|sub|
self.chunk_dim(sub)).max().unwrap_or(0);
+ let sub_stride = max_dsub
+ .checked_mul(ksub)
+ .expect("transposed codebook stride overflows usize");
+ let mut transposed = vec![0.0f32; m * sub_stride];
+ for sub in 0..m {
+ let dsub = self.chunk_dim(sub);
+ let c_base = self.centroid_chunk_base(sub);
+ let dst = &mut transposed[sub * sub_stride..sub * sub_stride +
dsub * ksub];
+ for j in 0..ksub {
+ for k in 0..dsub {
+ dst[k * ksub + j] = self.centroids[c_base + j * dsub + k];
+ }
+ }
+ }
+ (transposed, sub_stride)
+ }
+
+ /// Transposed-codebook encode for nbits=8. Rows are split into blocks
+ /// only for parallelism; each row is encoded independently.
+ fn encode_batch_8bit_transposed(&self, data: &[f32], n: usize, codes: &mut
[u8]) {
+ let d = self.d;
+ let m = self.m;
+ let ksub = self.ksub;
+ let cs = self.code_size();
+ debug_assert_eq!(cs, m);
+
+ let (transposed, sub_stride) = self
+ .transposed_codebook_cache
+ .get_or_init(|| self.build_transposed_codebook());
+ let sub_stride = *sub_stride;
+ let kernels: Vec<ScoreArgminKernel> = (0..m)
+ .map(|sub| score_argmin_kernel(self.chunk_dim(sub), ksub))
+ .collect();
+
+ let block_rows = encode_block_rows(n, rayon::current_num_threads());
+ codes[..n * cs]
+ .par_chunks_mut(block_rows * cs)
+ .enumerate()
+ .for_each_init(
+ || vec![0.0f32; ksub],
+ |scores, (block_idx, block_codes)| {
+ let row0 = block_idx * block_rows;
+ let rows = block_rows.min(n - row0);
+ let block_data = &data[row0 * d..(row0 + rows) * d];
+
+ for r in 0..rows {
+ let row = &block_data[r * d..(r + 1) * d];
+ for sub in 0..m {
+ let range = self.chunk_range(sub);
+ let dsub = range.len();
+ let q = &row[range];
+ let t = &transposed[sub * sub_stride..sub *
sub_stride + dsub * ksub];
+ block_codes[r * cs + sub] =
+ score_argmin(kernels[sub], q, t, ksub, scores);
+ }
+ }
+ },
+ );
+ }
+
/// Decode PQ codes back to an approximate vector.
pub fn decode(&self, codes: &[u8], x: &mut [f32]) {
for sub in 0..self.m {
@@ -612,20 +783,316 @@ fn argmin_code(distances: &[f32]) -> u8 {
best as u8
}
-/// Row block for the batched SGEMM encode path.
-const MAX_ENCODE_BLOCK_ROWS: usize = 512;
-const MIN_ENCODE_BLOCK_ROWS: usize = 4;
-const ENCODE_SGEMM_MIN_ROWS: usize = 32;
+#[cfg(target_arch = "x86_64")]
+pub(crate) fn supports_transposed_pq_encoding() -> bool {
+ is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma")
+}
-fn encode_sgemm_min_rows(workers: usize) -> usize {
- ENCODE_SGEMM_MIN_ROWS.max(workers.max(1) * MIN_ENCODE_BLOCK_ROWS)
+#[cfg(target_arch = "aarch64")]
+pub(crate) fn supports_transposed_pq_encoding() -> bool {
+ true
}
+#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
+pub(crate) fn supports_transposed_pq_encoding() -> bool {
+ false
+}
+
+/// Row block for the automatic batch-encode paths. 512 rows keeps scratch
+/// buffers hot and yields enough blocks to occupy the worker pool. Blocks
+/// only divide work; they never select another encoder.
+const MAX_ENCODE_BLOCK_ROWS: usize = 512;
+
fn encode_block_rows(rows: usize, workers: usize) -> usize {
rows.div_ceil(workers.max(1))
.clamp(1, MAX_ENCODE_BLOCK_ROWS)
}
+/// Squared-L2 argmin kernel over a transposed sub-codebook `t` laid out as
+/// `[dsub][ksub]` (stride-1 over `j`). `scores` is a reusable `ksub`-sized
+/// scratch buffer that some kernels do not touch.
+type ScoreArgminKernel = fn(&[f32], &[f32], usize, &mut [f32]) -> u8;
+
+/// Pick the kernel for one sub-quantizer. The choice depends only on `dsub`
+/// and CPU features, never on the batch size, which is what makes
+/// [`ProductQuantizer::encode_batch_blocked`] batch invariant. Every kernel
+/// returned here is byte-identical to [`score_argmin_scalar`].
+fn score_argmin_kernel(dsub: usize, ksub: usize) -> ScoreArgminKernel {
+ #[cfg(target_arch = "aarch64")]
+ {
+ if dsub == 4 && ksub.is_multiple_of(4) {
+ return score_argmin_neon_d4_entry;
+ }
+ if dsub >= 4 && ksub.is_multiple_of(16) {
+ return score_argmin_neon_generic_entry;
+ }
+ }
+ #[cfg(target_arch = "x86_64")]
+ {
+ if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma")
{
+ if dsub == 4 && ksub.is_multiple_of(8) {
+ return score_argmin_avx2_d4_entry;
+ }
+ return score_argmin_avx2_generic_entry;
+ }
+ }
+ let _ = (dsub, ksub);
+ score_argmin_scalar
+}
+
+/// Run `kernel` after checking the slice contract shared by all kernels.
+#[inline]
+fn score_argmin(
+ kernel: ScoreArgminKernel,
+ q: &[f32],
+ t: &[f32],
+ ksub: usize,
+ scores: &mut [f32],
+) -> u8 {
+ assert_eq!(q.len().checked_mul(ksub), Some(t.len()));
+ assert!(scores.len() >= ksub);
+ kernel(q, t, ksub, scores)
+}
+
+/// Direct squared-difference accumulation, one `mul_add` per dimension so
+/// every kernel rounds once per dimension in the same order.
+#[inline(always)]
+fn score_argmin_accumulate(q: &[f32], t: &[f32], ksub: usize, scores: &mut
[f32]) -> u8 {
+ let scores = &mut scores[..ksub];
+ scores.fill(0.0);
+ for (k, &qv) in q.iter().enumerate() {
+ let tk = &t[k * ksub..(k + 1) * ksub];
+ for (s, &tv) in scores.iter_mut().zip(tk) {
+ let diff = qv - tv;
+ *s = diff.mul_add(diff, *s);
+ }
+ }
+ argmin_code(scores)
+}
+
+/// Scalar oracle for every SIMD kernel.
+fn score_argmin_scalar(q: &[f32], t: &[f32], ksub: usize, scores: &mut [f32])
-> u8 {
+ score_argmin_accumulate(q, t, ksub, scores)
+}
+
+/// Generic AVX2+FMA kernel for any `dsub`: the same loop as the scalar
+/// oracle compiled with the target features enabled so the stride-1 inner
+/// loop vectorizes with `vfmadd`.
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2", enable = "fma")]
+unsafe fn score_argmin_avx2_generic(q: &[f32], t: &[f32], ksub: usize, scores:
&mut [f32]) -> u8 {
+ score_argmin_accumulate(q, t, ksub, scores)
+}
+
+#[cfg(target_arch = "x86_64")]
+fn score_argmin_avx2_generic_entry(q: &[f32], t: &[f32], ksub: usize, scores:
&mut [f32]) -> u8 {
+ // SAFETY: only handed out by `score_argmin_kernel` after AVX2 and FMA were
+ // detected on this CPU.
+ unsafe { score_argmin_avx2_generic(q, t, ksub, scores) }
+}
+
+#[cfg(target_arch = "x86_64")]
+fn score_argmin_avx2_d4_entry(q: &[f32], t: &[f32], ksub: usize, _scores: &mut
[f32]) -> u8 {
+ debug_assert_eq!(q.len(), 4);
+ debug_assert!(ksub.is_multiple_of(8));
+ // SAFETY: only handed out by `score_argmin_kernel` after AVX2 and FMA were
+ // detected on this CPU; slice bounds are checked by `score_argmin`.
+ unsafe { score_argmin_avx2_d4(q, t, ksub) }
+}
+
+#[cfg(target_arch = "aarch64")]
+fn score_argmin_neon_generic_entry(q: &[f32], t: &[f32], ksub: usize, _scores:
&mut [f32]) -> u8 {
+ debug_assert!(q.len() >= 4);
+ debug_assert!(ksub.is_multiple_of(16));
+ // SAFETY: NEON is baseline on aarch64; slice bounds are checked by
+ // `score_argmin`.
+ unsafe { score_argmin_neon_generic(q, t, ksub) }
+}
+
+#[cfg(target_arch = "aarch64")]
+fn score_argmin_neon_d4_entry(q: &[f32], t: &[f32], ksub: usize, _scores: &mut
[f32]) -> u8 {
+ debug_assert_eq!(q.len(), 4);
+ debug_assert!(ksub.is_multiple_of(4));
+ // SAFETY: NEON is baseline on aarch64; slice bounds are checked by
+ // `score_argmin`.
+ unsafe { score_argmin_neon_d4(q, t, ksub) }
+}
+
+/// Generic NEON kernel: score 16 centroids together, then update the running
+/// SIMD argmin without materializing the score table.
+#[cfg(target_arch = "aarch64")]
+unsafe fn score_argmin_neon_generic(q: &[f32], t: &[f32], ksub: usize) -> u8 {
+ use std::arch::aarch64::*;
+
+ unsafe {
+ let mut min_val = vdupq_n_f32(f32::MAX);
+ let mut min_idx = vdupq_n_u32(0);
+ let lane0: [u32; 4] = [0, 1, 2, 3];
+ let mut cur_idx = vld1q_u32(lane0.as_ptr());
+ let four = vdupq_n_u32(4);
+
+ for j in (0..ksub).step_by(16) {
+ let q0 = vdupq_n_f32(q[0]);
+ let t0 = t.as_ptr().add(j);
+ let d0 = vsubq_f32(q0, vld1q_f32(t0));
+ let d1 = vsubq_f32(q0, vld1q_f32(t0.add(4)));
+ let d2 = vsubq_f32(q0, vld1q_f32(t0.add(8)));
+ let d3 = vsubq_f32(q0, vld1q_f32(t0.add(12)));
+ let mut s0 = vmulq_f32(d0, d0);
+ let mut s1 = vmulq_f32(d1, d1);
+ let mut s2 = vmulq_f32(d2, d2);
+ let mut s3 = vmulq_f32(d3, d3);
+
+ for (k, &qv) in q.iter().enumerate().skip(1) {
+ let qk = vdupq_n_f32(qv);
+ let tk = t.as_ptr().add(k * ksub + j);
+ let d0 = vsubq_f32(qk, vld1q_f32(tk));
+ let d1 = vsubq_f32(qk, vld1q_f32(tk.add(4)));
+ let d2 = vsubq_f32(qk, vld1q_f32(tk.add(8)));
+ let d3 = vsubq_f32(qk, vld1q_f32(tk.add(12)));
+ s0 = vfmaq_f32(s0, d0, d0);
+ s1 = vfmaq_f32(s1, d1, d1);
+ s2 = vfmaq_f32(s2, d2, d2);
+ s3 = vfmaq_f32(s3, d3, d3);
+ }
+
+ for score in [s0, s1, s2, s3] {
+ let mask = vcltq_f32(score, min_val);
+ min_val = vbslq_f32(mask, score, min_val);
+ min_idx = vbslq_u32(mask, cur_idx, min_idx);
+ cur_idx = vaddq_u32(cur_idx, four);
+ }
+ }
+
+ let mut vals = [0.0f32; 4];
+ let mut idxs = [0u32; 4];
+ vst1q_f32(vals.as_mut_ptr(), min_val);
+ vst1q_u32(idxs.as_mut_ptr(), min_idx);
+ let mut best = idxs[0];
+ let mut best_val = vals[0];
+ for lane in 1..4 {
+ if vals[lane] < best_val || (vals[lane] == best_val && idxs[lane]
< best) {
+ best_val = vals[lane];
+ best = idxs[lane];
+ }
+ }
+ best as u8
+ }
+}
+
+/// dsub=4 NEON kernel: 4 broadcast-FMA rows, SIMD min+index tracking,
+/// horizontal reduce with smallest-index tie-break.
+#[cfg(target_arch = "aarch64")]
+#[inline]
+unsafe fn score_argmin_neon_d4(q: &[f32], t: &[f32], ksub: usize) -> u8 {
+ use std::arch::aarch64::*;
+
+ let t0 = t.as_ptr();
+ let t1 = unsafe { t0.add(ksub) };
+ let t2 = unsafe { t0.add(2 * ksub) };
+ let t3 = unsafe { t0.add(3 * ksub) };
+
+ unsafe {
+ let q0 = vdupq_n_f32(q[0]);
+ let q1 = vdupq_n_f32(q[1]);
+ let q2 = vdupq_n_f32(q[2]);
+ let q3 = vdupq_n_f32(q[3]);
+
+ let mut min_val = vdupq_n_f32(f32::MAX);
+ let mut min_idx = vdupq_n_u32(0);
+ let lane0: [u32; 4] = [0, 1, 2, 3];
+ let mut cur_idx = vld1q_u32(lane0.as_ptr());
+ let step = vdupq_n_u32(4);
+
+ for j in (0..ksub).step_by(4) {
+ let d0 = vsubq_f32(q0, vld1q_f32(t0.add(j)));
+ let d1 = vsubq_f32(q1, vld1q_f32(t1.add(j)));
+ let d2 = vsubq_f32(q2, vld1q_f32(t2.add(j)));
+ let d3 = vsubq_f32(q3, vld1q_f32(t3.add(j)));
+ let mut s = vmulq_f32(d0, d0);
+ s = vfmaq_f32(s, d1, d1);
+ s = vfmaq_f32(s, d2, d2);
+ s = vfmaq_f32(s, d3, d3);
+
+ // Strictly-smaller keeps the earliest index on equal scores.
+ let mask = vcltq_f32(s, min_val);
+ min_val = vbslq_f32(mask, s, min_val);
+ min_idx = vbslq_u32(mask, cur_idx, min_idx);
+ cur_idx = vaddq_u32(cur_idx, step);
+ }
+
+ let mut vals = [0.0f32; 4];
+ let mut idxs = [0u32; 4];
+ vst1q_f32(vals.as_mut_ptr(), min_val);
+ vst1q_u32(idxs.as_mut_ptr(), min_idx);
+ let mut best = idxs[0];
+ let mut best_val = vals[0];
+ for l in 1..4 {
+ if vals[l] < best_val || (vals[l] == best_val && idxs[l] < best) {
+ best_val = vals[l];
+ best = idxs[l];
+ }
+ }
+ best as u8
+ }
+}
+
+/// dsub=4 AVX2 kernel: mirrors the NEON version 8-wide.
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2", enable = "fma")]
+unsafe fn score_argmin_avx2_d4(q: &[f32], t: &[f32], ksub: usize) -> u8 {
+ use std::arch::x86_64::*;
+
+ let t0 = t.as_ptr();
+ let t1 = unsafe { t0.add(ksub) };
+ let t2 = unsafe { t0.add(2 * ksub) };
+ let t3 = unsafe { t0.add(3 * ksub) };
+
+ unsafe {
+ let q0 = _mm256_set1_ps(q[0]);
+ let q1 = _mm256_set1_ps(q[1]);
+ let q2 = _mm256_set1_ps(q[2]);
+ let q3 = _mm256_set1_ps(q[3]);
+
+ let mut min_val = _mm256_set1_ps(f32::MAX);
+ let mut min_idx = _mm256_setzero_si256();
+ let mut cur_idx = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+ let step = _mm256_set1_epi32(8);
+
+ for j in (0..ksub).step_by(8) {
+ let d0 = _mm256_sub_ps(q0, _mm256_loadu_ps(t0.add(j)));
+ let d1 = _mm256_sub_ps(q1, _mm256_loadu_ps(t1.add(j)));
+ let d2 = _mm256_sub_ps(q2, _mm256_loadu_ps(t2.add(j)));
+ let d3 = _mm256_sub_ps(q3, _mm256_loadu_ps(t3.add(j)));
+ let mut s = _mm256_mul_ps(d0, d0);
+ s = _mm256_fmadd_ps(d1, d1, s);
+ s = _mm256_fmadd_ps(d2, d2, s);
+ s = _mm256_fmadd_ps(d3, d3, s);
+
+ // Strictly-smaller keeps the earliest index on equal scores.
+ let mask = _mm256_cmp_ps::<_CMP_LT_OQ>(s, min_val);
+ min_val = _mm256_blendv_ps(min_val, s, mask);
+ min_idx = _mm256_blendv_epi8(min_idx, cur_idx,
_mm256_castps_si256(mask));
+ cur_idx = _mm256_add_epi32(cur_idx, step);
+ }
+
+ let mut vals = [0.0f32; 8];
+ let mut idxs = [0i32; 8];
+ _mm256_storeu_ps(vals.as_mut_ptr(), min_val);
+ _mm256_storeu_si256(idxs.as_mut_ptr().cast(), min_idx);
+ let mut best = idxs[0] as u32;
+ let mut best_val = vals[0];
+ for l in 1..8 {
+ let idx = idxs[l] as u32;
+ if vals[l] < best_val || (vals[l] == best_val && idx < best) {
+ best_val = vals[l];
+ best = idx;
+ }
+ }
+ best as u8
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -729,11 +1196,13 @@ mod tests {
pq.train(&data, n);
let cs = pq.code_size(); // m/2 = 4
- let mut codes = vec![0u8; n * cs];
- pq.encode_batch(&data, n, &mut codes);
+ let mut canonical = vec![0u8; n * cs];
+ pq.encode_batch(&data, n, &mut canonical);
+ let mut blocked = vec![0u8; n * cs];
+ pq.encode_batch_blocked(&data, n, &mut blocked);
- // Verify codes are non-trivial (not all zeros)
- assert!(codes.iter().any(|&b| b != 0));
+ assert_eq!(blocked, canonical);
+ assert!(blocked.iter().any(|&b| b != 0));
}
#[test]
@@ -764,116 +1233,444 @@ mod tests {
assert!((table_distance - decoded_distance).abs() < 1e-4);
}
- /// Reference per-vector encode used to pin the blocked batch path.
- fn encode_per_vector(pq: &ProductQuantizer, data: &[f32], n: usize) ->
Vec<u8> {
+ /// Scalar direct-L2 oracle for the transposed path: same transposed
+ /// table, always the scalar kernel, no threads.
+ fn encode_scalar_oracle(pq: &ProductQuantizer, data: &[f32], n: usize) ->
Vec<u8> {
let cs = pq.code_size();
+ assert_eq!(cs, pq.m);
+ let (transposed, sub_stride) = pq.build_transposed_codebook();
+ let mut scores = vec![0.0f32; pq.ksub];
let mut codes = vec![0u8; n * cs];
- for i in 0..n {
- pq.encode(
- &data[i * pq.d..(i + 1) * pq.d],
- &mut codes[i * cs..(i + 1) * cs],
- );
+ for r in 0..n {
+ let row = &data[r * pq.d..(r + 1) * pq.d];
+ for sub in 0..pq.m {
+ let range = pq.chunk_range(sub);
+ let dsub = range.len();
+ let t = &transposed[sub * sub_stride..sub * sub_stride + dsub
* pq.ksub];
+ codes[r * cs + sub] =
+ score_argmin(score_argmin_scalar, &row[range], t, pq.ksub,
&mut scores);
+ }
}
codes
}
- #[test]
- fn test_encode_batch_blocked_matches_per_vector() {
- let d = 32;
- let m = 8; // dsub = 4: hits the SIMD kernels
- let mut rng = StdRng::seed_from_u64(20260820);
+ const BATCH_SIZES: [usize; 11] = [1, 2, 7, 31, 32, 33, 511, 512, 513,
2048, 4096];
+
+ fn trained_pq(d: usize, m: usize, seed: u64) -> (ProductQuantizer, StdRng)
{
+ let mut rng = StdRng::seed_from_u64(seed);
let train: Vec<f32> = (0..3000 * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
let mut pq = ProductQuantizer::new(d, m);
pq.train(&train, 3000);
+ (pq, rng)
+ }
- // Below, exactly at, above, and misaligned against the block size.
- for n in [
- 1,
- 31,
- 32,
- 33,
- MAX_ENCODE_BLOCK_ROWS,
- MAX_ENCODE_BLOCK_ROWS + 7,
- 2048,
- ] {
- let data: Vec<f32> = (0..n * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
- let reference = encode_per_vector(&pq, &data, n);
- let mut batch = vec![0u8; n * pq.code_size()];
- pq.encode_batch_blocked(&data, n, &mut batch);
- assert_eq!(batch, reference, "n={n}");
+ #[test]
+ fn test_encode_batch_8bit_transposed_matches_scalar_oracle() {
+ // dsub = 4 (d4 kernel), 5 and 8 (generic kernel).
+ for (d, m) in [(32, 8), (25, 5), (40, 5)] {
+ let (pq, mut rng) = trained_pq(d, m, 20260820 + d as u64);
+ for n in BATCH_SIZES {
+ let data: Vec<f32> = (0..n * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
+ let oracle = encode_scalar_oracle(&pq, &data, n);
+ let mut batch = vec![0u8; n * pq.code_size()];
+ pq.encode_batch_8bit_transposed(&data, n, &mut batch);
+ assert_eq!(batch, oracle, "d={d} m={m} n={n}");
+ }
}
}
#[test]
- fn test_encode_batch_blocked_non_uniform_chunks() {
+ fn test_centroids_finite_cache_updates_and_clones() {
+ let mut pq = ProductQuantizer::new(4, 1);
+ for value in [0.0, f32::NAN, f32::INFINITY, f32::NEG_INFINITY, 1.0] {
+ let mut centroids = vec![0.0; 4 * 256];
+ centroids[4 * 256 - 1] = value;
+ pq.set_centroids(centroids);
+ assert_eq!(pq.centroids_are_finite, value.is_finite());
+ let cloned = pq.clone_without_transposed_cache();
+ assert_eq!(cloned.centroids_are_finite, value.is_finite());
+ if !supports_transposed_pq_encoding() {
+ let data = [1.0; 4];
+ let mut expected = [0];
+ if value.is_finite() {
+ pq.encode_batch_8bit_sgemm(&data, 1, &mut expected);
+ } else {
+ pq.encode_batch(&data, 1, &mut expected);
+ }
+ for encoder in [&pq, &cloned] {
+ let mut actual = [0];
+ encoder.encode_batch_blocked(&data, 1, &mut actual);
+ assert_eq!(actual, expected);
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn test_transposed_codebook_cache_reuses_and_invalidates() {
+ let (mut pq, mut rng) = trained_pq(32, 8, 20260908);
+ let data: Vec<f32> = (0..32).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
+ let mut expected = vec![0u8; pq.code_size()];
+
+ assert!(pq.transposed_codebook_cache.get().is_none());
+ pq.encode_batch_8bit_transposed(&data, 1, &mut expected);
+ let cached = pq.transposed_codebook_cache.get().unwrap().0.as_ptr();
+
+ let mut actual = vec![0u8; pq.code_size()];
+ pq.encode_batch_8bit_transposed(&data, 1, &mut actual);
+ assert_eq!(actual, expected);
+ assert_eq!(
+ pq.transposed_codebook_cache.get().unwrap().0.as_ptr(),
+ cached
+ );
+
+ let centroids = pq.centroids().to_vec();
+ pq.set_centroids(centroids);
+ assert!(pq.transposed_codebook_cache.get().is_none());
+ pq.encode_batch_8bit_transposed(&data, 1, &mut actual);
+ assert_eq!(actual, expected);
+ assert!(pq.transposed_codebook_cache.get().is_some());
+ }
+
+ #[test]
+ fn test_encode_batch_8bit_transposed_non_uniform_chunks() {
let d = 13;
let m = 3;
let mut rng = StdRng::seed_from_u64(20260821);
let train: Vec<f32> = (0..2000 * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
let mut pq = ProductQuantizer::with_nbits_balanced(d, m, 8);
pq.train(&train, 2000);
+ // Chunks of 5, 4 and 4: generic kernel next to the d4 kernel.
assert_eq!(pq.chunk_offsets, vec![0, 5, 9, 13]);
let n = MAX_ENCODE_BLOCK_ROWS + 13;
let data: Vec<f32> = (0..n * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
- let reference = encode_per_vector(&pq, &data, n);
+ let oracle = encode_scalar_oracle(&pq, &data, n);
let mut batch = vec![0u8; n * pq.code_size()];
- pq.encode_batch_blocked(&data, n, &mut batch);
- assert_eq!(batch, reference);
+ pq.encode_batch_8bit_transposed(&data, n, &mut batch);
+ assert_eq!(batch, oracle);
}
#[test]
- fn test_encode_batch_blocked_matches_canonical_large_offset() {
+ fn test_encode_batch_8bit_transposed_large_offset_is_correct() {
+ // The expanded form cancels here; the direct form must still pick the
+ // exact centroid, for every batch size.
let mut pq = ProductQuantizer::new(4, 1);
- pq.centroids = vec![100_000_016.0; pq.d * pq.ksub];
- pq.centroids[0..4].fill(100_000_008.0);
- pq.centroids[4..8].fill(100_000_000.0);
-
- let data = vec![100_000_000.0; 32 * pq.d];
- let mut canonical = vec![0; 32];
- pq.encode_batch(&data, 32, &mut canonical);
- let mut blocked = vec![0; 32];
- pq.encode_batch_blocked(&data, 32, &mut blocked);
+ let mut centroids = vec![100_000_016.0; pq.d * pq.ksub];
+ centroids[0..4].fill(100_000_008.0);
+ centroids[4..8].fill(100_000_000.0);
+ pq.set_centroids(centroids);
+
+ for n in BATCH_SIZES {
+ let data = vec![100_000_000.0; n * pq.d];
+ let mut codes = vec![0; n];
+ pq.encode_batch_8bit_transposed(&data, n, &mut codes);
+ assert!(codes.iter().all(|&code| code == 1), "n={n}");
+ }
+ }
- assert_eq!(blocked, canonical);
+ #[test]
+ fn test_encode_batch_8bit_transposed_is_batch_invariant() {
+ let mut pq = ProductQuantizer::new(4, 1);
+ let mut centroids = vec![100.0; pq.d * pq.ksub];
+ centroids[0..4].copy_from_slice(&[0.3658799, 0.06077051, -0.46501994,
-0.31766486]);
+ centroids[4..8].copy_from_slice(&[0.3658799, 0.06077051, -0.46501994,
-0.31766483]);
+ pq.set_centroids(centroids);
+
+ let max_n = *BATCH_SIZES.last().unwrap();
+ let mut largest = vec![0; max_n];
+ pq.encode_batch_8bit_transposed(&vec![0.0; max_n * 4], max_n, &mut
largest);
+ for n in BATCH_SIZES {
+ let mut codes = vec![0; n];
+ pq.encode_batch_8bit_transposed(&vec![0.0; n * 4], n, &mut codes);
+ assert_eq!(codes.as_slice(), &largest[..n], "n={n}");
+ }
}
#[test]
- fn test_encode_batch_blocked_is_batch_invariant() {
+ fn
test_encode_batch_8bit_transposed_non_finite_semantics_are_batch_invariant() {
let mut pq = ProductQuantizer::new(4, 1);
- pq.centroids = vec![100.0; pq.d * pq.ksub];
- pq.centroids[0..4].copy_from_slice(&[0.3658799, 0.06077051,
-0.46501994, -0.31766486]);
- pq.centroids[4..8].copy_from_slice(&[0.3658799, 0.06077051,
-0.46501994, -0.31766483]);
+ let mut centroids = vec![1.0; pq.d * pq.ksub];
+ centroids[0..4].fill(f32::NAN);
+ centroids[4..8].fill(0.0);
+ centroids[8..12].fill(f32::INFINITY);
+ centroids[12..16].fill(f32::NEG_INFINITY);
+ pq.set_centroids(centroids);
+
+ for n in BATCH_SIZES {
+ let mut codes = vec![0; n];
+ pq.encode_batch_8bit_transposed(&vec![0.0; n * pq.d], n, &mut
codes);
+ assert!(codes.iter().all(|&code| code == 1), "n={n}");
+ }
- let mut small = vec![0; 31];
- pq.encode_batch_blocked(&[0.0; 31 * 4], 31, &mut small);
- let mut large = vec![0; 32];
- pq.encode_batch_blocked(&[0.0; 32 * 4], 32, &mut large);
+ pq.set_centroids(vec![f32::NAN; pq.d * pq.ksub]);
+ for n in BATCH_SIZES {
+ let mut codes = vec![u8::MAX; n];
+ pq.encode_batch_8bit_transposed(&vec![0.0; n * pq.d], n, &mut
codes);
+ assert!(codes.iter().all(|&code| code == 0), "n={n}");
+ }
+ }
- assert_eq!(small.as_slice(), &large[..31]);
+ #[test]
+ fn test_encode_batch_blocked_small_dsub_uses_canonical_path() {
+ for dsub in [1, 2, 3] {
+ let m = 4;
+ let d = m * dsub;
+ let (pq, mut rng) = trained_pq(d, m, 20260903 + dsub as u64);
+ for n in [1, 33, 513] {
+ let data: Vec<f32> = (0..n * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
+ let mut canonical = vec![0u8; n * m];
+ pq.encode_batch(&data, n, &mut canonical);
+ let mut blocked = vec![0u8; n * m];
+ pq.encode_batch_blocked(&data, n, &mut blocked);
+ assert_eq!(blocked, canonical, "dsub={dsub} n={n}");
+ }
+ }
}
#[test]
- fn test_encode_batch_blocked_nan_centroid_is_batch_invariant() {
- let mut pq = ProductQuantizer::new(4, 1);
- pq.centroids = vec![1.0; pq.d * pq.ksub];
- pq.centroids[0] = f32::NAN;
- pq.centroids[4..8].fill(0.0);
+ fn test_encode_batch_8bit_transposed_is_thread_and_split_invariant() {
+ let d = 32;
+ let m = 8;
+ let (pq, mut rng) = trained_pq(d, m, 20260904);
+ let n = 3 * MAX_ENCODE_BLOCK_ROWS + 5;
+ let data: Vec<f32> = (0..n * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
+ let mut whole = vec![0u8; n * m];
+ pq.encode_batch_8bit_transposed(&data, n, &mut whole);
- let mut small = vec![0; 1];
- pq.encode_batch_blocked(&[0.0; 4], 1, &mut small);
- let mut large = vec![0; 32];
- pq.encode_batch_blocked(&[0.0; 32 * 4], 32, &mut large);
+ for threads in [1, 2, 11, 16] {
+ let pool = rayon::ThreadPoolBuilder::new()
+ .num_threads(threads)
+ .build()
+ .unwrap();
+ let mut codes = vec![0u8; n * m];
+ pool.install(|| pq.encode_batch_8bit_transposed(&data, n, &mut
codes));
+ assert_eq!(codes, whole, "threads={threads}");
+ }
- assert_eq!(small[0], large[0]);
+ for batch in [1, 7, MAX_ENCODE_BLOCK_ROWS] {
+ let mut codes = vec![0u8; n * m];
+ for start in (0..n).step_by(batch) {
+ let rows = batch.min(n - start);
+ pq.encode_batch_8bit_transposed(
+ &data[start * d..(start + rows) * d],
+ rows,
+ &mut codes[start * m..(start + rows) * m],
+ );
+ }
+ assert_eq!(codes, whole, "batch={batch}");
+ }
}
#[test]
- fn test_encode_block_rows_avoids_tiny_sgemm_blocks() {
+ fn test_encode_batch_8bit_sgemm_is_thread_and_split_invariant() {
+ let d = 32;
+ let m = 8;
+ let (pq, mut rng) = trained_pq(d, m, 20260905);
+ let n = 3 * MAX_ENCODE_BLOCK_ROWS + 5;
+ let data: Vec<f32> = (0..n * d).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
+ let mut whole = vec![0u8; n * m];
+ pq.encode_batch_8bit_sgemm(&data, n, &mut whole);
+
+ for threads in [1, 2, 11, 16] {
+ let pool = rayon::ThreadPoolBuilder::new()
+ .num_threads(threads)
+ .build()
+ .unwrap();
+ let mut codes = vec![0u8; n * m];
+ pool.install(|| pq.encode_batch_8bit_sgemm(&data, n, &mut codes));
+ assert_eq!(codes, whole, "threads={threads}");
+ }
+
+ for batch in [1, 7, MAX_ENCODE_BLOCK_ROWS] {
+ let mut codes = vec![0u8; n * m];
+ for start in (0..n).step_by(batch) {
+ let rows = batch.min(n - start);
+ pq.encode_batch_8bit_sgemm(
+ &data[start * d..(start + rows) * d],
+ rows,
+ &mut codes[start * m..(start + rows) * m],
+ );
+ }
+ assert_eq!(codes, whole, "batch={batch}");
+ }
+ }
+
+ #[test]
+ fn test_encode_block_rows_clamps() {
assert_eq!(encode_block_rows(2730, 12), 228);
assert_eq!(encode_block_rows(32768, 12), MAX_ENCODE_BLOCK_ROWS);
- assert_eq!(encode_sgemm_min_rows(8), 32);
- assert_eq!(encode_sgemm_min_rows(32), 128);
+ assert_eq!(encode_block_rows(5, 16), 1);
+ assert_eq!(encode_block_rows(0, 16), 1);
+ assert_eq!(encode_block_rows(100, 0), 100);
+ }
+
+ /// Every kernel `score_argmin_kernel` can hand out on this machine, plus
+ /// the scalar oracle.
+ fn kernels_under_test(dsub: usize, ksub: usize) -> Vec<(&'static str,
ScoreArgminKernel)> {
+ let mut kernels: Vec<(&'static str, ScoreArgminKernel)> = vec![
+ ("scalar", score_argmin_scalar),
+ ("dispatched", score_argmin_kernel(dsub, ksub)),
+ ];
+ #[cfg(target_arch = "x86_64")]
+ {
+ if is_x86_feature_detected!("avx2") &&
is_x86_feature_detected!("fma") {
+ kernels.push(("avx2_generic",
score_argmin_avx2_generic_entry));
+ if dsub == 4 && ksub.is_multiple_of(8) {
+ kernels.push(("avx2_d4", score_argmin_avx2_d4_entry));
+ }
+ }
+ }
+ #[cfg(target_arch = "aarch64")]
+ {
+ if dsub >= 4 && ksub.is_multiple_of(16) {
+ kernels.push(("neon_generic",
score_argmin_neon_generic_entry));
+ }
+ if dsub == 4 && ksub.is_multiple_of(4) {
+ kernels.push(("neon_d4", score_argmin_neon_d4_entry));
+ }
+ }
+ kernels
+ }
+
+ #[test]
+ fn test_score_argmin_kernels_match_scalar_oracle() {
+ let ksub = 256;
+ let mut rng = StdRng::seed_from_u64(20260905);
+ for dsub in [4, 5, 8, 12, 16] {
+ let mut t: Vec<f32> = (0..dsub * ksub)
+ .map(|_| rng.gen_range(-1.0f32..1.0))
+ .collect();
+ // Inject exact duplicates so some queries hit exact ties.
+ for j in (0..ksub).step_by(37) {
+ for k in 0..dsub {
+ t[k * ksub + (j + 1) % ksub] = t[k * ksub + j];
+ }
+ }
+ let mut scores = vec![0.0f32; ksub];
+ for _ in 0..500 {
+ let q: Vec<f32> = (0..dsub).map(|_|
rng.gen_range(-1.0f32..1.0)).collect();
+ let expected = score_argmin(score_argmin_scalar, &q, &t, ksub,
&mut scores);
+ for (name, kernel) in kernels_under_test(dsub, ksub) {
+ let got = score_argmin(kernel, &q, &t, ksub, &mut scores);
+ assert_eq!(got, expected, "kernel {name} dsub {dsub}");
+ }
+ }
+ // Queries equal to a centroid resolve to its smallest duplicate.
+ for j in (0..ksub).step_by(37) {
+ let q: Vec<f32> = (0..dsub).map(|k| t[k * ksub + j]).collect();
+ for (name, kernel) in kernels_under_test(dsub, ksub) {
+ assert_eq!(
+ score_argmin(kernel, &q, &t, ksub, &mut scores) as
usize,
+ j,
+ "kernel {name} dsub {dsub} j {j}"
+ );
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn test_score_argmin_non_finite_semantics() {
+ let ksub = 256;
+ for dsub in [4, 5] {
+ let q = vec![0.0f32; dsub];
+ let mut t = vec![1.0f32; dsub * ksub];
+ for k in 0..dsub {
+ t[k * ksub] = f32::NAN;
+ t[k * ksub + 1] = 0.0;
+ t[k * ksub + 2] = f32::INFINITY;
+ t[k * ksub + 3] = f32::NEG_INFINITY;
+ }
+ let mut scores = vec![0.0f32; ksub];
+ for (name, kernel) in kernels_under_test(dsub, ksub) {
+ assert_eq!(
+ score_argmin(kernel, &q, &t, ksub, &mut scores),
+ 1,
+ "kernel {name} dsub {dsub}"
+ );
+ }
+
+ t.fill(f32::NAN);
+ for (name, kernel) in kernels_under_test(dsub, ksub) {
+ assert_eq!(
+ score_argmin(kernel, &q, &t, ksub, &mut scores),
+ 0,
+ "all NaN: kernel {name} dsub {dsub}"
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn test_score_argmin_tie_prefers_smallest_index() {
+ let ksub = 256;
+ let dsub = 4;
+ let q = [0.25f32, -0.5, 0.75, -0.125];
+ let mut scores = vec![0.0f32; ksub];
+
+ // All centroids identical -> every score ties -> index 0 wins.
+ let all_equal = vec![0.5f32; dsub * ksub];
+ for (name, kernel) in kernels_under_test(dsub, ksub) {
+ assert_eq!(
+ score_argmin(kernel, &q, &all_equal, ksub, &mut scores),
+ 0,
+ "{name}"
+ );
+ }
+
+ // Exact ties at indices 1 and 8 (different AVX2 lanes) and at 8 and 9
+ // (the same lane group): the smallest index must win each time.
+ for (tie_a, tie_b) in [(1usize, 8usize), (8, 9), (7, 8), (255, 1)] {
+ let mut t = vec![3.0f32; dsub * ksub];
+ for k in 0..dsub {
+ t[k * ksub + tie_a] = q[k];
+ t[k * ksub + tie_b] = q[k];
+ }
+ let expected = tie_a.min(tie_b);
+ for (name, kernel) in kernels_under_test(dsub, ksub) {
+ assert_eq!(
+ score_argmin(kernel, &q, &t, ksub, &mut scores) as usize,
+ expected,
+ "{name} tie ({tie_a}, {tie_b})"
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn test_score_argmin_scalar_matches_squared_distance() {
+ let q = [0.4f32, -0.7, 0.2, 0.9, -0.3];
+ let ksub = 17;
+ let t: Vec<f32> = (0..q.len() * ksub)
+ .map(|i| ((i * 13 % 29) as f32 - 14.0) / 7.0)
+ .collect();
+ let distances: Vec<f32> = (0..ksub)
+ .map(|j| {
+ q.iter()
+ .enumerate()
+ .map(|(k, value)| (value - t[k * ksub + j]).powi(2))
+ .sum()
+ })
+ .collect();
+ let mut scores = vec![0.0; ksub];
+ assert_eq!(
+ score_argmin(score_argmin_scalar, &q, &t, ksub, &mut scores),
+ argmin_code(&distances)
+ );
+ }
+
+ #[test]
+ #[should_panic]
+ fn test_score_argmin_rejects_short_transposed_codebook() {
+ let mut scores = vec![0.0; 256];
+ score_argmin(
+ score_argmin_scalar,
+ &[0.0; 4],
+ &[0.0; 4 * 256 - 1],
+ 256,
+ &mut scores,
+ );
}
#[test]
@@ -907,8 +1704,7 @@ mod tests {
let m = 1;
let ksub = 256;
let mut pq = ProductQuantizer::new(d, m);
- pq.centroids = vec![0.0; ksub * d];
-
+ let mut centroids = vec![0.0; ksub * d];
let mut query = vec![0.0; d];
for i in 0..d {
let value = if i.is_multiple_of(2) {
@@ -917,9 +1713,9 @@ mod tests {
-1.0e10_f32 + i as f32
};
query[i] = value;
- pq.centroids[i] = value;
+ centroids[i] = value;
}
- pq.rebuild_norms_cache();
+ pq.set_centroids(centroids);
let mut table = vec![0.0; m * ksub];
pq.compute_distance_table(&query, MetricType::L2, &mut table);
diff --git a/core/src/vamana.rs b/core/src/vamana.rs
index fe037c7..f64ec30 100644
--- a/core/src/vamana.rs
+++ b/core/src/vamana.rs
@@ -346,8 +346,8 @@ impl VamanaGraph {
params: DiskAnnBuildParams,
) -> io::Result<(Self, VamanaBuildStats)> {
validate_build_inputs(vectors, count, dimension, params)?;
- if pq.d != dimension
- || !matches!(pq.nbits, 4 | 8)
+ if pq.d() != dimension
+ || !matches!(pq.nbits(), 4 | 8)
|| !pq.has_valid_layout()
|| pq_codes.len() != count.saturating_mul(pq.code_size())
{
@@ -427,8 +427,8 @@ impl VamanaGraph {
));
}
if let Some((pq, pq_codes)) = pq_build {
- if pq.d != dimension
- || !matches!(pq.nbits, 4 | 8)
+ if pq.d() != dimension
+ || !matches!(pq.nbits(), 4 | 8)
|| !pq.has_valid_layout()
|| pq_codes.len() != count.saturating_mul(pq.code_size())
{
@@ -1274,30 +1274,30 @@ impl<'a> PqBuildDistance<'a> {
metric: MetricType,
) -> io::Result<Self> {
let table_len = pq
- .m
- .checked_mul(pq.ksub)
- .and_then(|value| value.checked_mul(pq.ksub))
+ .m()
+ .checked_mul(pq.ksub())
+ .and_then(|value| value.checked_mul(pq.ksub()))
.ok_or_else(|| invalid_input("Vamana PQ build-distance table size
overflows usize"))?;
let mut centroid_distances = Vec::new();
centroid_distances
.try_reserve_exact(table_len)
.map_err(|_| invalid_input("Vamana PQ build-distance table
allocation failed"))?;
centroid_distances.resize(table_len, 0.0);
- for sub in 0..pq.m {
+ for sub in 0..pq.m() {
let chunk_dim = pq.chunk_dim(sub);
let sub_base = pq.centroid_chunk_base(sub);
- let table_base = sub * pq.ksub * pq.ksub;
- for left in 0..pq.ksub {
+ let table_base = sub * pq.ksub() * pq.ksub();
+ for left in 0..pq.ksub() {
let left_start = sub_base + left * chunk_dim;
- for right in left..pq.ksub {
+ for right in left..pq.ksub() {
let right_start = sub_base + right * chunk_dim;
let distance = fvec_distance(
- &pq.centroids[left_start..left_start + chunk_dim],
- &pq.centroids[right_start..right_start + chunk_dim],
+ &pq.centroids()[left_start..left_start + chunk_dim],
+ &pq.centroids()[right_start..right_start + chunk_dim],
metric,
);
- centroid_distances[table_base + left * pq.ksub + right] =
distance;
- centroid_distances[table_base + right * pq.ksub + left] =
distance;
+ centroid_distances[table_base + left * pq.ksub() + right]
= distance;
+ centroid_distances[table_base + right * pq.ksub() + left]
= distance;
}
}
}
@@ -1310,9 +1310,9 @@ impl<'a> PqBuildDistance<'a> {
Ok(Self {
codes,
code_size,
- m: pq.m,
- ksub: pq.ksub,
- nbits: pq.nbits,
+ m: pq.m(),
+ ksub: pq.ksub(),
+ nbits: pq.nbits(),
centroid_distances,
})
}
@@ -2498,13 +2498,15 @@ mod tests {
#[test]
fn
vamana_pq_build_distance_matches_decoded_centroid_distance_for_4bit_codes() {
let mut pq = ProductQuantizer::with_nbits(4, 2, 4);
- pq.centroids = (0..pq.d * pq.ksub)
- .map(|index| index as f32 * 0.125)
- .collect();
+ pq.set_centroids(
+ (0..pq.d() * pq.ksub())
+ .map(|index| index as f32 * 0.125)
+ .collect(),
+ );
let codes = [0x21, 0x43];
let distance = PqBuildDistance::new(&pq, &codes, 2,
MetricType::L2).unwrap();
- let mut left = vec![0.0; pq.d];
- let mut right = vec![0.0; pq.d];
+ let mut left = vec![0.0; pq.d()];
+ let mut right = vec![0.0; pq.d()];
pq.decode(&codes[..1], &mut left);
pq.decode(&codes[1..], &mut right);
@@ -2516,13 +2518,15 @@ mod tests {
#[test]
fn vamana_pq_build_distance_and_pruning_follow_inner_product_semantics() {
let mut pq = ProductQuantizer::with_nbits(2, 1, 4);
- pq.centroids = (0..pq.d * pq.ksub)
- .map(|index| index as f32 * 0.25 - 1.0)
- .collect();
+ pq.set_centroids(
+ (0..pq.d() * pq.ksub())
+ .map(|index| index as f32 * 0.25 - 1.0)
+ .collect(),
+ );
let codes = [0x01, 0x03];
let distance = PqBuildDistance::new(&pq, &codes, 2,
MetricType::InnerProduct).unwrap();
- let mut left = vec![0.0; pq.d];
- let mut right = vec![0.0; pq.d];
+ let mut left = vec![0.0; pq.d()];
+ let mut right = vec![0.0; pq.d()];
pq.decode(&codes[..1], &mut left);
pq.decode(&codes[1..], &mut right);
diff --git a/core/tests/storage_format_fixtures.rs
b/core/tests/storage_format_fixtures.rs
index df3d0e4..f2af87f 100644
--- a/core/tests/storage_format_fixtures.rs
+++ b/core/tests/storage_format_fixtures.rs
@@ -350,8 +350,9 @@ fn build_diskann_fixture() -> Vec<u8> {
..DiskAnnBuildParams::default()
},
);
- index.pq.centroids = (0..256).map(|code| code as f32 * 0.25).collect();
- index.pq.rebuild_norms_cache();
+ index
+ .pq
+ .set_centroids((0..256).map(|code| code as f32 * 0.25).collect());
index.ids = vec![7];
index.vectors = vec![0.0];
write_diskann_fixture(index)
@@ -375,10 +376,13 @@ fn build_diskann_compact_multipage_fixture() -> Vec<u8> {
..DiskAnnBuildParams::default()
},
);
- index.pq.centroids = (0..dimension)
- .flat_map(|coordinate| (0..256).map(move |code| code as f32 +
coordinate as f32 * 0.001))
- .collect();
- index.pq.rebuild_norms_cache();
+ index.pq.set_centroids(
+ (0..dimension)
+ .flat_map(|coordinate| {
+ (0..256).map(move |code| code as f32 + coordinate as f32 *
0.001)
+ })
+ .collect(),
+ );
index.ids = (0..count).map(|node| 10_000 + node as i64 * 7).collect();
index.vectors = (0..count)
.flat_map(|node| {
@@ -423,12 +427,13 @@ fn build_diskann_interleaved_4bit_fixture() -> Vec<u8> {
..DiskAnnBuildParams::default()
},
);
- index.pq.centroids = (0..dimension)
- .flat_map(|coordinate| {
- (0..16).map(move |code| code as f32 * 0.5 + coordinate as f32 *
0.001)
- })
- .collect();
- index.pq.rebuild_norms_cache();
+ index.pq.set_centroids(
+ (0..dimension)
+ .flat_map(|coordinate| {
+ (0..16).map(move |code| code as f32 * 0.5 + coordinate as f32
* 0.001)
+ })
+ .collect(),
+ );
index.ids = (0..count).map(|node| -500 + node as i64 * 11).collect();
index.vectors = (0..count)
.flat_map(|node| {
@@ -454,8 +459,9 @@ fn build_diskann_raw_row_ids_fixture() -> Vec<u8> {
..DiskAnnBuildParams::default()
},
);
- index.pq.centroids = (0..256).map(|code| code as f32).collect();
- index.pq.rebuild_norms_cache();
+ index
+ .pq
+ .set_centroids((0..256).map(|code| code as f32).collect());
index.ids = vec![i64::MIN, 0, i64::MAX];
index.vectors = vec![0.0, 1.0, 2.0];
write_diskann_fixture(index)
@@ -488,8 +494,11 @@ fn build_ivf_flat_fixture() -> Vec<u8> {
fn build_ivf_pq_fixture() -> Vec<u8> {
let mut index = IVFPQIndex::new(1, 2, 1, MetricType::L2, false);
index.set_quantizer_centroids(vec![0.0, 10.0]);
- index.pq.centroids = (0..index.pq.ksub).map(|code| code as f32 *
0.25).collect();
- index.pq.rebuild_norms_cache();
+ index.pq.set_centroids(
+ (0..index.pq.ksub())
+ .map(|code| code as f32 * 0.25)
+ .collect(),
+ );
index.ids = vec![vec![20, 10], vec![30]];
index.codes = vec![vec![1, 0], vec![0]];
@@ -501,10 +510,11 @@ fn build_ivf_pq_fixture() -> Vec<u8> {
fn build_ivf_pq_4bit_fixture() -> Vec<u8> {
let mut index = IVFPQIndex::with_nbits(2, 2, 2, 4, MetricType::L2, false);
index.set_quantizer_centroids(vec![0.0, 0.0, 10.0, 10.0]);
- index.pq.centroids = (0..index.pq.m)
- .flat_map(|_| (0..index.pq.ksub).map(|code| code as f32 * 0.5))
- .collect();
- index.pq.rebuild_norms_cache();
+ index.pq.set_centroids(
+ (0..index.pq.m())
+ .flat_map(|_| (0..index.pq.ksub()).map(|code| code as f32 * 0.5))
+ .collect(),
+ );
index.ids = vec![vec![8, 5], vec![30]];
index.codes = vec![vec![0x11, 0x00], vec![0x00]];
diff --git a/docs/api.html b/docs/api.html
index b1e004d..79bedf8 100644
--- a/docs/api.html
+++ b/docs/api.html
@@ -50,6 +50,7 @@
<div class="flow" aria-label="Unified API lifecycle"><div
class="flow-step"><small>01</small><strong>Create a Trainer<br>Parse and
validate options</strong></div><div
class="flow-step"><small>02</small><strong>Submit one or more<br>training
batches</strong></div><div class="flow-step"><small>03</small><strong>Finish
training and<br>create a one-shot Writer</strong></div><div
class="flow-step"><small>04</small><strong>Add row IDs / vectors<br>and write
the file</strong></div><di [...]
<ul><li>Vectors are contiguous <code>f32</code> values; length must
equal <code>vector_count × dimension</code>.</li><li>Training data may arrive
in batches. Every IVF trainer keeps a deterministic reservoir of at most
<code>max(65,536, 64 × resolved nlist)</code> vectors. DiskANN starts from a
50,000-row cap and lowers it when necessary so the retained sample, optional
cosine-normalized copy, codebook, and parallel PQ-training scratch fit
<code>diskann.memory-budget-bytes</cod [...]
<div class="callout warning"><strong>IVF coarse assignment is
approximate by default for large centroid matrices</strong>When <code>dimension
× nlist ≥ 1,000,000</code>, <code>ivf.coarse-assignment=auto</code> uses a
Vamana graph while training and adding vectors. Search still selects lists by
exact centroid distance, so graph assignment can lower recall at small
<code>nprobe</code> and does not guarantee that a vector is found by a
self-query with <code>nprobe=1</code>. Set <c [...]
+ <div class="callout warning"><strong>IVF-PQ encoding defaults to
auto</strong>For 8-bit PQ (<code>ksub=256</code>) where every subvector has at
least four dimensions, <code>ivf.pq-encoding=auto</code> uses the transposed
direct-L2 encoder on x86 with AVX2+FMA and on AArch64. Other CPUs use blocked
SGEMM for finite codebooks. Unsupported shapes, and non-finite codebooks on the
SGEMM fallback, use the canonical encoder. Codes can differ across these
backends; NaN and high-dynamic [...]
</section>
<section class="article-section" id="params">
diff --git a/docs/ivf-pq.html b/docs/ivf-pq.html
index 8ceeb65..d39de04 100644
--- a/docs/ivf-pq.html
+++ b/docs/ivf-pq.html
@@ -48,6 +48,7 @@ options.put("metric", "l2");
options.put("pq.code-ratio", "0.0625"); // optional; default 0.0625
options.put("use-opq", "false"); // optional; default false
// options.put("ivf.coarse-assignment", "exact"); // optional; default auto
+// options.put("ivf.pq-encoding", "canonical"); // optional; default auto
try (VectorIndexTraining training =
VectorIndexTrainer.train(options, trainingVectors, trainingCount);
@@ -72,7 +73,8 @@ try (VectorIndexTraining training =
<section class="article-section" id="parameters">
<h2>Parameters</h2>
- <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement /
default</th><th>Purpose</th><th>Tuning
meaning</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Input dimension
<code>d</code></td><td>The inferred or explicit <code>pq.m</code> must divide
it</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partitio
[...]
+ <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement /
default</th><th>Purpose</th><th>Tuning
meaning</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Input dimension
<code>d</code></td><td>The inferred or explicit <code>pq.m</code> must divide
it</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partitio
[...]
+ <p>All PQ encoding backends produce the same index format. The
accelerated automatic paths require 8-bit PQ (<code>ksub=256</code>) and at
least four dimensions in every subvector. Eligible shapes use transposed direct
squared L2 on x86 with AVX2+FMA and on AArch64; other CPUs use blocked SGEMM
when the codebook is finite. Unsupported shapes and non-finite codebooks on the
SGEMM fallback use canonical. SGEMM and canonical use the expanded form, so
codes can differ near ties; Na [...]
</section>
<section class="article-section" id="storage">
diff --git a/docs/releases.html b/docs/releases.html
index 75598c0..65a2fb2 100644
--- a/docs/releases.html
+++ b/docs/releases.html
@@ -51,7 +51,8 @@
<p>The repository is currently developing the 0.5.0 line. Until an
ASF vote passes and the signed source archive appears under Apache downloads,
code and packages from this line are development artifacts rather than an
Apache release.</p>
<p>IVF-SQ now pools residual training bounds, fuses residual
encoding, transposes output in bounded parallel batches, and reuses query heaps
with conservative L2 block pruning. Unified readers and language bindings also
cache decoded partitions within the existing reader memory budget. The <a
href="ivf-sq.html#benchmarks">SIFT1M/GIST1M/GloVe benchmarks</a> record build,
native query, and recall results with reproducible commands.</p>
<p>IVSQ v1 files remain compatible. Existing indexes receive reader
optimizations without a rebuild; retrain and rebuild to use the new
quantization bounds. Set the reader memory budget to zero to disable the SQ
cache; direct Rust <code>IVFSQIndexReader::open</code> remains uncached. See <a
href="api.html#reader-options">reader options</a> and <a
href="ivf-sq.html#storage">upgrade details</a>.</p>
- <div class="callout warning"><strong>Rust IVF API
migration</strong>Version 0.5.0 makes <code>quantizer_centroids</code> private
on <code>IVFFlatIndex</code>, <code>IVFPQIndex</code>, <code>IVFSQIndex</code>,
and <code>IVFRQIndex</code>. Replace direct reads with
<code>quantizer_centroids()</code> and direct assignments with
<code>set_quantizer_centroids(...)</code>. The setter validates the centroid
shape, rejects replacement after vectors are added, and refreshes cached deriv
[...]
+ <div class="callout warning"><strong>Rust IVF API
migration</strong>Version 0.5.0 makes <code>quantizer_centroids</code> private
on <code>IVFFlatIndex</code>, <code>IVFPQIndex</code>, <code>IVFSQIndex</code>,
and <code>IVFRQIndex</code>. Replace direct reads with
<code>quantizer_centroids()</code> and direct assignments with
<code>set_quantizer_centroids(...)</code>. The setter validates the centroid
shape, rejects replacement after vectors are added, and refreshes cached deriv
[...]
+ <div class="callout warning"><strong>IVF-PQ build encoding changes
by default</strong>For 8-bit PQ (<code>ksub=256</code>) where every subvector
has at least four dimensions, the default automatic encoder uses transposed
direct-L2 on x86 with AVX2+FMA and on AArch64. Other CPUs use blocked SGEMM
expanded-form encoding for finite codebooks. Unsupported shapes, and non-finite
codebooks on the SGEMM fallback, use canonical encoding. Codes can differ
across these backends at near t [...]
</section>
<section class="article-section" id="release-040">
diff --git
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java
index 761762e..b9bc6a3 100644
---
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java
+++
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java
@@ -346,6 +346,7 @@ public class VectorIndexJavaApiTest {
private static void testReaderAndWriterApiCompile() {
Map<String, String> options = ivfPqOptions(2, 4);
+ options.put("ivf.pq-encoding", "canonical");
VectorIndexReader closedReader =
VectorIndexReader.fromNativePointerForTesting(0L);
closedReader.close();
closedReader.close();
diff --git
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java
index a7fd445..e33692d 100644
---
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java
+++
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java
@@ -423,7 +423,14 @@ public class VectorIndexNativeValidationTest {
private static void testSupportedIndexRoundtrips() {
runRoundtrip("ivf_flat", ivfFlatOptions(ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST), 0, 0);
- runRoundtrip("ivf_pq", ivfPqOptions(ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST), 2, 8);
+ Map<String, String> automaticPq =
+ ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST);
+ automaticPq.put("ivf.pq-encoding", "auto");
+ runRoundtrip("ivf_pq", automaticPq, 2, 8);
+ Map<String, String> canonicalPq =
+ ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST);
+ canonicalPq.put("ivf.pq-encoding", "canonical");
+ runRoundtrip("ivf_pq", canonicalPq, 2, 8);
runRoundtrip("ivf_rq", ivfRqOptions(ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST), 0, 0);
runRoundtrip(
"ivf_sq",
diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py
index 950f714..d82490b 100644
--- a/python/tests/test_vindex.py
+++ b/python/tests/test_vindex.py
@@ -278,6 +278,18 @@ def test_python_ffi_roundtrips_supported_indexes():
"nlist": "4",
"metric": "l2",
"use-opq": "false",
+ "ivf.pq-encoding": "auto",
+ },
+ 16,
+ ),
+ (
+ {
+ "index.type": "ivf_pq",
+ "dimension": "16",
+ "nlist": "4",
+ "metric": "l2",
+ "use-opq": "false",
+ "ivf.pq-encoding": "canonical",
},
16,
),
