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 26044a3 Optimize IVF-RQ search hot paths (#57)
26044a3 is described below
commit 26044a3d371f81aaad534e188280f1493610eb6b
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Jul 10 10:20:01 2026 +0800
Optimize IVF-RQ search hot paths (#57)
---
core/src/ivfrq_io.rs | 267 ++++++++++++++++++++++++++++++++++++++++++---------
core/src/kmeans.rs | 59 +++++++++++-
core/src/rq.rs | 47 +++++++--
core/src/topk.rs | 69 ++++++++++---
4 files changed, 373 insertions(+), 69 deletions(-)
diff --git a/core/src/ivfrq_io.rs b/core/src/ivfrq_io.rs
index ea1bbcb..c4da6fa 100644
--- a/core/src/ivfrq_io.rs
+++ b/core/src/ivfrq_io.rs
@@ -39,6 +39,7 @@ const FLAG_DELTA_IDS: u32 = 1 << 0;
const REQUIRED_FLAGS: u32 = FLAG_DELTA_IDS;
const SUPPORTED_FLAGS: u32 = REQUIRED_FLAGS;
const FACTOR_BYTES: usize = 12;
+const MAX_RQ_BATCH_READ_BYTES: usize = 16 * 1024 * 1024;
pub const IVF_RQ_NUM_BITS_ONE: u32 = 1;
pub const IVF_RQ_ROTATION_TYPE_KAC: u32 = 1;
@@ -208,6 +209,15 @@ pub struct IVFRQIndexReader<R: SeekRead> {
loaded: bool,
}
+#[derive(Clone, Copy)]
+struct RQListPayloadMeta {
+ list_id: usize,
+ count: usize,
+ offset: u64,
+ id_bytes_len: usize,
+ payload_len: usize,
+}
+
impl<R: SeekRead> IVFRQIndexReader<R> {
pub fn open(mut reader: R) -> io::Result<Self> {
let mut cursor = PreadCursor::new(&mut reader, 0);
@@ -334,6 +344,93 @@ impl<R: SeekRead> IVFRQIndexReader<R> {
pub fn read_inverted_list(&mut self, list_id: usize) ->
io::Result<RQReadList> {
self.ensure_loaded()?;
+ let Some(meta) = self.list_payload_meta(list_id)? else {
+ return Ok(RQReadList {
+ list_id,
+ ids: Vec::new(),
+ codes: Vec::new(),
+ factors: Vec::new(),
+ });
+ };
+ let mut payload = vec![0u8; meta.payload_len];
+ self.reader
+ .pread(&mut [ReadRequest::new(meta.offset, &mut payload)])?;
+ self.decode_inverted_list_payload(meta, &payload)
+ }
+
+ fn read_inverted_lists(&mut self, list_ids: &[usize]) ->
io::Result<Vec<RQReadList>> {
+ self.ensure_loaded()?;
+
+ let mut results: Vec<Option<RQReadList>> = (0..list_ids.len()).map(|_|
None).collect();
+ let mut metas = Vec::new();
+ let mut payloads = Vec::new();
+
+ for (input_index, &list_id) in list_ids.iter().enumerate() {
+ let Some(meta) = self.list_payload_meta(list_id)? else {
+ results[input_index] = Some(RQReadList {
+ list_id,
+ ids: Vec::new(),
+ codes: Vec::new(),
+ factors: Vec::new(),
+ });
+ continue;
+ };
+ metas.push((input_index, meta));
+ payloads.push(vec![0u8; meta.payload_len]);
+ }
+
+ if !metas.is_empty() {
+ {
+ let mut requests: Vec<_> = payloads
+ .iter_mut()
+ .zip(metas.iter())
+ .map(|(payload, (_, meta))| {
+ ReadRequest::new(meta.offset, payload.as_mut_slice())
+ })
+ .collect();
+ self.reader.pread(&mut requests)?;
+ }
+
+ for ((input_index, meta), payload) in
metas.into_iter().zip(payloads) {
+ results[input_index] =
Some(self.decode_inverted_list_payload(meta, &payload)?);
+ }
+ }
+
+ results
+ .into_iter()
+ .map(|result| {
+ result.ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ "missing IVF-RQ inverted list read result",
+ )
+ })
+ })
+ .collect()
+ }
+
+ fn batch_read_end(&self, list_ids: &[usize]) -> io::Result<usize> {
+ let mut payload_bytes = 0usize;
+ for (index, &list_id) in list_ids.iter().enumerate() {
+ let Some(meta) = self.list_payload_meta(list_id)? else {
+ continue;
+ };
+ let next_payload_bytes =
+ payload_bytes.checked_add(meta.payload_len).ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ "IVF-RQ batched payload size overflow",
+ )
+ })?;
+ if index > 0 && next_payload_bytes > MAX_RQ_BATCH_READ_BYTES {
+ return Ok(index);
+ }
+ payload_bytes = next_payload_bytes;
+ }
+ Ok(list_ids.len())
+ }
+
+ fn list_payload_meta(&self, list_id: usize) ->
io::Result<Option<RQListPayloadMeta>> {
if list_id >= self.nlist {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
@@ -342,12 +439,7 @@ impl<R: SeekRead> IVFRQIndexReader<R> {
}
let count = self.list_counts[list_id] as usize;
if count == 0 {
- return Ok(RQReadList {
- list_id,
- ids: Vec::new(),
- codes: Vec::new(),
- factors: Vec::new(),
- });
+ return Ok(None);
}
if !self.delta_ids {
return Err(io::Error::new(
@@ -366,22 +458,46 @@ impl<R: SeekRead> IVFRQIndexReader<R> {
.ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidData, "IVF-RQ list
payload overflow")
})?;
- let mut payload = vec![0u8; payload_len];
- self.reader
- .pread(&mut [ReadRequest::new(offset, &mut payload)])?;
+ Ok(Some(RQListPayloadMeta {
+ list_id,
+ count,
+ offset,
+ id_bytes_len,
+ payload_len,
+ }))
+ }
+
+ fn decode_inverted_list_payload(
+ &self,
+ meta: RQListPayloadMeta,
+ payload: &[u8],
+ ) -> io::Result<RQReadList> {
+ if payload.len() != meta.payload_len {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!(
+ "IVF-RQ list {} payload length {} does not match expected
{}",
+ meta.list_id,
+ payload.len(),
+ meta.payload_len
+ ),
+ ));
+ }
let base_id = i64::from_le_bytes(payload[0..8].try_into().unwrap());
let encoded_len =
i32::from_le_bytes(payload[8..12].try_into().unwrap());
- if encoded_len < 0 || encoded_len as usize != id_bytes_len {
+ if encoded_len < 0 || encoded_len as usize != meta.id_bytes_len {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"IVF-RQ id_bytes_len mismatch",
));
}
- let ids = decode_delta_varint_ids(base_id, &payload[12..12 +
id_bytes_len], count)?;
- let code_start = 12 + id_bytes_len;
+ let ids =
+ decode_delta_varint_ids(base_id, &payload[12..12 +
meta.id_bytes_len], meta.count)?;
+ let code_start = 12 + meta.id_bytes_len;
+ let code_bytes = checked_list_bytes(meta.count, self.code_size)?;
let factor_start = code_start + code_bytes;
Ok(RQReadList {
- list_id,
+ list_id: meta.list_id,
ids,
codes: payload[code_start..factor_start].to_vec(),
factors: bytes_to_factors(&payload[factor_start..])?,
@@ -528,39 +644,39 @@ pub fn
search_batch_ivfrq_reader_filter_with_query_bits<R: SeekRead>(
}
let mut heaps: Vec<TopKHeap> = (0..nq).map(|_| TopKHeap::new(k)).collect();
- for list_id in unique_lists {
- let count = reader.list_counts[list_id] as usize;
- if count == 0 {
- continue;
- }
- let read_list = reader.read_inverted_list(list_id)?;
- if read_list.list_id != list_id {
- return Err(io::Error::new(
- io::ErrorKind::InvalidData,
- "IVF-RQ inverted list read returned wrong list id",
- ));
- }
- let centroid = &reader.quantizer_centroids[list_id *
reader.d..(list_id + 1) * reader.d];
- for &qi in &list_to_queries[list_id] {
- let query = &processed[qi * reader.d..(qi + 1) * reader.d];
- let rotated_query_residual =
- rotated_residual(query, centroid, reader.d, &reader.rotation);
- let distance_context =
reader.quantizer.prepare_distance_context_with_query_bits(
- rotated_query_residual,
- query,
- count >= RQ_BYTE_LUT_MIN_LIST_SIZE,
- query_bits,
- );
- scan_read_list(
- &read_list,
- &reader.quantizer,
- reader.code_size,
- reader.metric,
- &distance_context,
- filter,
- &mut heaps[qi],
- );
+ let mut batch_start = 0;
+ while batch_start < unique_lists.len() {
+ let batch_end = batch_start +
reader.batch_read_end(&unique_lists[batch_start..])?;
+ for read_list in
reader.read_inverted_lists(&unique_lists[batch_start..batch_end])? {
+ let list_id = read_list.list_id;
+ let count = read_list.ids.len();
+ if count == 0 {
+ continue;
+ }
+ let centroid =
+ &reader.quantizer_centroids[list_id * reader.d..(list_id + 1)
* reader.d];
+ for &qi in &list_to_queries[list_id] {
+ let query = &processed[qi * reader.d..(qi + 1) * reader.d];
+ let rotated_query_residual =
+ rotated_residual(query, centroid, reader.d,
&reader.rotation);
+ let distance_context =
reader.quantizer.prepare_distance_context_with_query_bits(
+ rotated_query_residual,
+ query,
+ count >= RQ_BYTE_LUT_MIN_LIST_SIZE,
+ query_bits,
+ );
+ scan_read_list(
+ &read_list,
+ &reader.quantizer,
+ reader.code_size,
+ reader.metric,
+ &distance_context,
+ filter,
+ &mut heaps[qi],
+ );
+ }
}
+ batch_start = batch_end;
}
let mut result_ids = vec![-1i64; nq * k];
@@ -962,8 +1078,41 @@ fn decode_roaring_filter(bytes: &[u8]) ->
io::Result<RoaringTreemap> {
#[cfg(test)]
mod tests {
use super::*;
- use crate::io::PosWriter;
- use std::io::Cursor;
+ use crate::io::{PosWriter, ReadRequest, SeekRead};
+ use std::io::{Cursor, Read, Seek, SeekFrom};
+ use std::sync::{Arc, Mutex};
+
+ #[derive(Default)]
+ struct ReaderStats {
+ max_ranges_per_batch: usize,
+ }
+
+ struct CountingPreadCursor {
+ inner: Cursor<Vec<u8>>,
+ stats: Arc<Mutex<ReaderStats>>,
+ }
+
+ impl CountingPreadCursor {
+ fn new(data: Vec<u8>, stats: Arc<Mutex<ReaderStats>>) -> Self {
+ Self {
+ inner: Cursor::new(data),
+ stats,
+ }
+ }
+ }
+
+ impl SeekRead for CountingPreadCursor {
+ fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> {
+ let mut stats = self.stats.lock().unwrap();
+ stats.max_ranges_per_batch =
stats.max_ranges_per_batch.max(ranges.len());
+ drop(stats);
+ for range in ranges {
+ self.inner.seek(SeekFrom::Start(range.pos))?;
+ self.inner.read_exact(range.buf)?;
+ }
+ Ok(())
+ }
+ }
#[test]
fn ivfrq_write_read_search_roundtrip() {
@@ -1027,6 +1176,30 @@ mod tests {
}
}
+ #[test]
+ fn ivfrq_batch_reader_reads_lists_in_one_pread_batch() {
+ let mut index = IVFRQIndex::new(8, 3, MetricType::L2);
+ index.quantizer_centroids = vec![0.0; index.nlist * index.d];
+ index.ids = vec![vec![10], vec![20], vec![30]];
+ index.codes = vec![vec![0b0000_0001], vec![0b0000_0010],
vec![0b0000_0100]];
+ index.factors = vec![
+ vec![RQCodeFactors::zero()],
+ vec![RQCodeFactors::zero()],
+ vec![RQCodeFactors::zero()],
+ ];
+
+ let mut bytes = Vec::new();
+ write_ivfrq_index(&index, &mut PosWriter::new(&mut bytes)).unwrap();
+ let stats = Arc::new(Mutex::new(ReaderStats::default()));
+ let stream = CountingPreadCursor::new(bytes, Arc::clone(&stats));
+ let mut reader = IVFRQIndexReader::open(stream).unwrap();
+
+ let (ids, _) = search_batch_ivfrq_reader(&mut reader, &[0.0; 8], 1, 1,
3).unwrap();
+
+ assert!(ids[0] >= 0);
+ assert_eq!(stats.lock().unwrap().max_ranges_per_batch, 3);
+ }
+
#[test]
fn ivfrq_reader_rejects_invalid_query_bits() {
let d = 8;
diff --git a/core/src/kmeans.rs b/core/src/kmeans.rs
index 2e1ee6a..8d03e8f 100644
--- a/core/src/kmeans.rs
+++ b/core/src/kmeans.rs
@@ -515,10 +515,13 @@ pub fn find_topk(
nprobe: usize,
) -> (Vec<usize>, Vec<f32>) {
let nprobe = nprobe.min(k);
+ if nprobe == 0 {
+ return (Vec::new(), Vec::new());
+ }
let mut dists: Vec<(f32, usize)> = (0..k)
.map(|c| (fvec_l2sqr(point, ¢roids[c * d..(c + 1) * d]), c))
.collect();
- dists.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
+ select_topk_prefix(&mut dists, nprobe);
let indices: Vec<usize> = dists[..nprobe].iter().map(|&(_, i)|
i).collect();
let distances: Vec<f32> = dists[..nprobe].iter().map(|&(d, _)|
d).collect();
(indices, distances)
@@ -535,6 +538,9 @@ pub fn find_topk_batch(
nprobe: usize,
) -> (Vec<Vec<usize>>, Vec<Vec<f32>>) {
let nprobe = nprobe.min(k);
+ if nprobe == 0 {
+ return (vec![Vec::new(); nq], vec![Vec::new(); nq]);
+ }
if nq == 1 {
let (indices, distances) = find_topk(&queries[..d], centroids, k, d,
nprobe);
@@ -565,7 +571,7 @@ pub fn find_topk_batch(
(dist.max(0.0), c)
})
.collect();
- dists.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
+ select_topk_prefix(&mut dists, nprobe);
all_indices.push(dists[..nprobe].iter().map(|&(_, i)| i).collect());
all_distances.push(dists[..nprobe].iter().map(|&(d, _)| d).collect());
@@ -574,6 +580,20 @@ pub fn find_topk_batch(
(all_indices, all_distances)
}
+fn select_topk_prefix(dists: &mut [(f32, usize)], nprobe: usize) {
+ debug_assert!(nprobe > 0 && nprobe <= dists.len());
+ if nprobe < dists.len() {
+ dists.select_nth_unstable_by(nprobe - 1, compare_distance_then_index);
+ }
+ dists[..nprobe].sort_by(compare_distance_then_index);
+}
+
+fn compare_distance_then_index(left: &(f32, usize), right: &(f32, usize)) ->
std::cmp::Ordering {
+ left.0
+ .total_cmp(&right.0)
+ .then_with(|| left.1.cmp(&right.1))
+}
+
// --- Streaming Coreset K-means ---
/// Streaming k-means trainer for very large datasets.
@@ -779,6 +799,41 @@ mod tests {
assert_eq!(indices[0], 0);
}
+ #[test]
+ fn test_find_topk_batch_matches_full_sort_with_ties() {
+ let d = 2;
+ let k = 32;
+ let nprobe = 5;
+ let centroids: Vec<f32> = (0..k)
+ .flat_map(|i| [i as f32 % 4.0, (i / 4) as f32])
+ .collect();
+ let queries = vec![0.5, 0.5, 2.5, 3.5, 1.5, 1.5];
+ let (actual_indices, actual_distances) =
+ find_topk_batch(&queries, 3, ¢roids, k, d, nprobe);
+
+ for qi in 0..3 {
+ let query = &queries[qi * d..(qi + 1) * d];
+ let mut expected: Vec<(f32, usize)> = (0..k)
+ .map(|ci| (fvec_l2sqr(query, ¢roids[ci * d..(ci + 1) *
d]), ci))
+ .collect();
+ expected.sort_by(compare_distance_then_index);
+ assert_eq!(
+ actual_indices[qi],
+ expected[..nprobe]
+ .iter()
+ .map(|&(_, index)| index)
+ .collect::<Vec<_>>()
+ );
+ assert_eq!(
+ actual_distances[qi],
+ expected[..nprobe]
+ .iter()
+ .map(|&(distance, _)| distance)
+ .collect::<Vec<_>>()
+ );
+ }
+ }
+
#[test]
fn test_find_nearest_batch_matches_scalar() {
let d = 5;
diff --git a/core/src/rq.rs b/core/src/rq.rs
index 2843ccd..c248c70 100644
--- a/core/src/rq.rs
+++ b/core/src/rq.rs
@@ -340,14 +340,19 @@ impl RaBitQuantizer {
for byte_idx in 0..code_size {
let dim_base = byte_idx * 8;
let dim_end = (dim_base + 8).min(self.d);
- for pattern in 0..256usize {
- let mut sum = 0.0f32;
- for dim in dim_base..dim_end {
- let bit = (pattern >> (dim - dim_base)) & 1;
- let value = rotated_query_residual[dim];
- sum += if bit != 0 { value } else { -value };
- }
- byte_signed_sums[byte_idx * 256 + pattern] = sum;
+ let lut = &mut byte_signed_sums[byte_idx * 256..(byte_idx + 1) *
256];
+ lut[0] = -rotated_query_residual[dim_base..dim_end]
+ .iter()
+ .sum::<f32>();
+ for pattern in 1..256usize {
+ let bit = pattern.trailing_zeros() as usize;
+ let previous = pattern & (pattern - 1);
+ let value = if dim_base + bit < dim_end {
+ rotated_query_residual[dim_base + bit]
+ } else {
+ 0.0
+ };
+ lut[pattern] = lut[previous] + 2.0 * value;
}
}
byte_signed_sums
@@ -595,6 +600,32 @@ mod tests {
}
}
+ #[test]
+ fn byte_lut_matches_scalar_signed_sum_for_every_pattern() {
+ let d = 13;
+ let quantizer = RaBitQuantizer::new(d);
+ let residual: Vec<f32> = (0..d).map(|i| i as f32 * 0.37 -
2.1).collect();
+ let lut = quantizer.build_byte_signed_sums(&residual);
+ let mut code = vec![0u8; quantizer.code_size()];
+
+ for first_byte in 0..=u8::MAX {
+ for second_byte in 0..=u8::MAX {
+ code[0] = first_byte;
+ code[1] = second_byte;
+ let scalar: f32 = residual
+ .iter()
+ .enumerate()
+ .map(|(dim, &value)| if get_bit(&code, dim) { value } else
{ -value })
+ .sum();
+ let actual = lut[first_byte as usize] + lut[256 + second_byte
as usize];
+ assert!(
+ (actual - scalar).abs() < 1e-5,
+ "code {first_byte:#010b} {second_byte:#010b}: {actual} !=
{scalar}"
+ );
+ }
+ }
+ }
+
#[test]
fn quantized_query_bit_planes_match_scalar_quantization() {
let d = 24;
diff --git a/core/src/topk.rs b/core/src/topk.rs
index 112f76d..6407d7f 100644
--- a/core/src/topk.rs
+++ b/core/src/topk.rs
@@ -39,26 +39,22 @@ impl TopKHeap {
if let Some(&idx) = self.positions.get(&id) {
if dist < self.data[idx].0 {
self.data[idx].0 = dist;
+ self.sift_down(idx);
}
return;
}
if self.data.len() < self.k {
self.positions.insert(id, self.data.len());
self.data.push((dist, id));
+ self.sift_up(self.data.len() - 1);
return;
}
- if let Some((worst_idx, _)) = self
- .data
- .iter()
- .enumerate()
- .max_by(|(_, a), (_, b)| a.0.total_cmp(&b.0))
- {
- if dist < self.data[worst_idx].0 {
- let old_id = self.data[worst_idx].1;
- self.positions.remove(&old_id);
- self.data[worst_idx] = (dist, id);
- self.positions.insert(id, worst_idx);
- }
+ if dist < self.data[0].0 {
+ let old_id = self.data[0].1;
+ self.positions.remove(&old_id);
+ self.data[0] = (dist, id);
+ self.positions.insert(id, 0);
+ self.sift_down(0);
}
}
@@ -70,6 +66,43 @@ impl TopKHeap {
pub(crate) fn len(&self) -> usize {
self.data.len()
}
+
+ fn sift_up(&mut self, mut index: usize) {
+ while index > 0 {
+ let parent = (index - 1) / 2;
+ if self.data[parent].0.total_cmp(&self.data[index].0).is_ge() {
+ break;
+ }
+ self.swap_entries(parent, index);
+ index = parent;
+ }
+ }
+
+ fn sift_down(&mut self, mut index: usize) {
+ loop {
+ let left = index * 2 + 1;
+ let right = left + 1;
+ let mut worst = index;
+ if left < self.data.len() &&
self.data[left].0.total_cmp(&self.data[worst].0).is_gt() {
+ worst = left;
+ }
+ if right < self.data.len() &&
self.data[right].0.total_cmp(&self.data[worst].0).is_gt()
+ {
+ worst = right;
+ }
+ if worst == index {
+ break;
+ }
+ self.swap_entries(index, worst);
+ index = worst;
+ }
+ }
+
+ fn swap_entries(&mut self, left: usize, right: usize) {
+ self.data.swap(left, right);
+ self.positions.insert(self.data[left].1, left);
+ self.positions.insert(self.data[right].1, right);
+ }
}
#[cfg(test)]
@@ -87,4 +120,16 @@ mod tests {
assert_eq!(heap.into_sorted(), vec![(1.0, 7), (3.0, 9)]);
}
+
+ #[test]
+ fn test_topk_heap_keeps_max_at_root_after_duplicate_update() {
+ let mut heap = TopKHeap::new(3);
+ heap.push(8.0, 1);
+ heap.push(6.0, 2);
+ heap.push(4.0, 3);
+ heap.push(1.0, 1);
+ heap.push(5.0, 4);
+
+ assert_eq!(heap.into_sorted(), vec![(1.0, 1), (4.0, 3), (5.0, 4)]);
+ }
}
